线性层的作用,类似于这边荧光绿的两层,就是fully connected layer
将1*1*4096的向量,变成1*1*1000的向量
import torch
import torchvision
from torch import nn
from torch.nn import Linear
from torch.utils.data import DataLoader
class Module(nn.Module):
def __init__(self):
super().__init__()
self.linear = Linear(196608,10)
def forward(self, input):
output = self.linear(input)
return output
test_data = torchvision.datasets.CIFAR10('dataset',train=False,transform=torchvision.transforms.ToTensor(),download=True)
# drop_last (bool, optional) – set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size.
# If False and the size of dataset is not divisible by the batch size, then the last batch will be smaller. (default: False)
# drop_last可以把最后不足64张图的batch补足
data_loader = DataLoader(test_data,batch_size=64,drop_last=True)
module = Module()
for data in data_loader:
images ,target =data
# print(images.shape)
output = torch.reshape(images,(1,1,1,-1))
print(output.shape)
output=module(output)
print(output.shape)
- 另外,将向量展开成一行有一个特殊的函数:torch.flatten()
output = torch.flatten(images)