帮我写一段CBAM代码

CBAM是一种用于计算机视觉中的注意力机制，可以用如下的Python代码实现：

import torch
import torch.nn as nn

class CBAM(nn.Module):
    def __init__(self, channels, reduction):
        super(CBAM, self).__init__()
        self.channels = channels
        self.reduction = reduction
        self.fc1 = nn.Linear(channels, channels // reduction, bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.fc2 = nn.Linear(channels // reduction, channels, bias=False)
        self.sigmoid_channel = nn.Sigmoid()
        self.conv1 = nn.Conv2d(channels, channels // reduction, kernel_size=1, stride=1, padding=0)
        self.conv2 = nn.Conv2d(channels // reduction, channels // reduction, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(channels // reduction, 1, kernel_size=3, stride=1, padding=1)
        self.sigmoid_spatial = nn.Sigmoid()

    def forward(self, x):
        x_channel_avg = torch.mean(x, dim=-1, keepdim=True)
        x_channel_avg = torch.mean(x_channel_avg, dim=-2, keepdim=True)
        x_channel_max = torch.max(x, dim=-1, keepdim=True)[0]
        x_channel_max = torch.max(x_channel_max, dim=-2, keepdim=True)[0]
        x_channel_att = self.fc2(self.relu(self.fc1(x_channel_avg + x_channel_max)))
        x_channel_att = self.sigmoid_channel(x_channel_att)

        x_spatial_att = self.conv1(x)
        x_spatial_att = self.relu(x_spatial_att)
        x_spatial_att = self.conv2(x_spatial_att)
        x_spatial_att = self.relu(x_spatial_att)
        x_spatial_att = self.conv3(x_spatial_att)
        x_spatial_att = self.sigmoid_spatial(x_spatial_att)

        x_att = x_channel_att * x_spatial_att
        return x_att * x

在这段代码中，我们定义了一个`CBAM`类，继承自`nn.Module`，并定义了`__init__`和`forward`两个方法。在`__init__`方法中，我们定义了需要用到的各种网络层，包括全连接层、卷积层等。在`forward`方法中，我们首先计算通道注意力和空间注意力，然后将两个注意力相乘，得到最终的注意力特征图。最后，我们将输入特征图和注意力特征图相乘，得到CBAM后的特征图，并返回。