使用Python写一个简单的时空图神经网络模型并对其进行训练。

以下是一个简单的时空图神经网络模型的示例代码：

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

class SpatioTemporalGraphConvolution(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(SpatioTemporalGraphConvolution, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=1)

    def forward(self, x, A):
        x = self.conv1(x)
        x = torch.matmul(A, x)
        x = self.conv2(x)
        x = F.relu(x)
        return x

class STGCN(nn.Module):
    def __init__(self, num_nodes, in_channels, num_timesteps, num_classes):
        super(STGCN, self).__init__()
        self.num_nodes = num_nodes
        self.num_timesteps = num_timesteps

        self.stgcn_layers = nn.ModuleList([
            SpatioTemporalGraphConvolution(in_channels, 64),
            SpatioTemporalGraphConvolution(64, 64),
            SpatioTemporalGraphConvolution(64, 64)
        ])

        self.fc = nn.Linear(64, num_classes)

    def forward(self, x, A):
        for layer in self.stgcn_layers:
            x = layer(x, A)
        x = x.mean(dim=(-2, -1))
        x = self.fc(x)
        x = F.log_softmax(x, dim=1)
        return x

# 构建数据
num_nodes = 10
in_channels = 3
num_timesteps = 5
num_classes = 2
batch_size = 32

x = torch.randn(batch_size, in_channels, num_nodes, num_timesteps)
A = torch.randn(batch_size, num_nodes, num_nodes)

# 构建模型
model = STGCN(num_nodes, in_channels, num_timesteps, num_classes)

# 训练模型
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.NLLLoss()

for epoch in range(10):
    optimizer.zero_grad()
    output = model(x, A)
    loss = criterion(output, torch.tensor([0, 1]*int(batch_size/2)))  # 示例标签
    loss.backward()
    optimizer.step()
    print('Epoch: {}, Loss: {}'.format(epoch+1, loss.item()))

这个示例中，我们定义了一个简单的时空图神经网络模型 STGCN，它由三个 SpatioTemporalGraphConvolution 层和一个全连接层组成。我们使用 nn.ModuleList 来管理层，并使用 nn.Conv2d 实现卷积操作。模型的输入数据 x 是一个四维张量，表示维度为 (batch_size, in_channels, num_nodes, num_timesteps) 的输入图数据。A 是一个二维张量，表示维度为 (batch_size, num_nodes, num_nodes) 的邻接矩阵。我们使用交叉熵损失函数 nn.NLLLoss 和 Adam 优化器 torch.optim.Adam 进行训练。

在训练过程中，我们通过迭代模型的前向传播、计算损失、反向传播、更新模型参数来训练模型