python 语言 创建一个GCN的模型

以下是一个简单的GCN模型的Python代码示例：

import torch
import torch.nn as nn

class GCN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(GCN, self).__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x, adj_matrix):
        x = self.linear1(x)
        x = torch.matmul(adj_matrix, x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

这个模型接受一个输入张量x和一个邻接矩阵adj_matrix。输入张量的形状为(batch_size, input_dim)，邻接矩阵的形状为(batch_size, num_nodes, num_nodes)，其中num_nodes是图中节点的数量。模型包含两个线性层和一个ReLU激活函数，第一个线性层将输入张量映射到隐藏维度，第二个线性层将隐藏维度映射到输出维度。在前向传递期间，模型首先将输入张量传递到第一个线性层，然后将邻接矩阵与线性层输出进行矩阵乘法，然后应用ReLU激活函数。最后，模型将结果传递到第二个线性层以获得最终输出。