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-geometric

2. 导入必要的库

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

3. 加载数据集

我们使用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_index
        x = 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 acc

accuracy = 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_index
        x = 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 acc

accuracy = test()
print(f'Accuracy: {accuracy:.4f}')
'''

9. 结果分析

通过上述步骤，我们成功地使用PyG构建并训练了一个图卷积神经网络（GCN）模型。
在训练过程中，模型逐步学习图结构数据的特征，最终在测试集上取得了较好的分类准确率。
这展示了PyG在图数据处理和模型构建方面的强大功能。

10. 参考文献

• PyTorch Geometric官方文档
• PyTorch Geometric教程

通过本教程，您可以了解如何使用PyG进行图神经网络的构建和训练，为进一步的研究和应用奠定基础。