系列教程GNN-algorithms之五：《注意力机制在图上的应用—GAT》

2020 年 8 月 7 日 专知

【导读】GAT是基于空域方法的图神经模型，它直接利用节点的空间邻域信息来学习节点的特征表示，灵活性强，这与基于谱域方法的GCN恰恰相反。在本教程中，我们将手把手教你如何构建基于Tensorflow的GAT模型在Cora数据集上实现节点分类任务。

系列教程《GNN-algorithms》

本文为系列教程《GNN-algorithms》中的内容，该系列教程不仅会深入介绍GNN的理论基础，还结合了TensorFlow GNN框架tf_geometric对各种GNN模型（GCN、GAT、GIN、SAGPool等）的实现进行了详细地介绍。本系列教程作者王有泽（https://github.com/wangyouze）也是tf_geometric框架的贡献者之一。

系列教程《GNN-algorithms》Github链接：

https://github.com/wangyouze/GNN-algorithms

TensorFlow GNN框架tf_geometric的Github链接：

https://github.com/CrawlScript/tf_geometric

前序讲解：

系列教程GNN-algorithms之一：《图卷积网络（GCN）的前世今生》

系列教程GNN-algorithms之二：《切比雪夫显神威—ChebyNet》

系列教程GNN-algorithms之三：《将图卷积简化进行到底—SGC》

系列教程GNN-algorithms之四：《Inductive Learning 大神—GraphSAGE》

前言

Graph Attention Networks（GAT）发表在ICLR2018上，利用masked self-attention 来学习中心节点与邻居节点之间的注意力权重，根据权重大小聚合邻居节点的空间信息来更新中心节点的特征表示，从而解决了基于卷积或者多项式近似卷积核等方法的固有缺陷。本教程将从具体实现的角度带你一起深入了解GAT。

GAT简介

GCN通过图的拉普拉斯矩阵来聚合邻居节点的特征信息，这种方式和图本身的结构紧密相关，这限制了GCN在训练时未见的图结构上的泛化能力。GAT利用注意力机制来对邻居节点特征加权求和，从而聚合邻域信息，完全摆脱了图结构的束缚，是一种归纳式学习方式。

GAT中的attention是self-attention，即Q(Query)，K(Key)，V(value)三个矩阵均来自统一输入。和所有的Attention机制一样，GAT的计算也分两步走：

计算注意力系数。

对于中心节点，我们需要逐个计算它与邻居节点之间的注意力系数：
首先对节点特征进行变换或者说是增加维度（常见的特征增强手段），然后将变换后的中心节点特征逐个与其邻居节点特征拼接后输入一个单层的神经网络a,所得结果就是中心节点与该邻居节点之间的注意力系数，当然，注意力系数还需要经过softmax归一化，将其转换为概率分布。
具体来说我们在实现过程中首先计算了中心节点Q向量与其邻居节点K向量之间的点乘，然后为了防止其结果过大，会除以一个尺度，其中为Query或者Key向量的维度。公式表示为：
通过加权求和的方式聚合节点信息。

利用上面计算所得的注意力系数对邻居节点的特征V进行线性组合作为中心节点的特征表示：
俗话说“一个篱笆三个桩，一个好汉三个帮”，GAT也秉承了这种思想，采用多头注意力机制（multi-head attention）来捕获邻居节点在不同的方面对中心节点影响力的强弱。我们将K 个head分别提取的节点特征表示进行拼接作为最终的节点表示：

我们可以看出GAT中的节点信息更新过程中是逐点运算的，每一个中心节点只与它的邻居节点有关，参数a和W也只与节点特征相关，与图结构无关。改变图的结构只需要改变节点的邻居关系重新计算，因此GAT与GraphSAGE一样都适用于inductive任务。而大名鼎鼎的GCN在节点特征的每一次更新上都需要全图参与，学习到的参数也很大程度上与图结构有关，因此GCN在inductive任务上不给力了。

GNN引入Attention机制有三大好处：

1. 参数少。Attention模型的复杂度与GCN等相比，复杂度更小，参数也更少，对算力的要求也就更小。

2. 速度快。Attention更加有利于并行计算。Attention机制中的每一步计算不依赖于前面一部的计算结果，并行性良好。

3. 效果好。Attention机制不会弱化对于长距离的信息记忆。

教程代码下载链接：

https://github.com/CrawlScript/tf_geometric/blob/master/demo/demo_gat.py

GAT论文地址：https://arxiv.org/pdf/1710.10903.pdf

文献参考：https://zhuanlan.zhihu.com/p/81350196

教程目录

开发环境
GAT的实现
模型构建
GAT训练
GAT评估

开发环境

操作系统: Windows / Linux / Mac OS
Python 版本: >= 3.5
依赖包:
tf_geometric（一个基于Tensorflow的GNN库）根据你的环境（是否已安装TensorFlow、是否需要GPU）从下面选择一条安装命令即可一键安装所有Python依赖:

        pip install -U tf_geometric # 这会使用你自带的TensorFlow，注意你需要tensorflow/tensorflow-gpu >= 1.14.0 or >= 2.0.0b1

	pip install -U tf_geometric[tf1-cpu] # 这会自动安装TensorFlow 1.x CPU版

	pip install -U tf_geometric[tf1-gpu] # 这会自动安装TensorFlow 1.x GPU版

	pip install -U tf_geometric[tf2-cpu] # 这会自动安装TensorFlow 2.x CPU版

	pip install -U tf_geometric[tf2-gpu] # 这会自动安装TensorFlow 2.x GPU版

教程使用的核心库是tf_geometric，一个基于TensorFlow的GNN库。tf_geometric的详细教程可以在其Github主页上查询：

https://github.com/CrawlScript/tf_geometric

GAT的实现

首先添加自环，即添加节点自身之间的连接边，这样中心节点在稍后的特征更新中也会计算自己原先的特征。

  num_nodes = x.shape[0]
    # self-attention    edge_index, edge_weight = add_self_loop_edge(edge_index, num_nodes)

row为中心节点序列，col为一阶邻居节点序列

 row, col = edge_index

将节点特征向量X通过不同的变换得到Q(Query)，K(Key)和V(value)向量。通过tf.gather得到中心节点的特征向量Q和相应的邻居节点的特征向量K。

    Q = query_activation(x @ query_kernel + query_bias)    Q = tf.gather(Q, row)
    K = key_activation(x @ key_kernel + key_bias)    K = tf.gather(K, col)
    V = x @ kernel

由于是multi-head attention，所以Q，K，V也需要划分为num_heads份，即每一个head都有自己相应的Q，K，V。相应的，为了计算方便，将图节点连接关系矩阵edge_index也进行扩展，每一个head都要对应整个graph。最后将Q，K矩阵相乘（每一个中心节点的特征向量与其邻居节点的特征向量相乘）得到的attention_score，通过segmen_softmax进行归一化操作。

# xxxxx_ denotes the multi-head style stuff    Q_ = tf.concat(tf.split(Q, num_heads, axis=-1), axis=0)    K_ = tf.concat(tf.split(K, num_heads, axis=-1), axis=0)    V_ = tf.concat(tf.split(V, num_heads, axis=-1), axis=0)    edge_index_ = tf.concat([edge_index + i * num_nodes for i in range(num_heads)], axis=1)
    att_score_ = tf.reduce_sum(Q_ * K_, axis=-1)    normed_att_score_ = segment_softmax(att_score_, edge_index_[0], num_nodes * num_heads)

将归一化后的attention系数当做边的权重来对邻居节点进行加权求和操作，从而更新节点特征。由于是multi-head attention，所以将同一个节点在每一个head attention下的节点特征拼接输出。

  h_ = aggregate_neighbors(        V_, edge_index_, normed_att_score_,        gcn_mapper,        sum_reducer,        identity_updater    )
    h = tf.concat(tf.split(h_, num_heads, axis=0), axis=-1)
    if bias is not None:        h += bias
    if activation is not None:        h = activation(h)
    return h

模型构建

导入相关库
本教程使用的核心库是tf_geometric，我们用它来进行图数据导入、图数据预处理及图神经网络构建。GAT的具体实现已经在上面详细介绍，另外我们后面会使用keras.metrics.Accuracy评估模型性能。
```
# coding=utf-8import osos.environ["CUDA_VISIBLE_DEVICES"] = "0"import tf_geometric as tfgimport tensorflow as tffrom tensorflow import keras
```

使用tf_geometric自带的图结构数据接口加载Cora数据集：

graph, (train_index, valid_index, test_index) = CoraDataset().load_data()

定义图模型。我们构建两层GAT，即GAT只聚合2-hop的邻居特征，Dropout层用来缓解模型过拟合(小数据集上尤其管用)。

gat0 = tfg.layers.GAT(64, activation=tf.nn.relu, num_heads=8, drop_rate=drop_rate, attention_units=8)
gat1 = tfg.layers.GAT(num_classes, drop_rate=0.6, attention_units=1)
dropout = keras.layers.Dropout(drop_rate)

def forward(graph, training=False):
    h = graph.x
    h = dropout(h, training=training)
    h = gat0([h, graph.edge_index], training=training)
    h = dropout(h, training=training)
    h = gat1([h, graph.edge_index], training=training)
    return h

GAT训练

模型的训练与其他基于Tensorflow框架的模型训练基本一致，主要步骤有定义优化器，计算误差与梯度，反向传播等。

optimizer = tf.keras.optimizers.Adam(learning_rate=5e-3)for step in range(2000):    with tf.GradientTape() as tape:        logits = forward(graph, training=True)        loss = compute_loss(logits, train_index, tape.watched_variables())
    vars = tape.watched_variables()    grads = tape.gradient(loss, vars)    optimizer.apply_gradients(zip(grads, vars))
    if step % 20 == 0:        accuracy = evaluate()        print("step = {}\tloss = {}\taccuracy = {}".format(step, loss, accuracy))

用交叉熵损失函数计算模型损失。注意在加载Cora数据集的时候，返回值是整个图数据以及相应的train_mask,valid_mask,test_mask。GAT在训练的时候的输入是整个Graph，在计算损失的时候通过train_mask来计算模型在训练集上的迭代损失。因此，此时传入的mask_index是train_index。由于是多分类任务，需要将节点的标签转换为one-hot向量以便于模型输出的结果维度对应。由于图神经模型在小数据集上很容易就会疯狂拟合数据，所以这里用L2正则化缓解过拟合。
```
def compute_loss(logits, mask_index, vars):
    masked_logits = tf.gather(logits, mask_index)
    masked_labels = tf.gather(graph.y, mask_index)
    losses = tf.nn.softmax_cross_entropy_with_logits(
        logits=masked_logits,
        labels=tf.one_hot(masked_labels, depth=num_classes)
    )

    kernel_vals = [var for var in vars if "kernel" in var.name]
    l2_losses = [tf.nn.l2_loss(kernel_var) for kernel_var in kernel_vals]

    return tf.reduce_mean(losses) + tf.add_n(l2_losses) * 5e-4
```

GAT评估

在评估模型性能的时候我们只需传入valid_mask或者test_mask，通过tf.gather函数就可以拿出验证集或测试集在模型上的预测结果与真实标签，用keras自带的keras.metrics.Accuracy计算准确率。

def evaluate(mask):
    logits = forward(graph)
    logits = tf.nn.log_softmax(logits, axis=-1)
    masked_logits = tf.gather(logits, mask)
    masked_labels = tf.gather(graph.y, mask)

    y_pred = tf.argmax(masked_logits, axis=-1, output_type=tf.int32)

    accuracy_m = keras.metrics.Accuracy()
    accuracy_m.update_state(masked_labels, y_pred)
    return accuracy_m.result().numpy()

运行结果

运行结果与论文中结果一致

step = 20	loss = 1.784507393836975	accuracy = 0.7839999794960022
step = 40	loss = 1.5089114904403687	accuracy = 0.800000011920929
step = 60	loss = 1.243167757987976	accuracy = 0.8140000104904175
...
step = 1120	loss = 0.8608425855636597	accuracy = 0.8130000233650208
step = 1140	loss = 0.8169388771057129	accuracy = 0.8019999861717224
step = 1160	loss = 0.7581816911697388	accuracy = 0.8019999861717224
step = 1180	loss = 0.8362383842468262	accuracy = 0.8009999990463257

完整代码

教程中完整代码链接：demo_gat.py:教程代码下载链接：

https://github.com/CrawlScript/tf_geometric/blob/master/demo/demo_gat.py

本教程（属于系列教程《GNN-algorithms》）Github链接：