谷歌正式发布TensorFlow 图神经网络

描述

发布人: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 图神经网络 (GNNs)

文章出处:【微信公众号:谷歌开发者】欢迎添加关注!文章转载请注明出处。

  审核编辑:彭菁
 

 

打开APP阅读更多精彩内容
声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉

全部0条评论

快来发表一下你的评论吧 !

×
20
完善资料,
赚取积分