物体及其关系在我们周围的世界中无处不在，关系对于理解一个物体可能与其自身属性同等重要——例如：交通网络、生产网络、知识图谱或社交网络。离散数学和计算机科学长期以来一直将此类网络形式化为图，由节点通过边以各种不规则方式任意连接而成。然而，大多数机器学习（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 等中被接受为一等公民。它同时存储图结构及其附加到节点、边和整个图上的特征。GraphTensor 的可训练变换可以定义为高级 Keras API（应用程序编程接口） 中的 Layers 对象，或直接使用 tfgnn.GraphTensor 原语。

GNN：为上下文中的对象进行预测

为了说明，让我们看一个 TF-GNN 的典型应用：预测由大型数据库的交叉引用表所定义的图中某一类型节点的属性。例如，一个计算机科学（CS）arXiv 论文的引用数据库，具有一对多的引用和多对一的被引用关系，我们希望预测每篇论文的主题领域。

像大多数神经网络一样，GNN 在包含许多标记示例（约数百万个）的数据集上进行训练，但每个训练步骤只包含一个小得多的训练示例批次（比如数百个）。为了扩展到数百万个，GNN 在来自底层图的一系列合理大小的子图上进行训练。每个子图包含足够的原始数据，以计算其中心标记节点的 GNN 结果并训练模型。这个过程——通常称为子图采样——对 GNN 训练极为重要。大多数现有工具以批处理方式完成采样，为训练生成静态子图。TF-GNN 提供了工具，通过动态和交互式采样来改进这一点。

|  |
| --- |
| 图中展示了子图采样的过程，即从更大的图中采样出小的、可处理的子图，以创建用于 GNN 训练的输入示例。 |

TF-GNN 1.0 首次推出灵活的 Python API，可在所有相关规模上配置动态或批量子图采样：在 Colab 笔记本中交互式进行（如此示例），用于高效采样存储在单个训练主机主内存中的小数据集，或通过 Apache Beam 分布式处理存储在网络文件系统上的海量数据集（多达数亿个节点和数十亿条边）。有关详细信息，请分别参阅我们的内存采样和基于 Beam 采样的用户指南。

在这些相同的采样子图上，GNN 的任务是计算根节点处的隐藏（或潜在）状态；该隐藏状态聚合并编码根节点邻域的相关信息。一种经典方法是消息传递神经网络。在每一轮消息传递中，节点沿入边从邻居接收消息，并据此更新自己的隐藏状态。经过 n 轮后，根节点的隐藏状态反映了 n 条边范围内所有节点的聚合信息（下图展示了 n = 2 的情况）。消息和新的隐藏状态由神经网络的隐藏层计算。在异构图（heterogeneous graph）中，通常有必要为不同类型的节点和边使用单独训练的隐藏层。

|  |
| --- |
| 图中展示了一个简单的消息传递神经网络，在每一步中，节点状态从外层节点传播到内层节点，在内层节点处进行池化以计算新的节点状态。一旦到达根节点，就可以做出最终预测。 |

训练设置通过在 GNN 的隐藏状态之上为有标签节点放置一个输出层、计算损失（以衡量预测误差）以及像任何神经网络训练中通常那样通过反向传播更新模型权重来完成。

除了监督训练（即最小化由标签定义的损失）之外，GNN 还可以以无监督方式（即不使用标签）进行训练。这使我们能够计算节点及其特征的离散图结构的连续表示（或嵌入）。这些表示随后通常用于其他机器学习系统。通过这种方式，由图编码的离散关系信息可以被纳入更典型的神经网络用例中。TF-GNN 支持为异构图精细指定无监督目标。

构建 GNN 架构

TF-GNN 库支持在不同抽象层次上构建和训练 GNN。

在最高层次上，用户可以使用该库捆绑提供的、以 Keras 层表达的任何预定义模型。除了一小部分来自研究文献的模型外，TF-GNN 还附带一个高度可配置的模型模板，该模板提供了一组精选的建模选择，我们发现这些选择在许多内部问题上提供了强大的基线。这些模板实现了 GNN 层；用户只需初始化 Keras 层。

import tensorflow_gnn as tfgnn
from tensorflow_gnn
.
models import mt_albis

def model_fn(graph_tensor_spec
:
tfgnn
.
GraphTensorSpec)
:
"""Builds a GNN as a Keras model
."""
graph
=
inputs
=
tf
.
keras
.
Input(type_spec
=
graph_tensor_spec)

  # Encode input features (callback omitted
for
brevity)
.
graph
=
tfgnn
.
keras
.
layers
.
MapFeatures(
      node_sets_fn
=
set_initial_node_states)(graph)

  #
For
each
round
of message passing
...
for
_ in range(
2
)
:
#
...
create
and
apply a Keras layer
.
graph
=
mt_albis
.
MtAlbisGraphUpdate(
        units
=128
, message_dim
=64
,
        attention_type
="
none", simple_conv_reduce_type
=
"mean"
,
        normalization_type
=
"layer"
, next_state_type
=
"residual"
,
        state_dropout_rate
=0.2
, l2_regularization
=1e-5
,
    )(graph)
return
tf
.
keras
.
Model(inputs, graph)

在最低层次上，用户可以使用在图周围传递数据的原语从头编写 GNN 模型，例如将数据从节点广播到其所有出边，或从所有入边将数据池化到节点中（例如，计算传入消息的总和）。TF-GNN 的图数据模型在特征或隐藏状态方面同等对待节点、边和整个输入图，使得不仅能够直接表达上述以节点为中心的模型（如 MPNN），还能表达更一般的 GraphNets 形式。这可以但不一定需要在核心 TensorFlow 之上使用 Keras 作为建模框架来完成。有关更多细节以及中间建模层次，请参阅 TF-GNN 用户指南和模型集合。

训练编排

虽然高级用户可以自由进行自定义模型训练，但 TF-GNN Runner 也提供了一种简洁的方式来在常见情况下编排 Keras 模型的训练。一个简单的调用可能如下所示：

from tensorflow_gnn import runner

runner
.
run
(
   task
=
runner
.
RootNodeBinaryClassification("papers",
...
),
   model_fn
=
model_fn,
   trainer
=
runner
.
KerasTrainer(tf
.
distribute
.
MirroredStrategy(), model_dir
=
"
/
tmp
/
model"),
   optimizer_fn
=
tf
.
keras
.
optimizers
.
Adam,
   epochs
=10
,
   global_batch_size
=128
,
   train_ds_provider
=
runner
.
TFRecordDatasetProvider("
/
tmp
/
train
*"
),
   valid_ds_provider
=
runner
.
TFRecordDatasetProvider("
/
tmp
/
validation
*"
),
   gtspec
=...
,
)

Runner 为诸如分布式训练和在 Cloud TPU 上为固定形状进行 tfgnn.GraphTensor padding 等机器学习痛点提供了即用型解决方案。除了在单个任务上训练（如上所示）之外，它还支持在多个（两个或更多）任务上联合训练。例如，无监督任务可以与监督任务混合，以应用特定的归纳偏置来形成最终的连续表示（或嵌入）。调用者只需将 task 参数替换为任务映射：

from tensorflow_gnn import runner
from tensorflow_gnn
.
models import contrastive_losses

runner
.
run
(
     task
=
{
"classification"
:
runner
.
RootNodeBinaryClassification("papers",
...
),
"dgi"
:
contrastive_losses
.
DeepGraphInfomaxTask(
"papers"
),
      },
...
)

此外，TF-GNN Runner 还包含用于模型归因的积分梯度实现。积分梯度的输出是一个 GraphTensor，其连接性与观测到的 GraphTensor 相同，但其特征被替换为梯度值，其中较大的值在 GNN 预测中的贡献大于较小的值。用户可以检查梯度值，以查看其 GNN 最常使用哪些特征。

结论

简而言之，我们希望 TF-GNN 能够有助于推动 GNN 在 TensorFlow 中的大规模应用，并促进该领域的进一步创新。如果您想了解更多信息，请尝试我们使用流行的 OGBN-MAG 基准测试的 Colab 演示（在浏览器中即可运行，无需安装），浏览我们的其他用户指南和 Colab，或查看我们的论文。

转载声明：本文为转载整理，原文版权归原作者及原站点所有。

原文来源：TensorFlow Blog

原创地址：https://blog.tensorflow.org/2024/02/graph-neural-networks-in-tensorflow.html