4、PyTorch Geometric (PyG) 入门：PyG环境安装与配置、Data对象详解、Dataset与DataLoader、内置数据集加载与可视化

好，咱们进入实战环节了。

前面聊了那么多图神经网络的理论，现在终于要动手了。我个人觉得，PyG（PyTorch Geometric）是目前最顺手的图神经网络框架。它把图数据封装得特别优雅，你不需要手动去搞那些邻接矩阵的稀疏运算，框架全给你包了。

这一节，咱们就把它装好，然后彻底搞懂它的核心数据结构——Data对象。嗯，这是所有后续工作的基石。

4.1 环境安装与配置

安装PyG其实是个体力活，主要坑在版本依赖上。我刚开始搞的时候，因为CUDA版本没对齐，折腾了一下午。你想想看，一个ImportError弹出来，说找不到某个动态链接库，多闹心。

我的建议是：先确认你的PyTorch版本，再装PyG。

先看PyTorch版本：

import torch
print(torch.__version__)  # 比如 2.1.0
print(torch.version.cuda) # 比如 11.8

然后去PyG官网找对应的安装命令。我一般用pip直接装，干净利落：

# 无CUDA版本（CPU only）
pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.1.0+cpu.html

# 有CUDA版本（以CUDA 11.8为例）
pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.1.0+cu118.html

装完之后，验证一下：

from torch_geometric.data import Data
print("PyG 安装成功！")

没报错？恭喜，环境搞定了。

4.2 Data对象详解

PyG里最核心的东西，就是这个Data对象。说白了，它就是一张图的容器。

一张图有什么？节点、边、还有节点和边上的特征。在PyG里，这些都变成了Data对象的属性：

属性名	类型	含义
x	Tensor (N, F)	节点特征矩阵，N个节点，每个节点F维特征
edge_index	LongTensor (2, E)	边索引，E条边，用两个节点编号表示一条边
edge_attr	Tensor (E, D)	边特征矩阵（可选）
y	Tensor	标签（节点级、图级或边级）
pos	Tensor (N, 2/3)	节点坐标（可视化用）

来，咱们手搓一个最简单的图：

import torch
from torch_geometric.data import Data

# 3个节点，每个节点有2维特征
x = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.float)

# 边：0-1, 1-2, 2-0（无向图，所以每条边要写两次）
edge_index = torch.tensor([[0, 1, 1, 2, 2, 0],
                           [1, 0, 2, 1, 0, 2]], dtype=torch.long)

# 图级标签
y = torch.tensor([1], dtype=torch.float)

data = Data(x=x, edge_index=edge_index, y=y)
print(data)
# 输出: Data(x=[3, 2], edge_index=[2, 6], y=[1])

注意这个edge_index的格式。它是[2, 边数]的矩阵，第一行是源节点，第二行是目标节点。我刚开始学的时候，老把它搞成[边数, 2]，结果模型怎么训都不收敛。后来才发现是数据格式错了。

4.3 Dataset与DataLoader

单个图好办，但实际项目中我们往往有成百上千张图。这时候就需要Dataset和DataLoader了。

PyG提供了两种数据集基类：

• InMemoryDataset：所有图一次性加载到内存。适合小数据集，比如Cora、PubMed。
• Dataset：图数据太大，不能全放内存，需要按需加载。适合大规模图数据。

咱们先看最简单的用法——加载内置数据集：

from torch_geometric.datasets import TUDataset

# 加载MUTAG数据集（包含188个分子图）
dataset = TUDataset(root='/tmp/MUTAG', name='MUTAG')

print(f'数据集大小: {len(dataset)}')  # 188
print(f'节点特征维度: {dataset.num_node_features}')  # 7
print(f'类别数: {dataset.num_classes}')  # 2

# 取第一张图
data = dataset[0]
print(data)
# 输出: Data(edge_index=[2, 38], x=[17, 7], edge_attr=[38, 4], y=[1])

然后配合DataLoader做批处理：

from torch_geometric.loader import DataLoader

loader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in loader:
    print(batch)
    # 输出: DataBatch(x=[512, 7], edge_index=[2, 1024], y=[32], batch=[512])
    # 注意这里 batch 属性记录了每个节点属于哪张图
    break

这里有个关键点：batch属性。当你把多张图拼成一个batch时，PyG会自动给每个节点打上「图编号」。这样模型就知道哪些节点属于同一张图，方便做图级别的池化操作。

4.4 内置数据集加载与可视化

PyG内置了很多经典数据集，比如：

• 节点分类：Cora、Citeseer、PubMed（论文引用网络）
• 图分类：MUTAG、PROTEINS、IMDB-BINARY（分子图、蛋白质图等）
• 链接预测：WN18RR、FB15k-237（知识图谱）

咱们用Cora数据集做个可视化，看看图长什么样：

from torch_geometric.datasets import Planetoid
import matplotlib.pyplot as plt
import networkx as nx

# 加载Cora数据集
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

# 转换成networkx图
from torch_geometric.utils import to_networkx
G = to_networkx(data, to_undirected=True)

# 可视化
plt.figure(figsize=(10, 10))
nx.draw(G, node_size=20, node_color=data.y, cmap='Set2')
plt.show()

运行这段代码，你会看到一张论文引用网络图。不同颜色代表不同类别的论文。嗯，这就是我们后面要训练的数据。

我个人习惯是，拿到一个新数据集，先可视化看看。有时候数据本身就有问题，比如孤立节点太多、标签分布不均。可视化一眼就能看出来，比跑模型发现不收敛再回头查数据要高效得多。

好，PyG的入门就到这儿。环境装好了，数据对象搞懂了，数据集也能加载了。接下来咱们就可以真正开始训练图神经网络了。

---

---