PyTorch 稀疏矩阵乘法与 graph message passing¶
PyTorch 通过 torch.sparse 提供稀疏矩阵运算,
Tiga 则通过沿图边的 message passing
描述计算。对于邻居加权求和,两者表达的是同一个运算;一般的自定义边函数则不一定能写成矩阵乘法。
稀疏矩阵如何表示图¶
在 compressed sparse row (CSR)格式中,每一行表示一个目标节点,存储的列索引表示源节点,对应的数值表示边权。 源节点 j 到目标节点 i 的边贡献如下:
当节点携带特征向量时,这就是稀疏矩阵与稠密矩阵乘法(SpMM)。
torch.sparse.mm(adjacency, features) 与 Tiga 中的 edge.weight * src.x
消息加 sum reducer 都表达该加权和。
图编程模型通过示意图解释连接关系;下例明确采用“行是目标、列是源”的约定。
GPU 示例:同时核对前向结果与梯度¶
先完成安装,包括与 CUDA 环境匹配的 Torch。 下例为两个实现创建独立的输入,核对输出、节点梯度和边权梯度。这是正确性检查,不是性能测试。
运行 python examples/pytorch_sparse_message_passing.py --device cuda。
同一份源码也支持 --device cpu。
"""Compare weighted graph aggregation with torch.sparse.mm, including gradients."""
import argparse
import torch
import tiga as tg
class WeightedAggregation(tg.MessagePassing):
reducer = tg.sum()
def edge(self, src, dst, edge):
return edge.weight * src.x
def run(device="cuda"):
# CSR rows are destinations; column indices identify source nodes.
row = torch.tensor([0, 2, 3, 5], dtype=torch.int64, device=device)
col = torch.tensor([0, 2, 1, 0, 1], dtype=torch.int64, device=device)
x = torch.tensor([1., 2., 3.], device=device, requires_grad=True)
weight = torch.tensor([2., 3., 4., 5., 6.], device=device,
requires_grad=True)
graph = tg.Graph.from_csr(row, col, num_src=3)
result = WeightedAggregation()(
graph=graph, src={"x": x}, dst={}, edge={"weight": weight})
# Independent leaves make the two gradient computations independent.
x_ref = x.detach().clone().requires_grad_()
weight_ref = weight.detach().clone().requires_grad_()
adjacency = torch.sparse_csr_tensor(
row, col, weight_ref, size=(3, 3), device=device, check_invariants=True)
reference = torch.sparse.mm(adjacency, x_ref[:, None]).squeeze(1)
torch.testing.assert_close(result, reference)
gradients = torch.autograd.grad(result.sum(), (x, weight))
reference_gradients = torch.autograd.grad(
reference.sum(), (x_ref, weight_ref))
for actual, expected in zip(gradients, reference_gradients):
torch.testing.assert_close(actual, expected)
print("Output:", result.detach().tolist()) # [11.0, 8.0, 17.0]
print("Node gradients:", gradients[0].tolist()) # [7.0, 10.0, 3.0]
print("Edge gradients:", gradients[1].tolist()) # [1., 3., 2., 1., 2.]
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
run(parser.parse_args().device)
两种实现都返回普通 torch.Tensor。五条存储边各自拥有独立的边权。
在这个 sum 运算中,没有入边的行输出零;相同源、目标之间的重复边按每条存储边分别累加。
其他 reducer 的空邻域规则需要单独检查。
Torch 的稀疏格式、设备及梯度支持以
torch.sparse.mm 文档为准;
Tiga 的覆盖范围见支持矩阵。
与图神经网络有什么关系¶
图神经网络(graph neural network,GNN) 把图聚合与特征变换、非线性运算组合起来。上面的加权和只是聚合模块, 不等于完整、已归一化的 graph convolutional network(GCN)层。 自环、度归一化、可学习变换与训练目标仍由具体模型定义。 参见多特征 GCN 形状的聚合示例 与 Torch 模型集成。
PyTorch Geometric(PyG)提供图学习层与训练工具;Tiga 聚焦于编译图程序, 包括自定义边消息及其已支持的梯度,不是所有 PyG 层的直接替代品。 同类对比进一步解释这些抽象边界。
不止固定的稀疏邻接矩阵¶
Radius 与 k-nearest-neighbor 图根据坐标或特征确定邻居,stencil 根据网格规则确定邻居。 动态关系示例介绍这些构造器、距离度量与边界条件。 执行时是物化边,还是在计算中生成候选,取决于关系与所选 backend。
性能比较需要固定图、特征、dtype、设备与梯度需求,并区分建图、编译、预热后执行和内存开销。
性能报告提供稀疏聚合、radius/kNN、大图与分布式案例;
这些结果不代表 Tiga 对任意图神经网络或 torch.sparse 运算都更快。