GNN节点回归任务python代码

在Python中，特别是使用PyTorch Geometric (PyG)库处理图神经网络（GNN）时，进行节点回归任务通常涉及构建一个模型，该模型接受一个图结构作为输入，并预测每个节点的连续值。以下是一个简单的例子，展示了如何使用Graph Convolutional Networks (GCN)进行节点属性预测：

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

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

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        # 层间信息传播
        x = self.conv1(x, edge_index)
        x = F.relu(x)  # 使用ReLU激活函数
        x = self.conv2(x, edge_index)
        # 输出层
        x = torch.mean(x, dim=1)  # 可能需要对每个节点取平均值
        return x

# 假设我们有8维输入特征、16个隐藏节点和4个输出通道
model = NodeRegressionModel(8, 16, 4)
# 初始化权重
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

data = ...  # 图数据对象，应包含x（节点特征），edge_index（边的连接）
for epoch in range(num_epochs):  # 假设num_epochs是你想要训练的轮数
    optimizer.zero_grad()
    output = model(data)
    loss = F.mse_loss(output, target_nodes)  # 假设target_nodes是目标节点的标签
    loss.backward()
    optimizer.step()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

在这个例子中，你需要先安装`torch_geometric`库，然后准备一个包含`x`（节点特征矩阵）和`edge_index`（表示边的二维张量）的数据集实例。