VGG11使用卷积层代替全连接层进行测试

小孟lp

已于 2024-07-11 13:00:44 修改

阅读量1.9k

点赞数 1

分类专栏：深度学习文章标签：深度学习神经网络 cnn pytorch

于 2022-11-09 16:21:36 首次发布

本文链接：https://blog.csdn.net/qq_52062754/article/details/127764354

版权

深度学习专栏收录该内容

2 篇文章 0 订阅

订阅专栏

一、使用全连接层训练好的VGG11模型改为全卷积层进行测试，在加载训练模型报错问题

size mismatch报错

RuntimeError: Error(s) in loading state_dict for VGG:
size mismatch for classifier.0.weight: copying a param with shape torch.Size([4096, 25088]) from checkpoint, the shape in current model is torch.Size([4096, 512, 7, 7]).
size mismatch for classifier.3.weight: copying a param with shape torch.Size([4096, 4096]) from checkpoint, the shape in current model is torch.Size([4096, 4096, 1, 1]).
size mismatch for classifier.6.weight: copying a param with shape torch.Size([5, 4096]) from checkpoint, the shape in current model is torch.Size([5, 4096, 1, 1]).

报错代码

weights_path = r"vgg11Net.pth"
state_dict = torch.load(weights_path)
model.load_state_dict(state_dict)

问题分析：

model.load_state_dict(state_dict)加载模型时，全连接层训练的模型权重torch.size和卷积层的torch.size维度是不同的。

二、解决方法

我们要将训练好的模型的权重进行升维的操作，才能与卷积层的torch.size匹配起来。

全连接层改为卷积层代码

self.classifier = nn.Sequential(
            # nn.Linear(512*7*7, 4096),   #  classifier.0
            # nn.ReLU(True),
            # nn.Dropout(p=0.5),
            # nn.Linear(4096, 4096),   #  classifier.3
            # nn.ReLU(True),
            # nn.Dropout(p=0.5),
            # nn.Linear(4096, num_classes)   #  classifier.6
            # 全连接层与卷积层要对应起来，可不使用dropout，可以使用池化层将两者对应起来
            nn.Conv2d(in_channels=512, out_channels=4096, kernel_size=7, stride=1),  #  classifier.0
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Conv2d(in_channels=4096, out_channels=4096, kernel_size=1, stride=1),  #  classifier.3
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Conv2d(in_channels=4096, out_channels=num_classes, kernel_size=1, stride=1),  #  classifier.6
        )

增加的代码

在加载模型state_dict = torch.load(weights_path)后面添加以下代码，使用reshape()来改变加载训练模型权重的维度，与卷积层所需要的维度一样。

    state_dict = torch.load(weights_path)
    for k, v in state_dict.items():
        if "classifier.0.weight" == k:
            state_dict[k] = v.reshape(4096, 512, 7, 7)
        if "classifier.3.weight" == k:
            state_dict[k] = v.reshape(4096, 4096, 1, 1)
        if "classifier.6.weight" == k:
            state_dict[k] = v.reshape(num_classes, 4096, 1, 1)
    model.load_state_dict(state_dict)

修改之后，就可以用全卷积进行测试了。

三、为什么测试要使用卷积层代替全连接层及注意事项

1. VGG使用全连接层训练时图片为224224，所以使用全连接层模型进行测试也要缩放图片大小224224。

2. 使用全卷积层则不需要对测试的图片进行缩放，可以直接使用原图，但是输入的图片的最短边不能小于224（如果图片最短边小于224，可以在替换的第一层卷积加上一个padding，保证图片有足够的尺寸大小来被7*7的卷积核进行卷积操作）。

3. 使用全卷积最后经过softmax得到的输出权重不一定是一维的向量（也可以是二维和三维的）。

例如：输入图片大小为256*256，最后输出的为2*2*num_classes大小的输出（num_classes为测试分类数量），全卷积测试时图片大小变化及最后输出权重的形状如下。
全卷积测试图片时的大小变换

实现代码

with torch.no_grad():
    output = torch.squeeze(model(img.to(device))).cpu()
    predict = torch.softmax(output, dim=0)
    '''判断输出的权重为几维的 '''
    if len(predict.shape) == 3:
        three_dim = np.array(predict)  # 将tensor数据转换为三维的numpy数组
        two_dim = np.sum(three_dim, -1)  # 按每个feature map进行行求和，降为二维的numpy数组
        one_dim = two_dim.sum(axis=1)  # axis=0 按列求和， axis=1 按行求和
        predict = torch.tensor(one_dim)  # 将numpy数组转换为tensor
    if len(predict.shape) == 2:
        two_dim = np.array(predict)
        one_dim = two_dim.sum(axis=1)
        predict = torch.tensor(one_dim)
    '''以上为需要添加的代码'''
    predict_cla = torch.argmax(predict).numpy()  # 获取最大权重分类的下标