我正在尝试使用预训练的 VGG16 模型在 pyTorch 上对 CIFAR10 进行分类。该模型最初是在 ImageNet 上训练的。
以下是我导入和修改模型的方法：
from torchvision import models
model = models.vgg16(pretrained=True).cuda()
model.classifier[6].out_features = 10
这是模型的总结
print(model)
VGG(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace=True)
(4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(6): ReLU(inplace=True)
(7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(8): ReLU(inplace=True)
(9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace=True)
(12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(13): ReLU(inplace=True)
(14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): ReLU(inplace=True)
(16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): ReLU(inplace=True)
(19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(20): ReLU(inplace=True)
(21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(22): ReLU(inplace=True)
(23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(25): ReLU(inplace=True)
(26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(27): ReLU(inplace=True)
(28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(29): ReLU(inplace=True)
(30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
(classifier): Sequential(
(0): Linear(in_features=25088, out_features=4096, bias=True)
(1): ReLU(inplace=True)
(2): Dropout(p=0.5, inplace=False)
(3): Linear(in_features=4096, out_features=4096, bias=True)
(4): ReLU(inplace=True)
(5): Dropout(p=0.5, inplace=False)
(6): Linear(in_features=4096, out_features=10, bias=True)
)
)
现在我想在 CIFAR10 上训练模型，我创建了批量大小 = 128 的火车装载机。 下面是训练函数，我添加了一些打印语句来检查一切是否正常。 问题在于模型为每个数据点输出 1000 个预测（作为原始版本 - 未修改版本）。
def train(model, optimizer, train_loader, epoch=5):
"""
This function updates/trains client model on client data
"""
model.train()
for e in range(epoch):
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
print("shape output: ", output.shape)
probs = F.softmax(output, dim=1) #get the entropy
print("entropy: ", probs)
_, predicted = torch.max(output.data, 1)
loss = criterion(output, target)
loss.backward()
if batch_idx % 100 == 0: # print every 100 mini-batches
print('[%d, %5d] loss: %.3f' %
(e + 1, batch_idx + 1, loss.item()))
optimizer.step()
return loss.item()
这是显示输出形状的输出的一部分：
shape output: torch.Size([128, 1000])
我不明白为什么模型会以这种方式输出结果。有什么我遗漏的吗？
答案 0 :(得分：2)
很简单，你在这里做的一切：
model = models.vgg16(pretrained=True).cuda()
model.classifier[6].out_features = 10
正在改变 out_features 层的属性 torch.nn.Linear，不改变实际执行计算的权重！
在最简单的情况下（参见输出形状的注释）：
import torch
layer = torch.nn.Linear(20, 10)
layer.out_features # 10
layer.weight.shape # (10, 20)
layer(torch.randn(64, 20)).shape # (64, 10)
layer.out_features = 100
layer.out_features # 100
layer.weight.shape # (10, 20)
layer(torch.randn(64, 20)).shape # (64, 10)
权重的形状相同，因为它们是在 __init__ 期间创建的（改变 out_features 不会改变任何东西）。
您必须重新创建最后一层（它将被随机初始化），如下所示：
model = torchvision.models.vgg16(pretrained=True)
# New layer
model.classifier[6] = torch.nn.Linear(4096, 10)