RuntimeError: size mismatch

RuntimeError: size mismatch出现于pytorch框架下加载VGG11预训练模型时出现的错误。笔者初期认为,出现该错误的主要原因是输出和输入的维度不匹配。带着疑惑,我们通过输出网络模型结构来观察:
RuntimeError: size mismatch_第1张图片
我们可以看到(18): Conv2d输出的维度是512,(avgpool): AdaptiveAvgPool2d输出的维度是7*7,(0): Linear输入的维度是25088。在这里,我们将25088分解,即25088=512*7*7。当时笔者设置的是model.avgpool = nn.AdaptiveAvgPool2d(1),因此512*1*1=512!=25088。随后,model.avgpool = nn.AdaptiveAvgPool2d(7)调试成功。
建议:加载预训练模型时若出现size mismatch的问题,先观察网络结构,重点观察
model.avgpool = nn.AdaptiveAvgPool2d(1)
model.fc = nn.Linear(2048,config.num_classes)
中的参数是否吻合原网络。

这里补充如果去除预训练网络某几层,自定义添加网络层方法。

import torchvision.models as models
from torchsummary import summary
import torch
import torchvision
import torch.nn.functional as F
from torch import nn
from config import config

class Net(nn.Module):
    def __init__(self, model):
        super(Net, self).__init__()
        # 去除原网络最后两层
        self.resnet_layer = nn.Sequential(*list(model.children())[:-2])
        # 自定义添加网络层
        self.transion_layer = nn.ConvTranspose2d(2048, 2048, kernel_size=14, stride=3)
        self.pool_layer = nn.MaxPool2d(32)
        self.Linear_layer = nn.Linear(2048, 8)

    def forward(self, x):
        x = self.resnet_layer(x)
        x = self.transion_layer(x)
        x = self.pool_layer(x)
        x = x.view(x.size(0), -1)
        x = self.Linear_layer(x)
        return x

定义好自己的网络后,加载预训练网络,传入到Net中即可。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg19bnft = models.vgg19_bn().to(device)
vgg19bnft = Net(vgg19bnft)
print(vgg19bnft)

若想单独修改最后的全连接层,可以尝试如下方法:

removed = list(vgg19bn.classifier.children())[:-1]
vgg19bn.classifier = torch.nn.Sequential(*removed)
vgg19bn.add_module('fc', torch.nn.Linear(4096, 7))
print(vgg19bn)

因为加载预训练模型的种类是1000,这里根据自己的需求,修改为了7。

你可能感兴趣的:(深度学习)