VGG16 模型输出不正确的维度 - 迁移学习

Question

我正在尝试使用预训练的 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])

不明白为什么模型会这样输出结果。有什么我遗漏的吗？

Answer 1

错误

这很简单，你在这里做的就是：

model = models.vgg16(pretrained=True).cuda()
model.classifier[6].out_features = 10

正在改变torch.nn.Linear层的属性out_features，没有改变实际执行计算的权重！

在最简单的情况下（请参阅输出形状的注释）：

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)

VGG16 模型输出不正确的维度 - 迁移学习

VGG16 Model Outputs Incorrect dimension - Transfer Learning

python

deep-learning

pytorch

transfer-learning

image-classification

错误

修复