TF2中的图神经网络

在TensorFlow 2.0中实现各种类型的图神经网络的代码和示例训练脚本。其中大部分基于tf-gnn-samples仓库中的代码。

代码由微软研究院的AI4Science团队维护。我们正在招聘。

目前这个包不在积极开发中，但它仍然作为我们分子生成模型的后端使用。

安装

您可以使用pip install tf2_gnn从Python包索引安装tf2_gnn模块。

或者（例如，用于开发），您可以检出此存储库，导航到它并运行pip install -e ./来将其安装为本地可编辑包。

您将能够使用 tf2_gnn.layers.GNN 类和相关实用工具。

此代码已在 Python 3.7、3.8 和 3.9 上进行了测试。

测试安装

为了测试所有组件是否正确设置，您可以首先在蛋白质-蛋白质相互作用（PPI）任务上进行简单的实验，该任务最早由 Zitnik & Leskovec, 2017 描述。您可以从 https://data.dgl.ai/dataset/ppi.zip 下载此任务的数据，并将其解压到本地目录中（例如，data/ppi）。然后，您可以使用方便的实用工具 tf2_gnn_train（有关选项的描述，请参阅 --help）来训练一个关系图卷积网络（RGCN）模型，如下所示

$ tf2_gnn_train RGCN PPI --max-epochs 10 data/ppi/
Setting random seed 0.
Trying to load task/model-specific default parameters from /dpuhome/files/users/mabrocks/Projects/TF2-GNN/tf2_gnn/cli_utils/default_hypers/PPI_RGCN.json ... File found.
 Dataset default parameters: {'max_nodes_per_batch': 10000, 'add_self_loop_edges': True, 'tie_fwd_bkwd_edges': False}
Loading data from data/ppi/.
 Loading PPI train data from data/ppi/.
 Loading PPI valid data from data/ppi/.
[...]
Dataset parameters: {"max_nodes_per_batch": 8000, "add_self_loop_edges": true, "tie_fwd_bkwd_edges": false}
Model parameters: {"gnn_aggregation_function": "sum", "gnn_message_activation_function": "ReLU", "gnn_hidden_dim": 320, "gnn_use_target_state_as_input": false, "gnn_normalize_by_num_incoming": true, "gnn_num_edge_MLP_hidden_layers": 0, "gnn_message_calculation_class": "RGCN", "gnn_initial_node_representation_activation": "tanh", "gnn_dense_intermediate_layer_activation": "tanh", "gnn_num_layers": 4, "gnn_dense_every_num_layers": 10000, "gnn_residual_every_num_layers": 10000, "gnn_use_inter_layer_layernorm": false, "gnn_layer_input_dropout_rate": 0.1, "gnn_global_exchange_mode": "gru", "gnn_global_exchange_every_num_layers": 10000, "gnn_global_exchange_weighting_fun": "softmax", "gnn_global_exchange_num_heads": 4, "gnn_global_exchange_dropout_rate": 0.2, "optimizer": "Adam", "learning_rate": 0.001, "learning_rate_decay": 0.98, "momentum": 0.85, "gradient_clip_value": 1.0}
Initial valid metric: Avg MicroF1: 0.368.
   (Stored model metadata to trained_model/RGCN_PPI__2020-02-25_11-10-38_best.pkl and weights to trained_model/RGCN_PPI__2020-02-25_11-10-38_best.hdf5)
== Epoch 1
 Train:  25.6870 loss | Avg MicroF1: 0.401 | 2.63 graphs/s
 Valid:  33.1668 loss | Avg MicroF1: 0.419 | 4.01 graphs/s
  (Best epoch so far, target metric decreased to -0.41886 from -0.36762.)
   (Stored model metadata to trained_model/RGCN_PPI__2020-02-25_11-10-38_best.pkl and weights to trained_model/RGCN_PPI__2020-02-25_11-10-38_best.hdf5)
[...]

训练完成后，可以使用 tf2_gnn_test trained_model/RGCN_PPI__2020-02-25_11-10-38_best.pkl data/ppi 测试训练好的模型。

代码结构

层

库的核心功能实现为 TensorFlow 2（Keras）层，使其易于集成到其他代码中。

tf2_gnn.layers.GNN

这实现了一个深度图神经网络，堆叠了多层消息传递。在构建时，需要提供一组超参数字典（默认值可以从 GNN.get_default_hyperparameters() 获取）。这些超参数配置了 GNN 层的确切堆叠

• "num_layers" 设置 GNN 消息传递层的数量（通常，一个介于 2 和 16 之间的数字）

• "message_calculation_class" 配置消息传递样式。这选择在每个步骤中使用的 tf2_gnn.layers.message_passing.* 层。

  我们目前支持以下内容

  • GGNN：门控图神经网络（《Li et al., 2015》）。
  • RGCN：关系图卷积网络（《Schlichtkrull et al., 2017》）。
  • RGAT：关系图注意力网络（《Veličković et al., 2018》）。
  • RGIN：关系图同构网络（《Xu et al., 2019》）。
  • GNN-Edge-MLP：带有边缘 MLP 的图神经网络 - RGCN 的一个变体，其中在边缘上的消息使用完整的 MLP 计算，而不仅仅是应用在源状态上的单层。
  • GNN-FiLM：具有特征线性调制（Feature-wise Linear Modulation）的图神经网络（《Brockschmidt, 2019》） - 带有 FiLM 层的 RGCN 的新扩展。

  其中一些公开了额外的超参数；请参阅它们的实现以获取详细信息。

• "hidden_dim" 设置所有消息传递层的输出大小。

• "layer_input_dropout_rate" 设置每个消息传递层输入的 dropout 率（在训练期间）。

• "residual_every_num_layers" 设置在消息传递层之间插入残差连接的频率。具体来说，k 的值表示每个层 l（仅这些！）不会接收层 l-1 的输出作为输入，而是接收层 l-1 和 l-k 的输出均值。

• "use_inter_layer_layernorm" 是一个布尔标志，表示是否在不同的消息传递层之间使用 LayerNorm。

• "dense_every_num_layers" 配置在消息传递层之间插入每个节点表示密集层的频率。将此值设置得很大（大于 "num_layers"）表示不会插入任何密集层。

  "dense_intermediate_layer_activation" 配置密集层之后的激活函数；默认的 "tanh" 可以帮助稳定深度 GNN 的训练。

• "global_exchange_every_num_layers" 配置进行图级别信息交换的频率。为此，计算图级表示（见下面的 tf2_gnn.layers.NodesToGraphRepresentation），然后使用它来更新每个节点的表示。此更新的样式由 "global_exchange_mode" 配置，提供三种模式

  • "mean"，它只是计算节点和图级表示的算术平均值。
  • "mlp"，使用MLP计算新的表示，以节点和图级别表示的连接作为输入。
  • "gru"，使用GRU单元，该单元以旧节点表示作为状态，以图表示作为输入。

GNN层以名为GNNInput的元组作为输入，该元组封装了初始节点特征、邻接列表和辅助信息。构建此类元组的最简单方法是将提供的数据集类与提供的模型一起使用。

tf2_gnn.layers.NodesToGraphRepresentation

这实现了计算图级别表示的任务，给定节点级别表示（例如，由GNN层获得的表示）。

目前，这仅由WeightedSumGraphRepresentation层实现，该层通过多个头部的加权求和（转换后的）节点表示来生成图表示，配置如下层的构造函数中的超参数：

• graph_representation_size设置计算表示的大小。将此设置为1，则该层可以直接用于实现图级别回归任务。
• num_heads配置要计算多少个并行（独立）的加权求和，其结果连接以获得最终结果。请注意，这意味着graph_representation_size需要是num_heads值的倍数。
• weighting_fun可以取两个值：
  • "sigmoid"独立地为每个节点计算权重，首先计算每个节点的得分，然后通过sigmoid进行压缩。这对于与在图中出现特征计数相关的任务来说是合适的，其中节点权重用于忽略某些节点。
  • "softmax"通过首先计算每个节点的得分，然后对所有得分执行softmax来计算所有图节点的权重。这对于需要识别图中重要部分的任务来说是合适的。
• scoring_mlp_layers、scoring_mlp_activation_fun和scoring_mlp_dropout_rate配置计算每个节点得分的MLP。
• transformation_mlp_layers、transformation_mlp_activation_fun和transformation_mlp_dropout_rate配置计算要加和的转换节点表示的MLP。

数据集

我们使用图的稀疏表示，这需要一个复杂的批处理策略，其中组成一个minibatch的图被连接成一个由许多不连通组件组成的单个图。可扩展的tf2_gnn.data.GraphDataset类实现了此过程，并可被继承以处理特定任务的数据集和附加属性。它公开了一个get_tensorflow_dataset方法，可以用来获取一个tf.data.Dataset，该数据集可用于训练/评估循环。

我们目前提供了三种实现：

• tf2_gnn.data.PPIDataset实现了读取Zitnik & Leskovec，2017首次使用的蛋白质-蛋白质相互作用（PPI）数据。
• tf2_gnn.data.QM9Dataset实现了读取Ramakrishnan等人，2014首次使用的量子化学数据。
• tf2_gnn.data.JsonLGraphPropertyDataset实现了读取由单个属性组成的通用数据集，该属性以JSONLines格式存储。
  • 期望"train.jsonl.gz"、"valid.jsonl.gz"和"test.jsonl.gz"文件存储训练/验证/测试数据集。
  • 每个文件都是一个压缩文本文件，其中每行都是一个有效的JSON字典，具有以下键：
    • 一个"graph"键，该键又指向一个字典，该字典有键：
      • "node_features"（数值初始节点标签的列表），
      • "adjacency_lists"（有向边对的列表），
    • 一个具有一个单精度浮点值的"Property"键。

模型

我们在tf2_gnn.models中提供了一些内置模型，这些模型可以直接重用，也可以作为其他模型的灵感来源。

• tf2_gnn.models.GraphRegressionTask实现了一个图级别的回归模型，例如在QM9任务中进行分子级别的预测。
• tf2_gnn.models.GraphBinaryClassificationTask实现了一个二分类模型。
• tf2_gnn.models.NodeMulticlassTask实现了一个节点级别的多分类模型，适用于实现PPI任务。

任务

任务是由数据集、模型和特定的超参数设置组合而成的。这些任务可以使用tf2_gnn.utils.task_utils中的工具注册（然后通过名称使用），并在tf2_gnn_train等工具中使用（其中还定义了一些默认任务）。

作者

• Henry Jackson-Flux
• Marc Brockschmidt
• Megan Stanley
• Pashmina Cameron
• Krzysztof Maziarz

参考文献

Brockschmidt, 2019

Marc Brockschmidt. GNN-FiLM: 图神经网络中的特征线性调制。(https://arxiv.org/abs/1906.12192)

Li et al., 2015

Yujia Li, Daniel Tarlow, Marc Brockschmidt, and Richard Zemel. 门控图序列神经网络。在2016年国际学习表示会议（ICLR）。(https://arxiv.org/pdf/1511.05493.pdf)

Ramakrishnan et al., 2014

Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole Von Lilienfeld. 134千分子的量子化学结构和性质。科学数据，1，2014。(https://www.nature.com/articles/sdata201422/)

Schlichtkrull et al., 2017

Michael Schlichtkrull, Thomas N. Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. 使用图卷积网络对关系数据进行建模。在2018年扩展语义网会议（ESWC）。(https://arxiv.org/pdf/1703.06103.pdf)

Veličković et al. 2018

Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. 图注意力网络。在2018年国际学习表示会议（ICLR）。(https://arxiv.org/pdf/1710.10903.pdf)

Xu et al. 2019

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. 图神经网络有多强大？在2019年国际学习表示会议（ICLR）。(https://arxiv.org/pdf/1810.00826.pdf)

Zitnik & Leskovec, 2017

Marinka Zitnik and Jure Leskovec. 通过多层组织网络预测多细胞功能。生物信息学，33，2017。(https://arxiv.org/abs/1707.04638)

贡献

本项目欢迎贡献和建议。大多数贡献都需要您同意贡献者许可协议（CLA），声明您有权，并且实际上确实授予我们使用您贡献的权利。有关详细信息，请访问https://cla.opensource.microsoft.com。

当您提交拉取请求时，CLA机器人将自动确定您是否需要提供CLA，并相应地装饰PR（例如，状态检查，注释）。只需遵循机器人提供的说明即可。您只需要在所有使用我们的CLA的仓库中这样做一次。

本项目采用了微软开源行为准则。有关更多信息，请参阅行为准则FAQ或联系opencode@microsoft.com，以提出任何其他问题或评论。