基于Pytorch Gemotric在昇腾上实现GAT图神经网络

本实验主要介绍了如何在昇腾上，使用pytorch对经典的图神经网络GAT在论文引用数据集Pubmed上进行分类训练的实战讲解。内容包括GAT网络创新点分析、图注意力机制原理与架构剖析、多头注意力机制分析与GAT网络模型代码实战分析等等。

本实验的目录结构安排如下所示：

• GAT网络创新点分析
• 图注意力机制原理与架构分析
• 多头注意力机制分析
• GAT网络用于Pubmed数据集分类实战

GAT网络创新点分析

注意机制已成功用于许多基于序列的任务，例如机器翻译，机器阅读等等。与GCN平等对待节点的所有邻居相比，注意力机制可以为每个邻居分配不同的注意力得分，从而识别出更重要的邻居。将注意力机制纳入图谱神经网络的传播步骤是很直观的。图注意力网络也可以看作是图卷积网络家族中的一种方法。其创新点如下：

• GAT是一种图结构化数据上操作的新型神经网络架构，利用掩码自注意力层来解决基于图卷积或其近似值的现有方法的缺点。
• GAT网络对不同的相邻节点分配相应的权重，既不需要矩阵运算，也不需要事先知道图结构。
• GAT网络克服了基于谱神经网络的几个关键挑战，使得模型更加适用于归纳问题以及转导问题。
• 在Cora、Citeseer和Pubmed引文网络数据集上取得了非常好的效果。

图注意力机制原理与架构分析

图中输入由N个节点组成，每个节点feature数为F(也就是feature vector长度)，下述两个公式分别表示输入的节点向量通过GAT网络进行注意力计算后得到的输出。

　　　　Input: $h = { \\vec{h_1}, \\vec{h_2}, . . . , \\vec{h_N} },\\vec h_i \\in R\^{F} $

　　　　Output: $h = { \\vec{h_1\^{'}}, \\vec{h_2\^{'}}, . . . , \\vec{h_N\^{'}} },\\vec h_i\^{'} \\in R\^{F\^{'}} $

为了将输入的特征转换为高维的特征，这里至少需要一个科学系的线性转换。在 (Velickovic et al.,2017)中，作者对于每一个点使用了一个共享的线性转换方式，同时介绍了一个权重矩阵$W$来参数化线性转换。对于节点中的每两个node间都有了相互关注机制（用来做加权平均，卷积时每个node的更新是其他的加权平均）。

计算节点$i$与节点$j$的相关性系数(注意力值)：

先初步计算节点i与节点j的相关性：$e_{ij} = a(W \* h_i ,W \* h_j)$其中a是一个共享的权重系数，执行的是$R\^{F\^{'}} X R\^{F\^{'}} -\> R\^{F} 转换，这里可以定义成一个前馈全连接神经网络$

计算节点i与节点j的相关性(归一化)$ \\alpha_{i, j} = softmax_j(e_{ij}) = \\frac {exp(e_{ij})} {\\sum_{k \\in N_i}exp(e_{ik})} $

论文中采取取的计算attention coefficient的函数a是一个单层的前馈网络，经LeakyReLU处理得最终的$\\alpha_{i, j}值$：

$\\alpha_{i, j} = \\frac {exp(LeakyReLU({\\vec a\^T}\[W\\vec{h_i} \|\| W \\vec{h_j}\]))} {\\sum_{k \\in N_i}exp(LeakyReLU(\\vec a\^T\[W \\vec{h_i} ,W \\vec{h_k}\])} $

式中 || 表示串联/连接，一旦获得，归一化的相互注意系数用来计算对应特征的线性组合，以用作每个节点的最终输出特征。

左图表示在模型中应用注意机制a(W * h_i ,W * h_j) 通过权重向量$a$参数化，应用LeakyReLU 激活输出。

右图表示$h_1\^′$在邻域中具有多端连接，不同的箭头样式表示独立的注意力计算，通过直连concat或平均avg获取$h_1\^′$。

多头注意力机制分析

对于图中多头注意力情况，对应上图中右图情景，不只用一个函数进行attention coefficient的计算，而是设置K个函数，每一个函数都能计算出一组attention coefficient，并能计算出一组加权求和用的系数，每一个卷积层中，K个attention机制独立的工作，分别计算出自己的结果后连接在一起，得到卷积的结果。

我们知道对于单个注意力层输出$\\vec h_i\^′$计算如下：

$\\vec h_i\^′ = \\delta (\\sum_{k \\in N_i}(a_{ij}W\\vec h_j))$

对于有k个独立的相互注意机制同时计算，则集中其特征，可得到特征表示如下：

分别遍历1\~k个头，每一个上按照下述公式计算:

$\\vec h_i\^′ = \\delta (\\sum_{k \\in N_i}(a_{ij}\^k W\^k \\vec h_j))$

对于网络的最后一层卷积层，如果还是使用multi-head attention机制，那么就不采取连接的方式合并不同的attention机制的结果了，而是采用求平均的方式进行处理。

$\\vec h_i\^′ = \\delta({\\frac {1}{K}}(\\sum_{k=1}\^K)(\\sum_{k \\in N_i}(a_{ij}\^k W\^k \\vec h_j)))$

GAT网络用于Pubmed数据集分类实战

#导入torch相关库
import torch
import torch.nn.functional as F

该实验需要跑在npu上，因此需要导入Npu相关库使得模型快速迁移到npu上运行

import torch_npu
from torch_npu.contrib import transfer_to_npu

/home/pengyongrong/miniconda3/envs/AscendCExperiments/lib/python3.9/site-packages/torch_npu/dynamo/__init__.py:18: UserWarning: Register eager implementation for the 'npu' backend of dynamo, as torch_npu was not compiled with torchair.
  warnings.warn(
/home/pengyongrong/miniconda3/envs/AscendCExperiments/lib/python3.9/site-packages/torch_npu/contrib/transfer_to_npu.py:164: ImportWarning: 
    *************************************************************************************************************
    The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..
    The torch.cuda.DoubleTensor is replaced with torch.npu.FloatTensor cause the double type is not supported now..
    The backend in torch.distributed.init_process_group set to hccl now..
    The torch.cuda.* and torch.cuda.amp.* are replaced with torch.npu.* and torch.npu.amp.* now..
    The device parameters have been replaced with npu in the function below:
    torch.logspace, torch.randint, torch.hann_window, torch.rand, torch.full_like, torch.ones_like, torch.rand_like, torch.randperm, torch.arange, torch.frombuffer, torch.normal, torch._empty_per_channel_affine_quantized, torch.empty_strided, torch.empty_like, torch.scalar_tensor, torch.tril_indices, torch.bartlett_window, torch.ones, torch.sparse_coo_tensor, torch.randn, torch.kaiser_window, torch.tensor, torch.triu_indices, torch.as_tensor, torch.zeros, torch.randint_like, torch.full, torch.eye, torch._sparse_csr_tensor_unsafe, torch.empty, torch._sparse_coo_tensor_unsafe, torch.blackman_window, torch.zeros_like, torch.range, torch.sparse_csr_tensor, torch.randn_like, torch.from_file, torch._cudnn_init_dropout_state, torch._empty_affine_quantized, torch.linspace, torch.hamming_window, torch.empty_quantized, torch._pin_memory, torch.Tensor.new_empty, torch.Tensor.new_empty_strided, torch.Tensor.new_full, torch.Tensor.new_ones, torch.Tensor.new_tensor, torch.Tensor.new_zeros, torch.Tensor.to, torch.nn.Module.to, torch.nn.Module.to_empty
    *************************************************************************************************************
    
  warnings.warn(msg, ImportWarning)

由于torch_geometric中集成了单层的GATConv模块，这里直接进行导入，若有兴趣可以自行实现该类，注意输入与输出对齐即可。此外，数据集用的是Pubmed，该数据集也直接集成在Planetoid模块中，这里也需要将其import进来。

from torch_geometric.nn import GATConv
from torch_geometric.datasets import Planetoid

/home/pengyongrong/miniconda3/envs/AscendCExperiments/lib/python3.9/site-packages/torch_npu/contrib/transfer_to_npu.py:124: RuntimeWarning: torch.jit.script will be disabled by transfer_to_npu, which currently does not support it.
  warnings.warn(msg, RuntimeWarning)

Pubmed数据集介绍

PubMed数据集是一个广泛用于图神经网络（GNN）研究的基准数据集，主要用于节点分类任务。其由生物医学文献组成，每篇文献被视为一个节点，引用关系被视为边。该数据集包含三类糖尿病相关的论文，每个节点都带有特征向量和标签。

数据集一共有19717个节点 ，每一个节点代表一篇生物医学文献，每个节点有一个500 维的特征向量，用来表示该医学文献的内容。总共含44338条边 ，每条边表示一篇文献对另一篇文献的引用关系，边与边之间是无向的，因此可以看做是对称的。总类别数包含Diabetes Mellitus Experiment、Diabetes Mellitus Type 1与Diabetes Mellitus Type2三类。

# 加载数据
print("===== begin Download Dadasat=====\n")
dataset = Planetoid(root='/home/pengyongrong/workspace/data', name='PubMed')
print("===== Download Dadasat finished=====\n")
print("dataset num_features  is:  ", dataset.num_features)
print("dataset.num_classes is：  ", dataset.num_classes)

print("dataset.edge_index is：  ", dataset.edge_index)

print("train data is:   ", dataset.data)
print("dataset0 is:  ", dataset[0])

print("train data mask is:   ", dataset.train_mask, "num train is： ", (dataset.train_mask ==True).sum().item())
print("val data mask is:   ",dataset.val_mask, "num val is： ", (dataset.val_mask ==True).sum().item())
print("test data mask is:   ",dataset.test_mask,  "num test is： ", (dataset.test_mask ==True).sum().item())

===== begin Download Dadasat=====

===== Download Dadasat finished=====

dataset num_features  is:   500
dataset.num_classes is：   3
dataset.edge_index is：   tensor([[ 1378,  1544,  6092,  ..., 12278,  4284, 16030],
        [    0,     0,     0,  ..., 19714, 19715, 19716]])
train data is:    Data(x=[19717, 500], edge_index=[2, 88648], y=[19717], train_mask=[19717], val_mask=[19717], test_mask=[19717])
dataset0 is:   Data(x=[19717, 500], edge_index=[2, 88648], y=[19717], train_mask=[19717], val_mask=[19717], test_mask=[19717])
train data mask is:    tensor([ True,  True,  True,  ..., False, False, False]) num train is：  60
val data mask is:    tensor([False, False, False,  ..., False, False, False]) num val is：  500
test data mask is:    tensor([False, False, False,  ...,  True,  True,  True]) num test is：  1000


/home/pengyongrong/miniconda3/envs/AscendCExperiments/lib/python3.9/site-packages/torch_geometric/data/in_memory_dataset.py:300: UserWarning: It is not recommended to directly access the internal storage format `data` of an 'InMemoryDataset'. If you are absolutely certain what you are doing, access the internal storage via `InMemoryDataset._data` instead to suppress this warning. Alternatively, you can access stacked individual attributes of every graph via `dataset.{attr_name}`.
  warnings.warn(msg)

下载后的pubmed数据集总共包含8个文件，分别是ind.pubmed.x、ind.pubmed.tx、ind.pubmed.all、ind.pubmed.y、ind.pubmed.ty、ind.pubmed.ally、ind.pubmed.graph与ind.pubmed.test.index。每个文件的作用说明如下：

ind.pubmed.x：训练集节点特征向量，大小（140,1433）

ind.pubmed.tx：测试集节点特征向量，实际展开后大小为（1000,1433）

ind.pubmed.allx：包含标签核无标签的训练节点特征向量（1708,1433）

ind.pubmed.y：one-hot表示的训练节点的标签

ind.pubmed.ty：one-hot表示的测试节点的标签

ind.pubmed.ally：one-hot表示的ind.cora.allx对应的标签

ind.pubmed.graph：保存节点之间边的信息

ind.pubmed.test.index：保存测试集节点的索引，用于后面的归纳学习设置

从打印结果可以看出，数据集的特点与上述描述的相对应，GAT_NET网络定义了一个两层的GAT网络，heads的数量设置成4。

开启Pubmed数据训练过程

class GAT_NET(torch.nn.Module):
    def __init__(self, features, hidden, classes, heads=4):
        super(GAT_NET, self).__init__()
        # 定义GAT层，使用多头注意力机制
        self.gat1 = GATConv(features, hidden, heads=4) 
        # 因为多头注意力是将向量拼接，所以维度乘以头数。
        self.gat2 = GATConv(hidden*heads, classes) 

    def forward(self, data):
        # 从输入数据集中获取x与边集相关信息
        x, edge_index = data.x, data.edge_index
        # 将输入传入GAT层中，获得第一层Gat层的输出
        x = self.gat1(x, edge_index)
        # 经过非线性激活与dropout，减少过拟合现象，增加模型的泛化能力
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        # 第二层GAT层，得到整个网络的输出送给分类器进行分类
        x = self.gat2(x, edge_index)
        return F.log_softmax(x, dim=1)

定义设备跑在Npu上，这里如果需要替换成Gpu或Cpu，则替换成'cuda'或'cpu'即可。

device = 'npu'

定义GAT_NET网络，中间隐藏层节点个数定义为16，'dataset.num_classes'为先前数据集中总的类别数，这里是7类。'to()'的作用是将该加载到指定模型设备上。优化器用的是'optim'中的'Adam'。

model = GAT_NET(dataset.num_node_features, 16, dataset.num_classes).to(device)  # 定义GraphSAGE
data = dataset[0].to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

开始训练模型，指定训练次数200次，训练后采用极大似然用作损失函数计算损失，然后进行反向传播更新模型的参数，训练完成后，用验证集中的数据对模型效果进行验证，最后打印模型的准确率。

model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

# 模型验证过程，对训练得到的模型效果进行评估，并打印准确率。
model.eval()
_, pred = model(data).max(dim=1)
correct = int(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / int(data.test_mask.sum())
print('GAT Accuracy: {:.4f}'.format(acc))

[W VariableFallbackKernel.cpp:51] Warning: CAUTION: The operator 'aten::scatter_reduce.two_out' is not currently supported on the NPU backend and will fall back to run on the CPU. This may have performance implications. (function npu_cpu_fallback)


GAT Accuracy: 0.3660

内存使用情况： 整个训练过程的内存使用情况可以通过"npu-smi info"命令在终端查看，因此本文实验只用到了单个npu卡（也就是chip 0），内存占用约573M，对内存、精度或性能优化有兴趣的可以自行尝试进行优化。

Reference

[1] Velikovi, Petar , et al. "Graph Attention Networks." (2017).