在密集硬件上快速训练稀疏图神经网络
机器学习
2019-06-28 v1 机器学习
摘要
近年来,图神经网络因其能够自然地编码关系型输入数据,并通过在图的邻接矩阵稀疏表示上进行运算而扩展到大规模图的能力,变得越来越流行。当我们考虑使用定制硬件来扩展这些模型时,一个自然的假设是,我们需要为稀疏操作和/或动态控制流量身定制的硬件。在这项工作中,我们通过使用一个针对固定大小数据的密集计算平台来扩展稀疏图神经网络,从而对这一假设提出了质疑。我们从稀疏矩阵数值算法的优化中汲取灵感,开发了一些技术,使得能够在 13 分钟内使用 512 核 TPUv2 Pod 训练来自 Allamanis 等人 [2018] 的稀疏图神经网络模型,而原始训练几乎需要一天时间。
引用
@article{arxiv.1906.11786,
title = {Fast Training of Sparse Graph Neural Networks on Dense Hardware},
author = {Matej Balog and Bart van Merriënboer and Subhodeep Moitra and Yujia Li and Daniel Tarlow},
journal= {arXiv preprint arXiv:1906.11786},
year = {2019}
}