PyTorch实现SENet（注意力机制）

最新推荐文章于 2024-04-09 02:27:24 发布

原创最新推荐文章于 2024-04-09 02:27:24 发布 · 944 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #pytorch

该博客介绍了一个名为SENet的网络结构，它通过通道注意力机制来增强特征表示能力。SE模块首先进行全局平均池化，然后通过线性变换调整通道维度，再通过Sigmoid激活函数得到每个通道的权重，最后将权重乘以原始特征图以提升模型对通道间信息的敏感度。

部署运行你感兴趣的模型镜像

import torch.nn as nn

class SENet(nn.Module):
    '''
    input:(batch_size, chs, h, w)
    output:(batch_size, chs, h, w)
    '''
    def __init__(self, chs, reduction=4):
        super(SENet, self).__init__()
        self.average_pooling = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        # (batch_size, chs, h, w) -> (batch, chs, 1, 1)
        self.fc = nn.Sequential(
            # First reduce dimension, then raise dimension.
            # Add nonlinear processing to fit the correlation between channels
            nn.Linear(chs, chs // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(chs // reduction, chs)
        )
        self.activation = nn.Sigmoid()  # Get the weight of each channel

    def forward(self, x):
        tmp = x
        batch_size, chs, h, w = x.shape
        x = self.average_pooling(x).view(batch_size, chs)
        # (batch_size, chs, h, w) -> (batch, chs, 1, 1) -> (batch, chs)
        x = self.fc(x).view(batch_size, chs, 1, 1)
        # (batch, chs) -> (batch, chs, 1, 1)
        # x = torch.clamp(x, min=0, max=1)
        # if min =< elem <= max, new_elem = elem, if elem > max, new_elem = max, if elem < min, new_elem = min
        x = self.activation(x)
        return x * tmp  # x: (batch_size, chs, 1, 1), tmp: (batch_size, chs, h, w)