PyTorch实现SENet（注意力机制）

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)