#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2022/5/6 08:41
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : nn_linear.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""
from torchvision.datasets import CIFAR10
from torchvision import transforms
from torch.utils.data import DataLoader
from torch import reshape, flatten
import torch.nn as nn

test_dataset = CIFAR10(root="cifai10", train=False, transform=transforms.ToTensor(), download=True)
dataloader = DataLoader(dataset=test_dataset, batch_size=64, drop_last=True)


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.linear1 = nn.Linear(196608, 10)

    def forward(self, x):
        return self.linear1(x)


my_net = MyNet()
for data in dataloader:
    images, targets = data
    print(images.shape)
    # output = reshape(images, (1, 1, 1, -1))
    output = flatten(images)
    print(output.shape)
    print(my_net(output).shape)