专业的编程技术博客社区

网站首页 > 博客文章 正文

DeepMind加持的GNN框架正式开源!TensorFlow进入图神经网络时代

baijin 2024-09-02 11:10:55 博客文章 4 ℃ 0 评论

编辑:好困 小咸鱼

【新智元导读】谷歌在垃圾邮件检测、流量估计以及YouTube内容标签等环境中使用了一种强大的工具GNN(图神经网络)。11月18日,谷歌联合DeepMind对外开源TensorFlow GNN工具,助力流量预测、谣言和假新闻检测、疾病传播建模、物理模拟等领域的基础研究。


11月18日,谷歌联合DeepMind发布了TensorFlow GNN(图神经网络)。



目前,谷歌已经在诸如垃圾邮件检测、流量估计以及YouTube内容标签等环境中用上了这个库的早期版本。



为什么要用GNN?


图(Graph)是用于表示对象之间关联关系的一种抽象数据结构,使用节点/顶点(Node/Vertex)和边(Edge)进行描述:顶点表示对象,边表示对象之间的关系。


而在现实世界和工程系统中,图无处不在。


比如,一组物体、地点或人以及它们之间的联系通常可以用图形来描述。


通常,在机器学习问题中看到的数据是结构化的或关系型的,因此也可以用图来描述。


经过几十年的基础研究,GNN已经在很多的领域都取得了进展,如流量预测、谣言和假新闻检测、疾病传播建模、物理模拟以及理解分子为什么有气味。


图模拟不同类型数据之间的关系:网页(左)、社交联系(中)或分子(右)


通过GNN,就可以回答那些关于图的多种特征的问题。比如在图中观察到的各种不同的「形状」:图中的圆圈,可能代表子分子,也可能代表密切的社会关系。


在节点级的任务中,GNN可以对图的节点进行分类,并预测图中的分区和亲和力,类似于图像分类或分割。


在边级别的任务中,可以使用GNN来发现实体之间的连接,比如用GNN「修剪」图中的边,从而识别场景中对象的状态。


TF-GNN的结构


TF-GNN为在TensorFlow中实现GNN模型提供了构建模块。


除了建模API之外,TF-GNN还围绕着处理图数据的困难任务提供了大量的工具:基于Tensor的图数据结构,数据处理管道,以及一些供用户快速上手的示例模型。


组成工作流程的TF-GNN的各个部分


TF-GNN库的初始版本包含了许多实用程序和功能,包括:


  • 一个高水平的Keras风格的API,用于创建GNN模型,可以很容易地与其他类型的模型组合。GNN经常与排名、深度检索(双编码器)或与其他类型的模型(图像、文本等)结合使用。


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


  • 一个定义明确的模式来声明一个图的拓扑结构,以及验证它的工具。这个模式描述了其训练数据的形状,并用于指导其他工具。


  • 一个保存了图数据的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)

])

此外,还可以在某些场景下让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确保了它的适用性,并且可以无缝地在异构图(具有各种类型的节点和边)上工作。


安装说明

这是目前安装tensorflow_gnn预览版的唯一方法。强烈建议使用虚拟环境。


  1. 克隆tensorflow_gnn


gt; git clone https://github.com/tensorflow/gnn.git tensorflow_gnn


  1. 安装TensorFlow


TF-GNN需要用到TensorFlow 2.7中的一个功能:tf.ExtensionTypes。


gt; pip install tensorflow
  1. 安装Bazel


构建TF-GNN的源代码需要用到Bazel。


  1. 安装GraphViz


TF-GNN将使用GraphViz作为可视化工具。安装方法因操作系统而异,例如,在Ubuntu中:


gt; sudo apt-get install graphviz graphviz-dev


  1. 安装tensorflow_gnn

gt; cd tensorflow_gnn && python3 -m pip install


参考资料:

https://blog.tensorflow.org/2021/11/introducing-tensorflow-gnn.html?m=1

https://github.com/tensorflow/gnn

本文暂时没有评论,来添加一个吧(●'◡'●)

欢迎 发表评论:

最近发表
标签列表