GNN小模型python代码

在Python中，特别是使用PyTorch Geometric (PyG)库来构建图神经网络（GNN）的小模型，通常会涉及到一些关键步骤，比如数据预处理、定义模型结构、训练和评估。下面是一个简单的GNN模型（如GCN或GraphSAGE）的示例：

import torch
from torch_geometric.nn import GCNConv

class SimpleGNN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(SimpleGNN, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))  # 首层GCN
        x = self.conv2(x, edge_index)  # 第二层GCN
        return F.log_softmax(x, dim=1)  # 输出层通常用softmax处理分类任务

# 使用模型
model = SimpleGNN(in_channels=16, hidden_channels=32, out_channels=10)  # 假设输入特征有16维
x = torch.randn(100, 16)  # 数据节点的特征
edge_index = torch.LongTensor([[0, 1], [1, 2]])  # 两个节点之间的边
output = model(x, edge_index)

在这个例子中，`in_channels`指输入特征的数量，`hidden_channels`和`out_channels`分别是隐藏层和输出层的神经元数量。`forward`函数描述了网络的前向传播过程。

注意：这只是一个基础示例，实际应用中可能还需要添加更多的功能，如损失函数、优化器、以及训练循环等。