import torchvision
from torch import nn
from torch.nn import Linear
from torch.utils.data import DataLoader
import torch
# 測試集

test_data=torchvision.datasets.CIFAR10(root="./CIFAR10_dataset",transform=torchvision.transforms.ToTensor(),train=False,download=True)
test_loader=DataLoader(dataset=test_data,batch_size=64)

#模型定義
class Tudui(nn.Module):
    def  __init__(self):
        super(Tudui, self).__init__()
        self.linear1=Linear(196608,10)

    def forward(self,input):
        output=self.linear1(input)
        return output

#宣告模型
tudui=Tudui()

for data in test_loader:
    imgs,targets=data
    print(imgs.shape) #torch.Size([64, 3, 32, 32])

    # output=torch.reshape(imgs,(1,1,1,-1))#-1表示寬隨大小自由設定
    # print(output.shape) #torch.Size([1, 1, 1, 196608])

    #展平輸入向量
    output=torch.flatten(imgs)
    print(output.shape) #torch.Size([196608])
    #呼叫模型
    output=tudui(output)
    print(output.shape) #torch.Size([10])