第四章 图的表示与存储:邻接表、邻接矩阵、稀疏表示、DGL/PyG 数据结构
做药物节点分类,第一步不是调模型,而是想清楚「图怎么存」。我见过不少同学把分子图塞进 GNN 之前,连邻接矩阵占多少内存都没算过,结果 10 万个药物分子还没训就把显存撑爆了。说白了,图的存储方式直接决定了你能处理多大规模的分子库、训练能跑多快。
这一章我们就掰开揉碎讲清楚四种主流存储方式——邻接矩阵、邻接表、稀疏表示,以及两个工业级框架 DGL 和 PyG 的数据结构。
4.1 邻接矩阵——最直觉,但也最“笨重”
邻接矩阵很好理解:假设图有 N 个节点,就建一个 N×N 的矩阵。第 i 行第 j 列是 1 就表示 i→j 有边,是 0 就没边。对于无向图,矩阵是对称的。
import numpy as np
# 4 个节点,边: 0→1, 1→2, 2→3, 3→0
N = 4
adj = np.zeros((N, N), dtype=np.int8)
edges = [(0,1), (1,2), (2,3), (3,0)]
for i, j in edges:
adj[i, j] = 1
adj[j, i] = 1 # 无向图
print(adj)
代码跑完你就会发现:N=4 时矩阵只有 16 个元素,很轻量。但换成 N=5000,矩阵就膨胀到 2500 万——差不多 200 MB。一个药物分子数据集动辄几万到十几万个分子,每个分子几百个原子(节点),用邻接矩阵存?显存根本扛不住。
我参与过一个激酶抑制剂的筛选项目,当时团队里有人坚持用邻接矩阵 + 全连接层做基线,结果 16GB 的 V100 连一个 batch 都塞不下。后来换成稀疏格式,同样 batch size 内存直接降到 1/20。
4.2 邻接表——省内存,遍历友好
邻接表只存“有边”的信息。每个节点维护一个邻居列表。对于稀疏图(药物分子图几乎全是稀疏图),邻接表比邻接矩阵省太多。
# 用字典实现邻接表
adj_list = {i: [] for i in range(N)}
for i, j in edges:
adj_list[i].append(j)
adj_list[j].append(i)
print(adj_list[0]) # 节点 0 的邻居
你想想看,药物分子图平均每个原子连 2–4 个键,用邻接表存储只需要 O(2×E) 的空间,E 是边数。对于 N=5000、E≈10000 的图,邻接表只要几千个整数,而邻接矩阵要 2500 万个。
但邻接表也有短板——查两个节点之间是否有边,得遍历邻居列表,最坏情况 O(N)。如果频繁做边查询,邻接矩阵 O(1) 的优势就体现出来了。
| 存储方式 | 空间复杂度 | 查边耗时 | 遍历邻居耗时 | 适合场景 |
|---|---|---|---|---|
| 邻接矩阵 | O(N²) | O(1) | O(N) | 稠密小图 |
| 邻接表 | O(N+E) | O(度) | O(度) | 稀疏图、GNN 消息传递 |
| 稀疏表示 (CSR) | O(N+E) | O(logN) | O(度) | 大规模图、GPU 加速 |
4.3 稀疏表示——大规模图的“工业标准”
实际做药物节点分类时,我最常用的是 Compressed Sparse Row (CSR) 格式。DGL 和 PyG 底层用的就是这种思路。
CSR 用三个数组表示图:
- row_ptr:长度 N+1,记录每个节点的邻居在 col_ind 中的起始偏移
- col_ind:长度 2E(无向图),存储所有邻居的索引
- val:可选,存储边权重
# 手动构造 CSR
row_ptr = [0, 2, 5, 7, 9] # 4 个节点,共 9 个邻居条目(无向图)
col_ind = [1, 3, 0, 2, 3, 1, 3, 0, 1]
# 查节点 2 的邻居:
start = row_ptr[2] # 5
end = row_ptr[3] # 7
print(col_ind[start:end]) # [1, 3]
我曾经用 CSR 格式处理过一个包含 200 万药物分子、每个分子平均 50 个原子的数据集。如果用邻接矩阵,需要 200 万 × 50² ≈ 50 亿个元素,根本不可能。但 CSR 只存了约 5 亿个整数,配合稀疏矩阵运算库(scipy.sparse、cuSPARSE),训练效率非常高。
scipy.sparse.csr_matrix 构建,然后转成 DGL/PyG 的图对象。别自己手撸 CSR 数组,容易下标越界。
4.4 DGL 的数据结构——从边列表到 DGLGraph
DGL(Deep Graph Library)把图封装成 DGLGraph 对象。我习惯用 dgl.graph() 直接从边列表构建。
import dgl
import torch
# 边列表:source 节点和 target 节点
src = torch.tensor([0, 1, 2, 3])
dst = torch.tensor([1, 2, 3, 0])
g = dgl.graph((src, dst))
# 添加节点特征
g.ndata['h'] = torch.randn(4, 64) # 每个节点 64 维特征
g.edata['w'] = torch.randn(4, 1) # 每条边 1 维权重
print(g)
DGL 内部自动选择最优存储——小图用邻接表,大图自动转为 CSR。我做项目时基本不操心底层格式,但有一点要注意:DGL 默认是单向边,处理无向分子图记得加 bidirectional=True。
g = dgl.graph((src, dst), bidirectional=True) # 自动补反向边
g.adjacency_matrix() 可以随时导出稀疏邻接矩阵(scipy.sparse),方便跟传统方法做对比。
4.5 PyG 的数据结构——Data 对象的“乐高”哲学
PyG(PyTorch Geometric)的核心是 torch_geometric.data.Data 对象。它不像 DGL 那样强调“图对象”,而是把图拆成 edge_index、x、edge_attr 等组件。
from torch_geometric.data import Data
import torch
edge_index = torch.tensor([
[0, 1, 2, 3],
[1, 2, 3, 0]
], dtype=torch.long)
x = torch.randn(4, 64) # 节点特征
edge_attr = torch.randn(4, 1) # 边特征
data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
print(data)
PyG 的 edge_index 形状是 2×E,第一行是源节点,第二行是目标节点。我一开始老记混,后来养成习惯:先写源,再写目标,跟 DGL 的 (src, dst) 顺序一致。
torch_geometric.utils.to_undirected() 转换。
from torch_geometric.utils import to_undirected
edge_index = to_undirected(edge_index) # 自动添加反向边
4.6 框架对比——选 DGL 还是 PyG?
两个框架我都深度用过。说实话,没有绝对的好坏,看你的场景:
| 对比维度 | DGL | PyG |
|---|---|---|
| 底层存储 | 自动 CSR / COO 切换 | 固定 edge_index (COO) |
| 异构图支持 | DGLHeteroGraph,API 一致 |
HeteroData,字典风格 |
| 消息传递写法 | update_all + 自定义函数 |
MessagePassing 基类 + propagate |
| 与 PyTorch 集成 | 需要额外 dgl.data 适配 |
原生兼容,DataLoader 即用 |
| 大规模图推理 | 内置 dgl.dataloading 采样 |
需 NeighborLoader 或手动采样 |
我个人习惯:如果做异构图(比如药物-靶点-疾病三元图),优先选 DGL,它的异构图 API 更统一。如果是纯节点分类或分子图分类,PyG 的 Data 对象更轻量、调试更直观。
4.7 一张图看懂图的存储体系
下面这张图梳理了从原始分子图到四种存储方式的脉络,以及 DGL/PyG 在其中的位置。建议你把它当作「图存储选型速查表」。
从左往右看,存储方式的「紧凑度」逐渐增加,但查询边的便利性逐渐下降。DGL 和 PyG 在顶层都帮你做了权衡——你只需要关注图的结构和特征,底层用 CSR 还是邻接表,框架帮你选。