📚 课程目录第3章 PyTorch Geometric入门

3. PyTorch Geometric入门

在金融图数据建模中，我们经常面对的是“账户—交易—设备—IP”这类多实体、多关系的异构图。PyTorch Geometric（简称 PyG）正是为这类场景设计的图深度学习库。本节将围绕环境安装、核心数据结构、数据集加载、图批处理与基本操作五个方面，建立可落地的 PyG 使用基础。

一、核心原理：PyG 的数据抽象

PyG 的核心思想是把图拆解为“节点特征 + 边索引 + 边属性”三要素，并用统一的张量容器管理。对于同构图，使用 Data 对象；对于包含多种节点类型和边类型的异构图，使用 HeteroData 对象。理解这两个对象，是后续搭建 GNN 模型的前提。

Data 对象的关键属性包括：x（节点特征矩阵，形状为 [num_nodes, num_features]）、edge_index（边索引，形状为 [2, num_edges]）、edge_attr（边特征）、y（节点或图标签）。HeteroData 则通过字典式访问区分类型，例如 data['account'].x、data['account', 'transfers', 'account'].edge_index。

二、环境安装与版本校验

PyG 依赖 PyTorch 与 torch-scatter、torch-sparse 等扩展包。推荐使用 conda 或 pip 安装，并严格匹配 PyTorch 与 CUDA 版本。以下命令以 PyTorch 2.1 + CUDA 12.1 为例：

# 创建虚拟环境
conda create -n pyg_env python=3.10 -y
conda activate pyg_env

# 安装 PyTorch
pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# 安装 PyG 及其扩展
pip install torch_geometric
pip install torch_scatter torch_sparse torch_cluster -f https://data.pyg.org/whl/torch-2.1.0+cu121.html

# 校验
python -c "import torch, torch_geometric; print(torch.__version__, torch_geometric.__version__)"

若安装 torch-scatter 报错，优先检查 PyTorch 与 CUDA 版本是否与 whl 链接一致；CPU 环境可将链接替换为 torch-2.1.0+cpu.html。

三、Data 与 HeteroData 的基本操作

下面以金融反欺诈场景为例，构造一个包含账户与设备的异构图，并演示常用操作。

import torch
from torch_geometric.data import Data, HeteroData

# 同构图：账户转账网络
edge_index = torch.tensor([[0, 1, 2, 3],
                           [1, 2, 3, 0]], dtype=torch.long)
x = torch.randn(4, 16)          # 4 个账户，16 维特征
data = Data(x=x, edge_index=edge_index)
print(data)
print("节点数:", data.num_nodes, "边数:", data.num_edges)

# 异构图：账户 + 设备
hetero = HeteroData()
hetero['account'].x = torch.randn(4, 16)
hetero['device'].x = torch.randn(3, 8)
hetero['account', 'transfers', 'account'].edge_index = edge_index
hetero['device', 'login', 'account'].edge_index = torch.tensor(
    [[0, 1, 2], [0, 1, 2]], dtype=torch.long)
print(hetero)
print("节点类型:", hetero.node_types)
print("边类型:", hetero.edge_types)

常用操作清单如下：

操作	Data 写法	HeteroData 写法
查看节点数	`data.num_nodes`	`hetero['account'].num_nodes`
查看边数	`data.num_edges`	`hetero['account','transfers','account'].num_edges`
添加自环	`data = T.AddSelfLoops()(data)`	对指定边类型分别调用
划分训练集	`data.train_mask`	按节点类型设置 mask
转无向图	`T.ToUndirected()(data)`	对目标边类型调用

四、数据集加载与图批处理

PyG 内置了 Planetoid、Amazon、OGB 等数据集，金融场景也可自定义 InMemoryDataset。加载后通常需要做标准化、划分掩码等预处理。

from torch_geometric.datasets import Planetoid
from torch_geometric.loader import DataLoader

dataset = Planetoid(root='./data/Cora', name='Cora')
data = dataset[0]
print(dataset.num_classes, dataset.num_node_features)

# 图批处理：将多个小图合并为一个大图
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
    print(batch.batch)   # 标记每个节点属于哪张图
    break

批处理的核心是 batch 向量：PyG 通过它区分不同图，并在池化时按图聚合。对于异构图，需使用 HeteroDataLoader，其 batch 中每种节点类型都有独立的 batch 向量。

常见误区：直接把多个 Data 对象拼成列表送入模型，会导致 batch 维度混乱。务必使用 DataLoader 或 HeteroDataLoader，并检查 batch.batch 的取值是否从 0 开始连续编号。

五、复盘与建议

本节的关键在于建立“数据对象—加载—批处理”的完整链路。建议在真实金融数据上先做小规模验证：用 Data 跑通同构图，再迁移到 HeteroData 表达多实体关系。遇到维度不匹配时，优先检查 edge_index 的 dtype 是否为 torch.long，以及节点特征行数是否与节点数一致。掌握这些基础操作后，后续的 GCN、GraphSAGE 与 HeteroGNN 建模才能顺利展开。