PyTorch Geometric（PyG）机器学习实战

在图神经网络（GNN）的研究和应用中，PyTorch Geometric（PyG）作为一个基于PyTorch的库，提供了高效的图数据处理和模型构建功能。
本文将通过一个节点分类任务，演示如何使用PyG进行机器学习实战。

1. 环境准备

首先，确保已安装PyTorch和PyG。可以使用以下命令进行安装：

pip install torch
pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric2. 导入必要的库import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv3. 加载数据集我们使用PyG自带的Planetoid数据集，这里以Cora数据集为例。dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]4. 定义GCN模型我们将构建一个包含两层图卷积层（GCNConv）的模型。class GCN(nn.Module):def __init__(self, in_channels, hidden_channels, out_channels):super(GCN, 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_indexx = self.conv1(x, edge_index)x = F.relu(x)x = self.conv2(x, edge_index)return F.log_softmax(x, dim=1)5. 初始化模型和优化器model = GCN(in_channels=dataset.num_node_features,hidden_channels=16,out_channels=dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)6. 训练模型def train():model.train()optimizer.zero_grad()out = model(data)loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])loss.backward()optimizer.step()return loss.item()for epoch in range(200):loss = train()if epoch % 10 == 0:print(f'Epoch {epoch}, Loss: {loss:.4f}')7. 测试模型def test():model.eval()out = model(data)pred = out.argmax(dim=1)correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()acc = int(correct) / int(data.test_mask.sum())return accaccuracy = test()
print(f'Accuracy: {accuracy:.4f}')8. 完整代码import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv# 加载数据集
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]# 定义GCN模型
class GCN(nn.Module):def __init__(self, in_channels, hidden_channels, out_channels):super(GCN, 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_indexx = self.conv1(x, edge_index)x = F.relu(x)x = self.conv2(x, edge_index)return F.log_softmax(x, dim=1)# 初始化模型和优化器
model = GCN(in_channels=dataset.num_node_features,hidden_channels=16,out_channels=dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)# 训练模型
def train():model.train()optimizer.zero_grad()out = model(data)loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])loss.backward()optimizer.step()return loss.item()for epoch in range(200):loss = train()if epoch % 10 == 0:print(f'Epoch {epoch}, Loss: {loss:.4f}')# 测试模型
def test():model.eval()out = model(data)pred = out.argmax(dim=1)correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()acc = int(correct) / int(data.test_mask.sum())return accaccuracy = test()
print(f'Accuracy: {accuracy:.4f}')
'''9. 结果分析通过上述步骤，我们成功地使用PyG构建并训练了一个图卷积神经网络（GCN）模型。
在训练过程中，模型逐步学习图结构数据的特征，最终在测试集上取得了较好的分类准确率。
这展示了PyG在图数据处理和模型构建方面的强大功能。10. 参考文献• PyTorch Geometric官方文档
• PyTorch Geometric教程