GeoT：通过 GPU 高效段归约实现的图神经网络以张量为中心的库
分布式、并行与集群计算
2024-04-09 v2 机器学习
摘要
近年来，图神经网络 (GNNs) 引发了创新的浪潮，显著增强了对图、点云和网格等几何数据结构的处理。随着该领域的不断发展，一系列框架和库正在被开发，以将 GNN 的效率推向新高度。尽管以图为中心的库在过去取得了成功，但高效张量编译器的出现凸显了对以张量为中心的库的迫切需求。然而，由于在 GNN 环境中实现段归约时遇到独特的挑战和限制，用于 GNN 的高效以张量为中心的框架仍然稀缺。我们介绍了 GeoT，这是一个专门为 GNN 设计的尖端以张量为中心的库，通过高效的段归约实现。GeoT 首次推出了创新的并行算法，不仅引入了新的设计原则，还扩展了可用的设计空间。重要的是，GeoT 的设计旨在计算图内实现直接融合，确保与当代以张量为中心的机器学习框架和编译器的兼容性。GeoT 设定了新的性能基准，通过展示 1.80 倍的平均算子加速和 1.68 倍的端到端加速，标志着显著的进步。
引用
@article{arxiv.2404.03019,
title = {GeoT: Tensor Centric Library for Graph Neural Network via Efficient Segment Reduction on GPU},
author = {Zhongming Yu and Genghan Zhang and Hanxian Huang and Xin Chen and Jishen Zhao},
journal= {arXiv preprint arXiv:2404.03019},
year = {2024}
}