TensorFlow官方出了个GNN框架，YYDS！

2021 年 11 月 20 日 图与推荐

今天，我们很高兴发布TensorFlow 图神经网络(GNN)，这是一个旨在使使用 TensorFlow 轻松处理图结构化数据的库。我们已经在 Google 的各种环境中（例如，垃圾邮件和异常检测、流量估计、YouTube 内容标记）在 Google 的生产中使用了该库的早期版本，并作为我们可扩展图挖掘管道中的一个组件。特别是，考虑到 Google 的数据种类繁多，我们的库在设计时就考虑到了异构图。我们发布这个库的目的是鼓励与行业研究人员的合作。

为什么要使用 GNN？

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

图可以为许多不同类型的数据之间的关系建模，包括网页（左）、社交关系（中）或分子（右）。

图表示实体集合（节点或顶点）之间的关系（边）。我们可以表征每个节点、边或整个图，从而将信息存储在图的每个部分中。此外，例如，我们可以将方向性归因于边缘来描述信息或交通流。

GNN 可用于回答有关这些图的多个特征的问题。通过在图级别工作，我们尝试预测整个图的特征。我们可以识别某些“形状”的存在，例如图中可能代表子分子或可能代表密切社会关系的圆圈。GNN 可用于节点级任务，对图的节点进行分类，并预测图中的分区和亲和性，类似于图像分类或分割。最后，我们可以在边缘级别使用 GNN 来发现实体之间的连接，也许使用 GNN 来“修剪”边缘以识别场景中对象的状态。

结构

TF-GNN 提供了在 TensorFlow 中实现 GNN 模型的构建块。除了建模 API 之外，我们的库还围绕处理图形数据的艰巨任务提供了广泛的工具：基于张量的图形数据结构、数据处理管道和一些供用户快速入门的示例模型。

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

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

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

一个定义良好的模式来声明一个图的拓扑结构，以及验证它的工具。该模式描述了其训练数据的形状，并用于指导其他工具。
一种 GraphTensor复合张量类型，它保存图形数据，可以批处理，并具有可用的图形操作例程。
GraphTensor结构上的操作库：

节点和边上的各种高效广播和池化操作，以及相关工具。
标准烘焙卷积库，机器学习工程师/研究人员可以轻松扩展。
产品工程师可以快速构建 GNN 模型而不必担心其细节的高级 API。

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

示例用法

在下面的示例中，我们使用 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 仍可确保其适用并无缝处理异构图（具有各种类型的节点和边）。

下一步

您可以查看 TF-GNN GitHub 存储库以获取更多信息。要保持最新状态，您可以阅读 TensorFlow博客、加入 TensorFlow 论坛（位于Discussion.tensorflow.org ）、关注twitter.com/tensorflow或订阅youtube.com/tensorflow。如果您已经构建了想要分享的内容，请通过goo.gle/TFCS将其提交给我们的社区聚焦。如需反馈，请在GitHub 上提交问题。谢谢！

致谢

此处描述的工作是 Google 的 Oleksandr Ferludin、Martin Blais、Jan Pfeifer、Arno Eigenwillig、Dustin Zelle、Bryan Perozzi 和 Da-Cheng Juan 与 Sibon Li、Alvaro Sanchez-Gonzalez、Peter Battaglia、Kevin Villela 之间的研究合作DeepMind 的 Jennifer She 和 David Wong。