介绍一下TensorFlow图谱神经网络

语言: CN / TW / HK

今天,我们很高兴地发布了TensorFlow图形神经网络(GNNs),这是一个旨在使用TensorFlow轻松处理图形结构化数据的库。我们已经在谷歌的生产中使用了这个库的早期版本,在各种情况下(例如,垃圾邮件和异常检测,流量估计,YouTube内容标签),并作为我们可扩展的图形挖掘管道中的一个组成部分。特别是,鉴于谷歌的数据类型繁多,我们的库在设计时考虑到了异质图。我们发布这个库的目的是为了鼓励与业界的研究人员合作。

为什么使用GNNs?

图在我们身边,在现实世界和我们的工程系统中无处不在。一组物体、地点或人以及它们之间的联系通常可以被描述为一个图。更多时候,我们在机器学习问题中看到的数据是结构化或关系化的,因此也可以用图来描述。虽然关于GNN的基础研究可能已经有几十年的历史了,但现代GNN能力的最新进展已经在交通预测谣言和假新闻检测疾病传播建模、 物理模拟理解分子气味的原因等不同领域带来了进展。

Graphs can model the relationships between many different types of data, including web pages (left), social connections (center), or molecules (right).
图可以模拟许多不同类型的数据之间的关系,包括网页(左)、社会联系(中)或分子(右)。

图表示一组实体(节点或顶点)之间的关系(边)。我们可以对每个节点、边或整个图进行定性,从而将信息存储在图的每一块中。此外,我们可以赋予边以方向性,以描述信息或交通流,例如。

GNN可以用来回答关于这些图的多种特征的问题。通过在图层面的工作,我们试图预测整个图的特征。我们可以识别某些 "形状 "的存在,比如图中的圆圈可能代表子分子,也可能代表密切的社会关系。GNN可以用于节点级的任务,对图的节点进行分类,并预测图中的分区和亲和力,类似于图像分类或分割。最后,我们可以在边缘层面上使用GNN来发现实体之间的联系,也许使用GNN来 "修剪 "边缘,以确定场景中物体的状态。

结构

TF-GNN提供了在TensorFlow中实现GNN模型的构建块。除了建模API,我们的库还围绕着处理图数据的困难任务提供了广泛的工具:一个基于Tensor的图数据结构,一个数据处理管道,以及一些供用户快速上手的例子模型。

The various components of TF-GNN that make up the workflow.
组成工作流程的TF-GNN的各个组成部分。

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

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

    • 用于异质图的GNN API。我们在谷歌和现实世界中处理的许多图问题都包含不同类型的节点和边。因此,我们选择提供一种简单的方法来建模。
  • 一个定义明确的模式来声明图的拓扑结构,以及验证它的工具。这个模式描述了其训练数据的形状,并用于指导其他工具。

  • 一个GraphTensor 复合张量类型,它持有图数据,可以被分批处理,并有图操作例程可用。

  • 一个关于GraphTensor 结构的操作库:

    • 节点和边上的各种有效的广播和集合操作,以及相关工具。
    • 一个标准的烘焙卷积库,可以由ML工程师/研究人员轻松扩展。
    • 一个高层次的API,供产品工程师快速建立GNN模型,而不一定要担心其细节问题。
  • 一个在磁盘上的图形训练数据的编码,以及一个用于将这些数据解析为数据结构的库,你的模型可以从中提取各种特征。

使用实例

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

我们使用ConvGNNBuilder 方法来指定边缘和节点的配置类型,即对边缘使用WeightedSumConvolution (定义见下文)。而在每次通过GNN时,我们将通过Dense互联层更新节点值:

```python” 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 类来汇集边缘值,作为所有边缘的权重之和:

```python” 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也能确保它的适用性,并能在异质图(有各种类型的节点和边缘)上无缝工作。