如果你正在使用图神经网络（GNN），并且在寻找加速模型的方法，那么这篇文章正适合你。

我们很高兴地宣布Graphcore拟未IPU上对PyTorch Geometric（PyG）的支持，并将在本文中引导你了解你的问题是否与GNN和IPU吻合，以及如何开始利用这个强大的组合。

请注意：本文假定你已经熟悉PyG和GNN的基本知识。如果你不熟悉这些话题，我们建议你从GNN的介绍博客和PyTorch Geometric的官方文档或教程开始了解。

IPU上的PyTorch Geometric

我们一直在努力确保PyTorch Geometric在拟未IPU上能够尽可能地无缝运行。我们最新发布的Poplar SDK 3.2包含了对PyG的扩展。该扩展被称为PopTorch Geometric，它支持用户在IPU上运行PyG。

这意味着你的模型将利用Poplar SDK中的高效收集和分散操作。这些操作（例如PyG的消息传递层）经常被用于GNN中。

我们目前的软件栈依赖于完整的模型编译，这与你可能熟悉的其他加速器中的即刻执行（Eager Execution）相反。即刻执行中，模型被切分成小内核，并按顺序执行。预编译支持我们应用优化，例如自动融合或并行化操作，这在按内核逐个部署时是不可能的。

此版本中，我们提供了数据加载器和transforms，使你能够便捷地静态编译你的PyG模型，支持你轻松获取完整模型编译所带来的优化。

最后，我们已经在IPU上验证了许多PyG层的功能，并正在继续扩大我们的覆盖范围。与Poplar SDK的整合也释放了其他许多强大功能，包括通过流水线和复制轻松扩展你的模型。

PopTorch Geometric library目前是实验性的。随着时间的推移，我们将继续添加更多来自上游PyG的功能。

有赖于Paperspace Gradient，现在使用PyG与Poplar SDK比以前更容易了。

Paperspace在PyG + Poplar SDK环境中提供了一个基于notebook的界面。这种用户友好的环境使开发人员和研究人员能够利用IPU的力量，而无需在硬件上进行前期投资。

对那些希望为其GNN工作负载尝试IPU能力的人来说，Paperspace Gradient以其随用随付的模式和为新用户提供的免费选项，成为了一个有吸引力的选项。

为了帮助用户入门，我们在运行时中包含了一系列的教程和示例，并将在本文中展开讨论。

欢迎你在Paperspace Gradient上探索用于IPU的PyG运行时。你也可以查看我们的博客，了解代码演练的示例。

为什么IPU擅长GNN？

IPU在GNN上表现十分出色，而IPU独特的架构（MIMD）在其中发挥了重要的作用。凭借大量的核心（每个处理器1472个）、独特的程序执行、分布式的处理器内存储、对具有不规则数据访问模式和控制流应用程序的高效执行，IPU在图学习任务中的表现超过了CPU（SISD）和GPU（SIMD）。

CPU在单线程性能和控制为主的代码方面表现出色，但受限于复杂核心数量少、存储层次深和时延降低技术等因素。

另一方面，GPU拥有较小的核心，在每台设备上的核心数量较多，它的架构也相对简单，但其SIMD执行模式和对常规、密集、数值、数据流为主的工作负载的关注，使其在GNN上没有那么高效。

IPU在处理收集和分散操作方面特别有效。这都是图学习中的常用操作，对信息传递至关重要。信息传递是GNN的一个关键特征，其中，节点特征通过非线性转换传递给相邻的节点，并在这些节点中聚合。

算术强度是一个用来描述算法或计算工作负载的计算效率的指标。它被定义为执行的浮点运算总数（FLOPs）与在存储和处理单元之间传输的数据总量（通常以字节为单位）的比率。

简言之，算术强度衡量每移动一个字节的数据需要进行的计算的次数。

低算术强度意味着每传输一个字节的数据都要执行少量的运算，这表明计算更受存储或带宽的限制。在这种情况下，算法的性能受到存储访问速度和可用存储带宽的约束。

提高处理能力可能不会带来实质性的性能提升，因为该算法仍然会受到存储子系统的瓶颈限制。

算术强度低是许多图机器学习任务的特点。

基准收集和分散操作

大量的片上SRAM和快速互连系统的结合使IPU非常适合处理受存储约束的工作负载。

为了证明这一点，我们在IPU上对一系列不同规模的收集和分散操作进行了基准测试，正如上文所讨论的，这让我们得以很好地了解信息传递的性能。

我们在张量大小的较低范围内观察到了明显的加速，而随着张量大小的增加，GPU的性能逐渐接近IPU。值得注意的是，在稀疏访问的情况下（即，小的分散输入尺寸），相同条件下IPU的性能是领先GPU的16倍以上。

值得强调的是，在许多情况下，不仅在IPU上运行基于图的工作负载速度更快，在云端租用拟未硬件的成本明显也更低。截至2023年第一季度，云平台Paperspace上A100-80GB的定价为3.18美元/小时，而Bow-IPU的价格为1.67美元/小时。

你可以在Paperspace上查看完整结果，并自己运行基准。

下图强调了拟未Bow IPU与NVIDIA A100（1个，蓝色平面）相比所提供的优势。最上面一行是关于scatter_add的，最下面一行则是不同大小的收集操作。左列显示了我们之前收集的全部性能数据，其中包括每对输入-输出的特征大小范围，作为散点图；而右边则显示了所有特征大小的平均速度提升。

分析IPU上的图学习任务

让我们从支持的特征和IPU上性能的角度来考虑以下几个图学习任务类别：图预测、节点预测和链接预测。

图层面的预测

图层面的预测是指预测单个图的属性。如果你的数据集是一个小型图的集合，那你可能会常常发现此类问题。

那么，IPU对图层面的预测任务有什么用？

通过用流水线将工作负载分配到多个设备上，小型图的问题可以轻易地扩展。这样的情况下，执行方案往往是许多与存储副本交织的小矩阵的混合体，这些操作在IPU上表现得特别好。

在此基础上，PopTorch Geometric提供的一套工具可以通过使用填充或打包，轻松对小型图进行批处理和加载，生产固定大小的批次，为IPU做好准备。

打包是指通过将图打包到任何本来会被填充的空间中，在一个固定大小的批次中充分利用任何浪费的空间，使得批次和计算更高效。

查看我们关于这些批处理方法的教程《教程Notebook：在IPU上使用填充进行小型图批处理》和《教程Notebook：在IPU上使用打包进行小型图批处理》。

我们还提供了一些GNN的示例，供你在Paperspace上运行。

其中一个示例是SchNet。美国能源部PNNL国家实验室在IPU上使用该模型进行分子性质预测，大大缩短了训练时间，提高了成本效率。下面的SchNet示例还展示了如何使用复制来在IPU上扩展你的PyG模型。

你也可以尝试用GIN网络来预测分子性质。GIN网络使用打包来充分利用小型分子图批次。

查看两个GNN的示例：

关于在IPU上成功完成图层级预测任务的进一步佐证，可参见我们的博客。

链接预测

链接预测处理的问题涉及预测图中节点之间的连接是否缺失或将来是否存在。链接预测任务的重要例子包括预测社交网络、知识图谱或推荐系统中的连接，通常在单个较大的图上运行。我们考虑了静态和动态（即图的结构或节点和边的特征随时间变化）的情况。

如何针对这些任务使用IPU？

与消息传递神经网络的其他应用一样，链接预测任务的特点往往是算术强度低、需要高存储带宽，因此IPU非常合适。

PyTorch Geometric与PopTorch的整合使得利用IPU完成这些任务变得更加轻松。

什么是时间图？

大多数图神经网络（GNN）的设计都假定是静态的图，但这种假设在许多现实世界的场景中可能并不成立，因为这些场景的基础系统是动态的，导致图随着时间的推移而演变。

社交网络或推荐系统等应用尤其如此。在这些应用中，用户与内容的互动会迅速改变。

为应对该挑战，研究人员已经提出了几种可以处理动态图的GNN架构，如作用于图结构不断演进的时间图网络（TGN）和处理动态变化的交通数据的时空图卷积网络。

在作用于动态变化的图形的应用中，经常可以观察到模型的准确性与用于训练和推理的批大小之间的依赖关系：批越大，模型的预测窗口中缺少的短期语境就越多，导致预测能力下降。

当批的规模较小时，大多数硬件平台会受到计算单元利用不足的影响，而IPU受到的影响则明显较小，这使得它们非常适合TGN。

你可以以Paperspace Gradient Notebook运行TGN训练。

如欲进一步了解时间图在IPU上的应用，可以参考我们的博客。

另外，你也可以参考我们的博客，了解新加坡国立大学是如何使用IPU的。

图形补全的其他方法

网络框架，结合了传统的基于路径的方法和现代图神经网络的优点。这种方法有助于归纳设置中的概括性和可解释性。NBFNet适用于同质图和知识图谱中的链接预测任务。

该模型现在可以作为Paperspace Gradient Notebook运行。

节点级预测

节点级的预测涉及到预测网络中一个或一组特定节点的属性。在很多情况下，这些任务涉及到大型图，因此往往需要采样技术来进行训练。

那么，为什么要在节点层级的任务中使用IPU？

在IPU上训练用于节点分类的大型图时，聚类是一种适合的采样方法。聚类涉及将大型图切割成多个子图，然后在训练模型（例如Cluster-GCN）时，重新组合子图并重新建立连接子图的原始边。

因为可以选择组合集群以适应IPU SRAM，从而在低算术强度状态下充分利用上述性能优势，这种方法非常适合IPU。

现在，使用PopTorch Geometric提供的新工具可以轻松地对具有适合IPU的静态大小的集群进行采样。这种采样方法可以很好地扩展，因为模型可以跨IPU复制，每个IPU采样一组不同的集群。你可以通过PopTorch轻松获取此功能，只需更改一些代码，即可以这种方式扩展模型。

现在就在Paperspace上尝试我们用集群采样进行节点分类的示例。

结论

IPU与PyTorch Geometric和Poplar SDK一起使用时，可为广泛的图学习任务提供强大的加速功能。

即刻免费体验，探索我们在Paperspace上的教程和示例，在你的应用程序中释放GNN的潜力。我们在本文中展示了IPU上的PyTorch Geometric，而这只是一个开始，我们还计划继续扩大我们的支持和示例。

如果您有任何的改进建议，欢迎加入我们的Slack频道。期待您的成果！