RuntimeError: size mismatch

RuntimeError: size mismatch出现于pytorch框架下加载VGG11预训练模型时出现的错误。笔者初期认为,出现该错误的主要原因是输出和输入的维度不匹配。带着疑惑,我们通过输出网络模型结构来观察:
在这里插入图片描述
我们可以看到(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。

最后,倘若尝试以上方法未果,可以留意输入图片尺寸。注意训练时输入图片尺寸与加载模型时使用的输入图片尺寸是否一致。一般来说输入图片的大小为224的倍数(要找一个72的指数次方,72的5次方等于7*32=224)

发布了25 篇原创文章 · 获赞 28 · 访问量 3万+

猜你喜欢

转载自blog.csdn.net/qq_30680871/article/details/90733811
今日推荐