第四章 图的表示与存储:邻接表、邻接矩阵、稀疏表示、DGL/PyG 数据结构

做药物节点分类,第一步不是调模型,而是想清楚「图怎么存」。我见过不少同学把分子图塞进 GNN 之前,连邻接矩阵占多少内存都没算过,结果 10 万个药物分子还没训就把显存撑爆了。说白了,图的存储方式直接决定了你能处理多大规模的分子库、训练能跑多快。

这一章我们就掰开揉碎讲清楚四种主流存储方式——邻接矩阵、邻接表、稀疏表示,以及两个工业级框架 DGL 和 PyG 的数据结构。

4.1 邻接矩阵——最直觉,但也最“笨重”

邻接矩阵很好理解:假设图有 N 个节点,就建一个 N×N 的矩阵。第 i 行第 j 列是 1 就表示 i→j 有边,是 0 就没边。对于无向图,矩阵是对称的。

适用场景:小图(节点数 < 1 万)、稠密图、需要快速随机访问边的场景。
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。

避坑提醒:邻接矩阵的存储复杂度是 O(N²)。对于药物分子图(通常 N < 200),全量存储浪费不大;但对于蛋白质-配体复合图或超大规模虚拟筛选,千万别用。

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),训练效率非常高。

实用技巧:在 Python 里可以直接用 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)  # 自动补反向边
惯用操作:DGL 的 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) 顺序一致。

注意:PyG 默认不自动补无向边。如果做分子图分类,需要手动把边列表镜像一份,或者用 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 在其中的位置。建议你把它当作「图存储选型速查表」。

图的表示与存储方式 图的表示与存储方式 原始药物分子图 邻接矩阵 邻接表 稀疏表示 (CSR) DGLGraph PyG Data 根据图规模 & 任务选择存储格式 小图用邻接矩阵,大图用 CSR,框架封装了细节

从左往右看,存储方式的「紧凑度」逐渐增加,但查询边的便利性逐渐下降。DGL 和 PyG 在顶层都帮你做了权衡——你只需要关注图的结构和特征,底层用 CSR 还是邻接表,框架帮你选。

一句话总结:写药物节点分类的代码时,直接用 DGL 或 PyG 的高层 API,别自己手搓存储结构。但理解了邻接矩阵、邻接表、CSR 的差异,你才能诊断性能瓶颈、调优数据加载。这一章的知识,说白了就是帮你打好图存储的地基。