SENET通道注意力机制源代码+注释

import torch
import torch.nn as nn
class SENET(nn.Module):
    def __init__(self, channel, r=0.5):  # channel为输入的维度, r为全连接层缩放比例->控制中间层个数
        super(SENET, self).__init__()
        # 全局均值池化
        self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(channel, int(channel * r)),  # int(channel * r)取整数   #通道压缩
            nn.ReLU(),          #relu激活函数进行激活 （）激励
            nn.Linear(int(channel * r), channel),     #展开
            nn.Sigmoid(),  #折算成0到1之间的权重
        )
    def forward(self, x):
        # 对x进行分支计算权重, 进行全局均值池化
        branch = self.global_avg_pool(x)   #前向传播先平均池化
        branch = branch.view(branch.size(0), -1)   #展开
        # 全连接层得到权重
        weight = self.fc(branch)    #经过全连接得到权重
        # 将维度为b, c的weight, reshape成b, c, 1, 1 与 输入x 相乘 即乘以权重
        hi, wi = weight.shape
        weight = torch.reshape(weight, (hi, wi, 1, 1))
        # 乘积获得结果
        scale = weight * x    #weight为权重 x为输入的  输出结果
        return scale
# alexnet1.add_module("linear",nn.Linear(1000 , 10))
alexnet1 = SENET(channel = 3) #通道数等于输入图片的通道数
print(alexnet1)
test1 = torch.ones(64, 3, 120, 120)  #输入64batch  3通道 120*120的图片
test1 = alexnet1 (test1)
print(test1.shape)   #输出无变化