【DAY8-代码】Pytorch缝合模块之Resnet和SE注意力模块，以及修改模型后如何加载预训练权重

咸鱼神经蛙

于 2024-09-27 21:35:04 发布

阅读量453

点赞数 5

分类专栏：每日学习文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/weixin_55940866/article/details/142601003

版权

每日学习专栏收录该内容

13 篇文章 1 订阅

订阅专栏

准备代码

首先把这篇文章的代码copy下来。【DL系列】ResNet网络结构详解、完整代码实现
在这里插入图片描述
也可以直接看resnet源码（复杂一点）。resnet源码

再看SE注意力模块代码注意力机制（一）SE模块（Squeeze-and-Excitation Networks）论文总结和代码实现

把SE模块复制到resnet的model文件中，塞到某个子模块下面就行。注意，basicblock和bottleneck都是resnet网络的子模块，ResNet才是主网络。
在这里插入图片描述

#全局平均池化+1*1卷积核+ReLu+1*1卷积核+Sigmoid
class SE_Block(nn.Module):
    def __init__(self, inchannel, ratio=16):
        super(SE_Block, self).__init__()
        # 全局平均池化(Fsq操作)
        self.gap = nn.AdaptiveAvgPool2d((1, 1))
        # 两个全连接层(Fex操作)
        self.fc = nn.Sequential(
            nn.Linear(inchannel, inchannel // ratio, bias=False),  # 从 c -> c/r
            nn.ReLU(),
            nn.Linear(inchannel // ratio, inchannel, bias=False),  # 从 c/r -> c
            nn.Sigmoid()
        )
 
    def forward(self, x):
            # 读取批数据图片数量及通道数
            b, c, h, w = x.size()
            # Fsq操作：经池化后输出b*c的矩阵
            y = self.gap(x).view(b, c)
            # Fex操作：经全连接层输出（b，c，1，1）矩阵
            y = self.fc(y).view(b, c, 1, 1)
            # Fscale操作：将得到的权重乘以原来的特征图x
            return x * y.expand_as(x)

看看SE模块的输入和输出形状，这里我设置的输入通道是2，模块的输出也是2。
在这里插入图片描述

开始缝合

缝合模块只需要看模块的输入输出通道维数是否与缝合部位的维数匹配。因此我们先随便缝在一个地方，看看该部位的特征维数是多少。
首先，把SE_Block添加进ResNet的初始化里面（和self放一起就行）。
在这里插入图片描述
然后找到resnet的forward函数，假设要把se加在conv5下面，那么我们打印一下这一层的输出。

可以看到输出的通道数为512，因此我们的se模块应该也是512输入，512输出。

添加se模块。

输出的shape与没加模块之前一样，说明加得没问题。
在这里插入图片描述