在我们的世界中，物体及其相互关系无处不在，理解一个物体，其关系往往和单独的属性一样重要，例如，如交通网络、生产网络，知识图谱，或是社交网络等。离散数学和计算机科学长期以来一直在将这些网络以图的形式进行建模，图以多种复杂方式互连的节点和边构成。但是，大多数机器学习（ML）算法仅支持输入对象之间规则和统一的关系，比如像素的网格、一系列的单词，或者完全没有关系。

图神经网络（GNN）已经成为一种强大的技术，不仅可以利用图的连接性（如老旧算法 DeepWalk 和 Node2Vec 所做的那样），还能利用各节点和边上的输入特征。GNN 能够对整个图（例如，这个分子是否以特定方式反应？）、图中的单个节点（这篇文档的主题是什么，给出它的引用来源？）或是潜在的边（这个产品是否可能与那个产品一起被购买？）进行预测。除了对图进行预测，GNN 也是强大的工具，可以帮助过渡到更常见的神经网络应用场景。它们将图的离散、关系信息编码为连续形式，使得这些信息可以自然地融入到其他深度学习系统中。

现在，我们非常高兴地宣布 TensorFlow GNN 1.0（TF-GNN）发布了，这是一个经过生产级测试的库，专为大规模构建 GNN 而设计。它不仅支持在 TensorFlow 中进行模型构建和训练，还支持从庞大的数据存储中提取输入图。TF-GNN 从基础上为处理异构图而构建，这些图中的对象和关系类型通过不同的节点和边集合来表示。在现实世界中，不同类型的对象及其关系普遍存在，TF-GNN 的设计重点放在异构图上，使得这种多样性的表示成为可能。

在 TensorFlow 内部，图结构是通过一种名为 tfgnn.GraphTensor 的对象表示的。这是一个复合张量类型（即一个 Python 类中包含的一组张量），它在 tf.data.Dataset、tf.function 等中作为一等公民被接受。它不仅存储了图的结构，还存储了附加在节点、边缘以及整个图上的特征。GraphTensors 的可训练转换可以在高级 Keras API 中定义为层（Layers）对象，或者直接使用 tfgnn.GraphTensor 基础设施来定义。

#01 GNN：在上下文中对对象进行预测

以 TF-GNN 的一个典型应用为例：预测一个庞大数据库的交叉引用表格中定义的图上某个类型节点的属性。举个例子，就像是计算机科学（CS）领域 arXiv 论文的引用数据库，它有着一对多的引用以及多对一的被引用关系，我们希望能预测每篇论文所属的研究领域。

像大多数神经网络一样，GNN 在大量标记样本（大约数百万）的数据集上进行训练，但每次训练只涉及一小批训练样本（例如，数百个）。为了能够处理数百万级别的数据，GNN 通过从底层图中不断取出相对较小的子图流进行训练。每个子图包含了足够多的原始数据，能够计算出中心标记节点的 GNN 结果，并据此训练模型。这一过程 —— 通常称为子图采样 —— 对于 GNN 的训练至关重要。大多数现有的工具都是以批处理方式进行采样，生成静态的子图用于训练。TF-GNN 则提供了一种通过动态和交互式方式进行采样的改进工具。

TF-GNN 1.0 推出了一个灵活的 Python API，允许用户在所有相关的规模上配置动态或批处理子图采样：无论是在 Colab 笔记本中进行交互式操作，高效采样存储在单个训练主机主内存中的小型数据集，还是通过 Apache Beam 分布式处理存储在网络文件系统中的大型数据集（高达数亿个节点和数十亿条边）。更多细节，请参见：

内存采样用户指南：https://github.com/tensorflow/gnn/blob/main/tensorflow_gnn/docs/guide/inmemory_sampler.md

基于 beam 的采样用户指南：https://github.com/tensorflow/gnn/blob/main/tensorflow_gnn/docs/guide/beam_sampler.md

在这些被采样的子图上，GNN 的任务是计算根节点处的一个隐藏（或潜在）状态；这个隐藏状态会聚合并编码根节点周围邻域的相关信息。一个典型的方法是消息传递神经网络。在每一轮消息传递中，节点通过传入的边缘接收来自邻居的消息，并基于这些消息更新自己的隐藏状态。经过 n 轮更新后（如下图所示，n = 2），根节点的隐藏状态就能反映出所有在 n 跳边内节点的综合信息。这些消息和新的隐藏状态是由神经网络的隐藏层计算出来的。在一个多元化的图中，为不同类型的节点和边缘使用专门训练的隐藏层往往是有益的。

TensorFlow GNN 1.0（TF-GNN）_神经网络

如图所示，一个简易的消息传递神经网络过程中，节点状态会从外围节点向中心节点传递，并在此过程中聚合这些信息以形成新的节点状态。当这一过程触及到根节点时，我们便能根据累积的信息做出最终的预测。

训练过程通过在标记节点的 GNN 隐藏状态之上增加一个输出层来实现，这个输出层负责计算损失值（用来评估预测的准确性），并通过标准的反向传播方法更新模型权重，这与常规的神经网络训练流程相同。

除了标准的监督式训练 —— 即通过最小化标签定义的损失来训练 ——GNN 也可以进行非监督式训练，即在没有标签的情况下进行。这种训练方式允许我们为节点及其特征的离散图结构创建一个连续的表示（或称为嵌入），这些嵌入表示随后可以在其他机器学习系统中使用。这样，图中编码的离散和关系信息就可以融入到更常见的神经网络应用场景中。TF-GNN 支持为异构图定制非监督学习目标的细节。

#02 构建 GNN 架构

TF-GNN 库提供了多种级别的支持，以构建和训练 GNN。

在最高层次，用户可以选择库中预先定义的模型，这些模型是以 Keras 层的形式表达的。除了一些研究文献中的模型外，TF-GNN 还提供了一个高度可配置的模型模板，这个模板集成了一系列我们发现在众多内部问题上能提供强有力基准的建模选项。这些模板实现了 GNN 的各个层，用户只需负责初始化这些 Keras 层即可。 

TensorFlow GNN 1.0（TF-GNN）_tensorflow_02

在最基础层面，用户可以完全从零开始构建 GNN 模型，利用图数据传递的基础操作，如从一个节点向所有外连边广播数据，或者将所有进入边的数据汇总到一个节点（比如，累加所有传入的信息）。TF-GNN 的图数据模型平等对待节点、边缘以及整个输入图的特征或隐藏状态，这让构建从节点中心模型（如前文讨论的 MPNN）到更泛化的 GraphNets 都变得简单直接。这种模型构建既可以在 TensorFlow 的核心上直接使用 Keras 作为建模框架，也可以不用。 

想了解更多详细信息和中级建模方法，请参考：

TF-GNN 用户指南：https://github.com/tensorflow/gnn/blob/main/tensorflow_gnn/docs/guide/gnn_modeling.md

模型集合：https://github.com/tensorflow/gnn/tree/main/tensorflow_gnn/models

#03 训练编排

尽管高级用户可以自行定制模型训练，但 TF-GNN Runner 提供了一种高效编排 Keras 模型训练的简洁方法。一个典型的调用示例如下：

TensorFlow GNN 1.0（TF-GNN）_人工智能_03

Runner 为机器学习的常见难题，比如分布式训练和在 Cloud TPU 上针对固定形状进行 tfgnn。GraphTensor 填充，提供了即插即用的解决方案。除了单任务训练外，它还支持同时针对多个（两个或更多）任务的联合训练。比如，可以将非监督任务与监督任务结合起来，共同优化最终的连续表示（或嵌入），这些表示融入了特定应用的归纳偏见。调用者仅需将任务参数换成任务映射即可：

TensorFlow GNN 1.0（TF-GNN）_子图_04

此外，TF-GNN Runner 还实现了集成梯度方法，用于模型的归因分析。集成梯度的输出是一个 GraphTensor，其连接性与观测到的 GraphTensor 相同，但特征被梯度值替换，这些梯度值中的较大值对 GNN 的预测贡献更大。用户可以通过检查梯度值来识别他们的 GNN 最依赖哪些特征。

#04 结论

总之，我们希望 TF-GNN 能够帮助推动 GNN 在 TensorFlow 中的大规模应用，并进一步促进该领域的创新。如果你想深入了解，请尝试我们基于流行的 OGBN-MAG 基准的 Colab 演示（无需安装，直接在浏览器中运行）。

用户指南和 Colab 演示：https://github.com/tensorflow/gnn/blob/main/tensorflow_gnn/docs/guide/overview.md

论文：https://arxiv.org/abs/2207.03522

原文链接： https://blog.research.google/2024/02/graph-neural-networks-in-tensorflow.html