iSpLib：使用自动调优稀疏操作加速图神经网络的库
机器学习
2024-03-25 v1 分布式、并行与集群计算
性能
摘要
图神经网络 (GNN) 训练和推理中的核心计算通常被映射为稀疏矩阵操作，例如稀疏-稠密矩阵乘法。这些稀疏操作难以通过手动调优来优化，因为其性能显著依赖于输入图的稀疏性、GNN 模型和计算平台。为了应对这一挑战，我们提出了 iSpLib，这是一个基于 PyTorch 的 C++ 库，配备了自动调优的稀疏操作。iSpLib 通过启用缓存的反向传播加速 GNN 训练，该机制将中间矩阵存储在本地缓存中。该库提供了一个用户友好的 Python 插件，允许用户开箱即用地利用我们优化的 PyTorch 操作，只需增加两行代码即可用于任何基于线性代数的流行 GNN（图卷积网络、GraphSAGE、图推理网络等）的现有 PyTorch 实现。我们证明，在 CPU 上，与等效的 PyTorch 2.1.0 和 PyTorch Geometric 2.4.0 实现相比，iSpLib 可获得高达 27 倍的总体训练加速。我们的库已在 https://github.com/HipGraph/iSpLib (https://doi.org/10.5281/zenodo.10806511) 公开提供。
引用
@article{arxiv.2403.14853,
title = {iSpLib: A Library for Accelerating Graph Neural Networks using Auto-tuned Sparse Operations},
author = {Md Saidul Hoque Anik and Pranav Badhe and Rohit Gampa and Ariful Azad},
journal= {arXiv preprint arXiv:2403.14853},
year = {2024}
}