本文也分享在 Google Research 博客

对象及其关系在周围的世界中无处不在，关系在理解对象时可能与孤立地查看其属性一样重要 - 例如：交通网络、生产网络、知识图谱或社交网络。离散数学和计算机科学在将此类网络形式化为 图 方面有着悠久的历史，图由节点组成，这些节点通过边以各种不规则的方式任意连接。然而，大多数机器学习 (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 原语。

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）。消息和新的隐藏状态由神经网络的隐藏层计算。在异构图中，使用单独训练的隐藏层来处理不同类型的节点和边通常是有意义的。

图示，一个简单消息传递神经网络，其中，在每一步，节点状态从外节点传播到内节点，在那里进行池化以计算新的节点状态。到达根节点后，可以进行最终预测。

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

除了监督训练（即，最小化由标签定义的损失）之外，GNN 还可以以无监督的方式进行训练（即，没有标签）。这使我们能够计算离散图结构的节点及其特征的连续表示（或嵌入）。然后，这些表示通常用于其他 ML 系统。通过这种方式，图编码的离散关系信息可以包含在更典型的网络网络用例中。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 填充。除了对单个任务进行训练（如上所示），它还支持对多个（两个或更多）任务进行联合训练。例如，无监督任务可以与监督任务混合，以告知具有特定于应用的归纳偏差的最终连续表示（或嵌入）。调用者只需要用任务映射替换任务参数

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 中的应用，并推动该领域的进一步创新。如果您想了解更多信息，请尝试我们的 Colab 演示，其中包含流行的 OGBN-MAG 基准测试（在您的浏览器中，无需安装），浏览我们的其他 用户指南和 Colabs，或者查看我们的 论文。

致谢

TF-GNN 版本 1.0 是由Google Research（Sami Abu-El-Haija、Neslihan Bulut、Bahar Fatemi、Johannes Gasteiger、Pedro Gonnet、Jonathan Halcrow、Liangze Jiang、Silvio Lattanzi、Brandon Mayer、Vahab Mirrokni、Bryan Perozzi、Anton Tsitsulin、Dustin Zelle）、Google Core ML（Arno Eigenwillig、Oleksandr Ferludin、Parth Kothari、Mihir Paradkar、Jan Pfeifer、Rachael Tamakloe）和Google DeepMind（Alvaro Sanchez-Gonzalez 和 Lisa Wang）合作开发的。