视频地址神经网络-线性层及其他层介绍_哔哩哔哩_bilibili
视频中称为线性层
import torch
import torchvision
from torch import nn
from torch.nn import Linear
from torch.utils.data import DataLoaderdataset = torchvision.datasets.CIFAR10("CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),download=True)dataloader = DataLoader(dataset, 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 outputtudui = Tudui()for data in dataloader:imgs, targets = dataprint(imgs.shape) # torch.Size([64, 3, 32, 32])output = torch.reshape(imgs, (1, 1, 1, -1))print(output.shape) # torch.Size([1, 1, 1, 196608])output = tudui(output)print(output.shape) # torch.Size([1, 1, 1, 10])
