推出 TensorFlow 图神经网络 (GNNs)

发布人：Sibon Li、Jan Pfeifer、Bryan Perozzi 和 Douglas Yarrington

日前，我们很高兴发布了 TensorFlow 图神经网络 (Graph Neural Networks, GNNs)，此库可以帮助开发者利用 TensorFlow 轻松处理图结构化数据。我们已在 Google 的多个生产环境中使用了该库的早期版本（例如，垃圾邮件和异常检测、数据流量估计、YouTube 内容标签），并作为可扩容的图挖掘管道的一个组成部分。特别是，鉴于 Google 的数据类型繁多，我们的库在设计时就考虑到了异构图。发布此库的初衷是为了鼓励与业界研究人员的合作。

TensorFlow 图神经网络

https://github.com/tensorflow/gnn

为何使用 GNNs？

在现实世界和我们的工程系统中，“图”无处不在。一组物体、地点或人以及它们之间的联系通常都可以用图来表述。通常情况下，我们在机器学习问题中看到的数据是结构化或关系化的，因此也可以用图来表述。虽然关于 GNNs 的基础研究可能只有几十年的历史，但当代 GNNs 最近取得的功能进展已经帮助推动了多个领域的进步，其中包括数据流量预测、谣言和假新闻检测、疾病传播建模、物理学模拟和理解分子有气味的原因等。

数据流量预测

https://deepmind.com/blog/article/traffic-prediction-with-advanced-graph-neural-networks

谣言和假新闻检测

https://arxiv.org/abs/2108.03548

疾病传播建模

https://arxiv.org/abs/2007.03113

物理学模拟

http://proceedings.mlr.press/v80/sanchez-gonzalez18a/sanchez-gonzalez18a.pdf

理解分子有气味的原因

https://arxiv.org/abs/1910.10685

图可以对多种不同类型数据之间的关系进行建模，

包括网页（左）、社交关系（中）或分子（右）等

图代表了一组实体（节点或顶点）之间的关系（边）。我们可以通过描述每个节点、边或整个图的特征，从而将信息存储在图的每一块中。此外，我们还可以赋予边方向性，来描述信息或数据流，等等。

若这些图有多个特征，则可以用 GNNs 来解决这类问题。通过图层级的研究，我们尝试预测整个图的特征。我们可以识别某些“形状”的存在，如图中的圆圈，可能代表亚分子，也可能代表密切的社会关系。GNNs 可以用于节点级的任务，对图的节点进行分类，并预测图中的分区和相似性，类似于图像分类或分割。最后，我们可以在边层级上使用 GNNs 来发现实体之间的联系，或许可以使用 GNNs 来“修剪”边，以确定场景中对象的状态。

结构

TF-GNN 提供了基本模块，以便在 TensorFlow 中实现 GNN 模型。除了建模 API，我们的库还针对处理图数据的困难任务提供了大量的工具：基于 Tensor 的图数据结构、数据处理管道，以及一些供用户快速上手的示例模型。

组成工作流的各种 TF-GNN 组件

TF-GNN库的初始版本包含一些实用程序和功能，初学者和有经验的用户都能使用，其中包括：

TF-GNN 库

https://github.com/tensorflow/gnn

高阶 Keras 式 API，用于创建 GNN 模型，可以轻松地与其他类型的模型组合。GNNs 经常与排名、深度检索（双编码器）结合使用或与其他类型的模型（图像、文本等）混合使用。

用于异构图的 GNN API。我们在 Google 和现实世界中处理的许多图问题都包含不同类型的节点和边。因此，我们选择提供一种简单的方法来对此建模。

定义明确的架构，用于声明图的拓扑，以及验证架构的工具。此架构描述了其训练数据的形状，并用于指导其他工具。

GraphTensor 复合张量类型，它持有图数据，可以分批处理，还有图操作例程可用。

关于 GraphTensor 结构的操作库：

对节点和边的各种有效的广播和池化运算，以及相关工具。

标准并入的卷积库，ML 工程师/研究人员可以轻松地对其进行扩展。

高阶 API，帮助产品工程师快速构建 GNN 模型，而不必担心其细节问题。

磁盘上的图形训练数据的编码，以及用来将该数据解析为数据结构的库，您的模型可以从该库中提取各种特征。

示例用法

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

我们使用 ConvGNNBuilder 方法来指定边类型和节点配置，即对边使用 WeightedSumConvolution（定义如下）。每次通过 GNN 时，我们将通过 Dense 互连层来更新节点值：

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)
    ])

以上代码很好用，但有时我们可能想要为 GNNs 使用更强大的自定义模型架构。例如，在之前的用例中，我们可能想指定某些电影或类型，让它们在我们进行推荐时拥有更多权重。在下列片段中，我们用自定义图卷积定义了一个更高级的 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 repo，以获得更多信息。

GitHub repo

https://github.com/tensorflow/gnn

欢迎加入 TensorFlow 论坛，随时关注 TensorFlow  博客、Twitter 或 YouTube，获悉最新消息。

TensorFlow 论坛

https://discuss.tensorflow.google.cn/

博客

https://blog.tensorflow.google.cn/

Twitter

http://twitter.com/tensorflow

YouTube

http://youtube.com/tensorflow

您还可以通过 Community Spotlight 计划向我们提交作品，分享您的构建成果。如需提供反馈，您可以在 GitHub 上提交问题。我们欢迎您的贡献和参与，谢谢！

Community Spotlight 计划

http://goo.gle/TFCS

GitHub

https://github.com/tensorflow/gnn/issues

致谢

本文所介绍的研究来自以下各位的合作成果：来自 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。

点击屏末 | 阅读原文 | 访问 TensorFlow 官网