gpu版本的GNN的demo

news2024/12/23 11:12:19

1、当涉及到在GPU上运行图神经网络（GNN）时，通常使用深度学习框架，如PyTorch或TensorFlow。在这里，我将为您提供一个使用PyTorch Geometric库实现GNN的简单示例。

首先，确保您已经安装了PyTorch和PyTorch Geometric。您可以使用以下命令安装它们：

pip install torch
pip install torch-geometric

2、使用以下代码作为GPU版本的GNN的简单示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv
from torch_geometric.utils import to_networkx

# 检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 创建一个简单的图数据
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long, device=device)
x = torch.tensor([[0], [1], [2]], dtype=torch.float, device=device)

data = Data(x=x, edge_index=edge_index)

# 定义一个简单的图卷积网络（GCN）模型
class GCN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(input_size, hidden_size)
        self.conv2 = GCNConv(hidden_size, output_size)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = torch.relu(x)
        x = self.conv2(x, edge_index)
        return x

# 初始化模型、损失函数和优化器
model = GCN(input_size=1, hidden_size=16, output_size=1).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 将数据移动到GPU
data = data.to(device)

# 训练模型
for epoch in range(1000):
    model.train()
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, data.x)
    loss.backward()
    optimizer.step()
    print(f'Epoch {epoch + 1}/{1000}, Loss: {loss.item():.4f}')

# 在测试集上评估模型
model.eval()
with torch.no_grad():
    test_output = model(data)

print("Test output:", test_output.cpu().numpy())