推出 TensorFlow 图神经网络(GNNs)

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

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

为何使用 GNNs?

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

图可以对多种不同类型数据之间的关系进行建模,包括网页(左)、社交关系(中)或分子(右)等

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

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

结构

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

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

TF-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,以获得更多信息。

欢迎加入 TensorFlow 论坛,随时关注 TensorFlow 博客TwitterYouTube,获悉最新消息。

您还可以通过 Community Spotlight 计划 向我们提交作品,分享您的构建成果。如需提供反馈,您可以在 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。

原文:Introducing TensorFlow Graph Neural Networks
中文:TensorFlow 公众号