作者：Sibon Li、Jan Pfeifer 和 Bryan Perozzi 以及 Douglas Yarrington

今天，我们很高兴发布 TensorFlow 图神经网络 (GNN)，这是一个旨在简化使用 TensorFlow 处理图结构化数据的库。我们在 Google 的生产环境中使用过该库的早期版本，用于各种场景（例如，垃圾邮件和异常检测、流量估计、YouTube 内容标记）以及我们可扩展图挖掘管道的组件。特别地，鉴于 Google 数据的多样性，我们的库的设计考虑了异构图。我们发布这个库是为了鼓励与行业研究人员合作。

为什么要使用 GNN？

图无处不在，无论是在现实世界中还是在我们的工程系统中。一组对象、地点或人以及它们之间的连接通常可以描述为一个图。机器学习问题中我们所看到的数据往往是有结构或关系的，因此也可以用图来描述。虽然 GNN 的基础研究可能已经开展了几十年，但现代 GNN 能力的最新进展已导致各个领域取得进步，例如 交通预测、谣言和假新闻检测、疾病传播建模、物理模拟 以及 了解分子为什么有气味。

图可以模拟不同类型数据之间的关系，包括网页（左）、社交连接（中）或分子（右）。

图表示实体（节点或顶点）集合之间的关系（边）。我们可以描述每个节点、边或整个图，并由此在图的这些部分存储信息。此外，我们可以为边分配方向性，以描述信息或流量流，例如。

GNN 可用于回答有关这些图的多个特征的问题。通过在图级别工作，我们尝试预测整个图的特征。我们可以识别某些“形状”的存在，例如图中的圆圈可能代表子分子或密切的社会关系。GNN 可用于节点级任务，对图的节点进行分类，并预测图中的分区和亲和力，类似于图像分类或分割。最后，我们可以使用 GNN 在边级别发现实体之间的连接，也许使用 GNN 来“修剪”边以识别场景中对象的狀態。

结构

TF-GNN 提供构建块，用于在 TensorFlow 中实现 GNN 模型。除了建模 API 外，我们的库还提供了广泛的工具，用于解决与图数据相关的复杂任务：基于张量的图数据结构、数据处理管道以及一些示例模型，供用户快速上手。

构成工作流的 TF-GNN 的各个组件。

TF-GNN 库 的初始版本包含许多实用程序和功能，供初学者和经验丰富的用户使用，包括

• 一个高级 Keras 风格的 API，用于创建 GNN 模型，这些模型可以轻松地与其他类型的模型组合在一起。GNN 通常与排名、深度检索（双编码器）一起使用，或与其他类型的模型（图像、文本等）混合使用。
  • 用于异构图的 GNN API。我们在 Google 和现实世界中处理的许多图问题都包含不同类型的节点和边。因此，我们选择提供一种简便的方法来对此进行建模。

• 一个明确定义的模式，用于声明图的拓扑结构，以及用于验证它的工具。此模式描述其训练数据的形状，并用于指导其他工具。
  
• 一种名为 GraphTensor 的复合张量类型，它保存图数据，可以进行批处理，并且具有可用的图操作例程。
  
• 一个针对 GraphTensor 结构的操作库
  • 各种对节点和边进行的有效广播和池化操作，以及相关工具。
  • 一个标准烘焙卷积库，ML 工程师/研究人员可以轻松地扩展它。
  • 一个高级 API，供产品工程师快速构建 GNN 模型，而无需担心其细节。
    
• 磁盘上图形状的训练数据的编码，以及用于将此数据解析为数据结构的库，您的模型可以从中提取各种特征。

示例用法

在下面的示例中，我们使用 TF-GNN Keras API 构建了一个模型，根据用户观看的电影和他们喜欢的类型向用户推荐电影。

我们使用 ConvGNNBuilder 方法指定边和节点配置的类型，即对边使用 WeightedSumConvolution（在下面定义）。并且在每次通过 GNN 时，我们都会通过一个密集的互连层来更新节点值。

    import tensorflow as tf
    import tensorflow_gnn as tfgnn

    # Model hyper-parameters:
    h_dims = {'user': 256, 'movie': 64, 'genre': 128}
    
    # Model builder initialization:
    gnn = tfgnn.keras.ConvGNNBuilder(
      lambda edge_set_name: WeightedSumConvolution(),
      lambda node_set_name: tfgnn.keras.layers.NextStateFromConcat(
         tf.keras.layers.Dense(h_dims[node_set_name]))
    )
    
    # Two rounds of message passing to target node sets:
    model = tf.keras.models.Sequential([
        gnn.Convolve({'genre'}),  # sends messages from movie to genre
        gnn.Convolve({'user'}),  # sends messages from movie and genre to users
        tfgnn.keras.layers.Readout(node_set_name="user"),
        tf.keras.layers.Dense(1)
    ])

上面的代码效果很好，但有时我们可能希望为我们的 GNN 使用更强大的自定义模型架构。例如，在我们之前的用例中，我们可能希望指定某些电影或类型在提供推荐时具有更大的权重。在下面的代码段中，我们定义了一个更高级的 GNN，它具有自定义图卷积，在本例中为加权边。我们定义了 WeightedSumConvolution 类，将边值作为所有边权重的总和进行池化

class WeightedSumConvolution(tf.keras.layers.Layer):
  """Weighted sum of source nodes states."""

  def call(self, graph: tfgnn.GraphTensor,
           edge_set_name: tfgnn.EdgeSetName) -> tfgnn.Field:
    messages = tfgnn.broadcast_node_to_edges(
        graph,
        edge_set_name,
        tfgnn.SOURCE,
        feature_name=tfgnn.DEFAULT_STATE_NAME)
    weights = graph.edge_sets[edge_set_name]['weight']
    weighted_messages = tf.expand_dims(weights, -1) * messages
    pooled_messages = tfgnn.pool_edges_to_node(
        graph,
        edge_set_name,
        tfgnn.TARGET,
        reduce_type='sum',
        feature_value=weighted_messages)
    return pooled_messages

请注意，即使卷积的编写仅考虑了源节点和目标节点，TF-GNN 也会确保它适用于异构图（具有各种类型的节点和边），并且能够无缝地工作。

后续步骤

您可以在 TF-GNN 的 GitHub 存储库 中找到更多信息。要保持最新状态，您可以阅读 TensorFlow 的 博客，加入 TensorFlow 论坛，网址为 discuss.tensorflow.org，关注 twitter.com/tensorflow，或订阅 youtube.com/tensorflow。如果您构建了一些想分享的东西，请将其提交给我们的社区亮点，网址为 goo.gle/TFCS。对于反馈，请在 GitHub 上提交问题。谢谢！

致谢

这里描述的工作是 Google 的 Oleksandr Ferludin‎、Martin Blais、Jan Pfeifer‎、Arno Eigenwillig、Dustin Zelle、Bryan Perozzi 和 Da-Cheng Juan 以及 DeepMind 的 Sibon Li、Alvaro Sanchez-Gonzalez、Peter Battaglia、Kevin Villela、Jennifer She 和 David Wong 之间的研究合作。