gnn图神经网络实例代码

GNN（Graph Neural Network，图神经网络）是一种用于处理图结构数据的机器学习模型。下面是一个简单的GNN实例代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GNN(nn.Module):
    def __init__(self, num_features, num_classes):
        super(GNN, self).__init__()
        self.conv1 = GCNConv(num_features, 16)
        self.conv2 = GCNConv(16, num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

# 创建一个GNN模型实例
model = GNN(num_features, num_classes)

# 定义损失函数和优化器
criterion = nn.NLLLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(num_epochs):
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, data.y)
    loss.backward()
    optimizer.step()

# 使用训练好的模型进行预测
model.eval()
with torch.no_grad():
    output = model(data)
    predicted_labels = output.argmax(dim=1)

这是一个使用PyTorch Geometric库实现的简单的GNN模型。代码中使用了两个GCNConv层来构建图卷积网络，通过前向传播计算输出，并使用负对数似然损失函数进行训练。在训练过程中，使用Adam优化器进行参数更新。最后，使用训练好的模型进行预测。