线性层 就是拉平

import torch
import torchvision
from torch import nn

from torch.nn import MaxPool2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from torch.nn import Linear
from torch.nn import Flatten


input = torch.tensor([
    [1,2,0,3,1],
    [0,1,2,3,1],
    [1,2,1,0,0],
    [5,2,3,1,1],
    [2,1,0,1,1]
],dtype=torch.float32)

print(input.shape)


data = torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor())
dataloader = DataLoader(data,batch_size=64)

# for data in dataloader:
#     img,target = data
#     print(img.shape)
#     print(    img.reshape(1,1,1,-1).shape )




class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.liner1 = Linear(196608,10)    #不够3*3的边缘，要

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

tudui = Tudui()

for data in dataloader:
    img,target = data
    print(img.shape)


    output = img.reshape(1, 1, 1, -1)
    # output = torch.flatten(img)

    output = tudui(output)
    print(output.shape)
    print(output)