在GCN代码的全连接层后加入空间注意力机制层的代码

假设你已经实现了GCN的全连接层，现在要在它的后面加入空间注意力机制层。以下是参考代码：

import torch.nn as nn
import torch.nn.functional as F

class SpatialAttention(nn.Module):
    def __init__(self, in_channels):
        super(SpatialAttention, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels // 8, kernel_size=1)
        self.conv2 = nn.Conv2d(in_channels // 8, 1, kernel_size=1)
        self.gamma = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        f = self.conv1(x)
        g = self.conv1(x)
        h = self.conv1(x)

        s = torch.matmul(g.view(g.size(0), -1, g.size(3)), h.view(h.size(0), -1, h.size(3)).permute(0, 2, 1))
        beta = F.softmax(s, dim=-1)
        o = torch.matmul(beta, f.view(f.size(0), -1, f.size(3)))
        o = o.view(x.size())
        x = self.gamma * o + x

        return x

class GCN(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(GCN, self).__init__()
        self.fc = nn.Linear(in_channels, out_channels)
        self.spatial_att = SpatialAttention(out_channels)

    def forward(self, x, adj):
        x = self.fc(x)
        x = F.relu(x)
        x = torch.matmul(adj, x)
        x = self.spatial_att(x) # 加入空间注意力机制层
        return x

这里实现了一个简单的空间注意力机制层`SpatialAttention`，它的输入是GCN全连接层的输出，输出是加上了空间注意力机制的特征向量。在GCN的`forward`函数中，我们先将全连接层的输出输入到`SpatialAttention`中，再输出到下一层。注意，这里的输入`x`是一个二维张量，所以`SpatialAttention`使用的是二维卷积。

`SpatialAttention`的实现是基于论文[CBAM: Convolutional Block Attention Module](https://arxiv.org/abs/1807.06521)中的空间注意力机制，它包括一个卷积层和一个权重计算层，通过计算输入特征图的通道间相似度，来为每个通道分配不同的权重，以提高空间上的注意力。