class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio=4): super(ChannelAttention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) # 利用1x1卷积代替全连接 self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False) self.relu1 = nn.ReLU() self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x)))) max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x)))) out = avg_out + max_out return self.sigmoid(out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super(SpatialAttention, self).__init__() assert kernel_size in (3, 7), 'kernel size must be 3 or 7' padding = 3 if kernel_size == 7 else 1 self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) x = torch.cat([avg_out, max_out], dim=1) x = self.conv1(x) return self.sigmoid(x)

这段代码定义了两个 PyTorch 模块：通道注意力模块（ChannelAttention）和空间注意力模块（SpatialAttention）。这两个模块都是用来增强卷积神经网络的表达能力的。

通道注意力模块的输入是一个四维张量 x，其形状为 [batch_size, in_planes, height, width]，其中 in_planes 是输入特征图的通道数。在该模块中，首先对输入特征图进行了平均池化和最大池化操作，然后将两个池化结果分别经过一个 1x1 卷积层和激活函数，并将它们相加得到最终的特征图。最后，通过一个 sigmoid 激活函数将特征图的像素值归一化到 [0,1] 范围内。

空间注意力模块的输入也是一个四维张量 x，其形状与通道注意力模块相同。在该模块中，通过对输入特征图进行平均池化和最大池化操作，并将两个池化结果拼接在一起，得到了一个二通道的特征图。然后，将该特征图经过一个卷积层和 sigmoid 激活函数，得到一个 [0,1] 范围内的权重图。最后，将输入特征图与权重图相乘，得到了加强了空间细节的特征图。

这两个模块可以组合起来使用，以增强卷积神经网络的表达能力。