图注意力神经网络的pytorch代码解析_图注意力机制公式

1.图注意力神经网络的原理简介

图注意网络的原理介绍有很多，可以参考知乎文章：向往的GAT（图注意力模型）。作者是清华大学的一个博士，他写的图卷积原理非常透彻，这里对于图注意力的描述也很好。

为了让后面的代码介绍更清楚，本文再重述一下注意力公式和多头注意力的原理。

1.1 注意力机制的公式

（1）计算注意力系数（attention coefficient）
对于顶点$i$ ，逐个计算它的邻居们（ $j\in N_i$）和它自己之间的相似系数：
$$e_{ij}=a([W{h}_i||W{h}_j]),j\in N_i$$
$h_i$与$h_j$分别为中心节点及其邻居节点的特征。$W$的作用在于对特征进行映射，提高特征的表达能力，$[\ast ||\ast ]$表示拼接，将映射之后的特征进行组合，并通过$a(\ast )$映射成一个实数，作者通过单层前馈神经网络实现。然后通过类似于softmax的方法求解注意力系数：
$${\alpha }_{ij}=\frac{exp(LeakyReLU(e_{ij}))}{{\sum }_{k\in N_i}exp(LeakyReLU(e_{ik}))}$$

（2）特征的聚合
将计算好的注意力系数作为融合权重，对邻居节点的特征进行聚合：
$$h_i^{{}^{\prime }}(K)=\delta (\sum _{j\in N_i}{\alpha }_{ij}W{h}_j)$$
其中$h_i^{{}^{\prime }}$为融合完邻居节点后的顶点$i$的新特征，$\delta (\ast )$为激活函数。

（3）多头注意力机制
多头注意力机制的主要意义在于，多找几个人一起干活，分为治之。
$$h_i^{{}^{\prime }}(K)=\prod _{k=1}^K\delta (\sum _{j\in N_i}{\alpha }_{ij}^k{W}^k{h}_j)$$

1.2 代码中公式的应用差异

多头注意力机制的结合方式

多头的本质是多个独立的attention计算，每个注意力机制函数只负责最终输出序列中一个子空间，且相互独立。所以在经典的多头注意力机制中，输入的参数是需要降维的，最后再拼接到一起。在论文attention is all your need中，输入序列是完全一样的，但最后会只对应于一个子空间，并且相互独立。

这里的注意力有所区别，首先将图的特征经过每个注意力进行学习，然后将结果进行拼接，作为节点的新特征，再次进行注意力学习，后面将在代码中进行进一步阐述。

2.GAT的pytorch代码解析

解析的代码来自gihub，当然论文的源码是用tensorflow写的。用tensorflow的同学可以再挖一挖源码。

2.1 导入需要的包和参数设定

这里参考了GCN的代码设定方式

from __future__ import division
from __future__ import print_function

import os
import glob
import time
import random
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

from utils import load_data, accuracy
from models import GAT, SpGAT

# Training settings
parser = argparse.ArgumentParser()
parser.add_argument('--no-cuda', action='store_true', default=False, help='Disables CUDA training.')
parser.add_argument('--fastmode', action='store_true', default=False, help='Validate during training pass.')
parser.add_argument('--sparse', action='store_true', default=False, help='GAT with sparse version or not.')
parser.add_argument('--seed', type=int, default=72, help='Random seed.')
parser.add_argument('--epochs', type=int, default=10000, help='Number of epochs to train.')
parser.add_argument('--lr', type=float, default=0.005, help='Initial learning rate.')
parser.add_argument('--weight_decay', type=float, default=5e-4, help='Weight decay (L2 loss on parameters).')
parser.add_argument('--hidden', type=int, default=8, help='Number of hidden units.')
parser.add_argument('--nb_heads', type=int, default=8, help='Number of head attentions.')
parser.add_argument('--dropout', type=float, default=0.6, help='Dropout rate (1 - keep probability).')
parser.add_argument('--alpha', type=float, default=0.2, help='Alpha for the leaky_relu.')
parser.add_argument('--patience', type=int, default=100, help='Patience')

args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

random.seed(args.seed) #本函数没有返回值，目的在于给random设置种子，这样生成的随机数会是一样的，后面的np.random.seed和torch.manual_seed是一样的
np.random.seed(args.seed)
torch.manual_seed(args.seed) #为CPU设置种子用于生成随机数，以使得结果是确定的
if args.cuda:
    torch.cuda.manual_seed(args.seed)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

argparse 模块是 Python 内置的一个用于命令项选项与参数解析的模块，argparse 模块可以让人轻松编写用户友好的命令行接口。通过在程序中定义好我们需要的参数，然后 argparse 将会从 sys.argv 解析出这些参数。argparse 模块还会自动生成帮助和使用手册，并在用户给程序传入无效参数时报出错误信息。argparse 模块解释

2.2 加载数据

数据可以从github的源码中下载，这里用的core数据集，做论文分类。
2.2.1 数据格式
数据共有两个文件，为content和Cites。
① centent文件
共有2708行，每一行代表一个样本点，即一篇论文。格式为编号-特征-类别

• 论文的编号，如31336；
• 论文的词向量，一个有1433位的二进制, 表示1433个词汇中的每个单词在文章中是存在(由1表示)还是不存在(由0表示)
• 论文的类别，如Neural_Networks。

31336      0    0.....    0    0    0    0    0    0    0    0    0    0    0    0    Neural_Networks
1061127    0    0.....    0    0    0    0    0    0    0    0    0    0    0    0    Rule_Learning
1106406    0    0.....    0    0    0    0    0    0    0    0    0    0    0    0    Reinforcement_Learning
1
2
3

因此该数据的特征应该有 1433 个维度，另外加上第一个字段 idx，最后一个字段 label， 一共有 1433 + 2 个维度。
② cites文件
共5429行， 每一行有两个论文编号，表示第一个编号的论文先写，第二个编号的论文引用第一个编号的论文。

35    1033
35    103482
35    103515
1
2
3

如果将论文看做图中的点，那么这5429行便是点之间的5429条边。
2.2.3 数据处理代码
加载数据的函数如下：

def load_data(path="./data/cora/", dataset="cora"):
    """Load citation network dataset (cora only for now)"""
    print('Loading {} dataset...'.format(dataset))

    idx_features_labels = np.genfromtxt("{}{}.content".format(path, dataset), dtype=np.dtype(str))
    features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)
    labels = encode_onehot(idx_features_labels[:, -1])

    # build graph
    idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
    idx_map = {j: i for i, j in enumerate(idx)}
    edges_unordered = np.genfromtxt("{}{}.cites".format(path, dataset), dtype=np.int32)
    edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), dtype=np.int32).reshape(edges_unordered.shape)
    adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), shape=(labels.shape[0], labels.shape[0]), dtype=np.float32)

    #scipy.sparse.scs_matrix和coo_matrix都是稀疏矩阵的压缩方式，只不过csr是按行压缩，csc是按列压缩
    #都有data,indptr,indices三个向量，利用indpter向量找到data里的数据和indices的索引，从而按行或者按列构建矩阵

    # build symmetric adjacency matrix
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj) #把矩阵变换成对称矩阵

    features = normalize_features(features)
    adj = normalize_adj(adj + sp.eye(adj.shape[0])) #加入自连接后，构建拉普拉斯矩阵

    idx_train = range(140)
    idx_val = range(200, 500)
    idx_test = range(500, 1500)

    adj = torch.FloatTensor(np.array(adj.todense())) #todense作用是转成稠密矩阵
    features = torch.FloatTensor(np.array(features.todense()))
    labels = torch.LongTensor(np.where(labels)[1]) #np.where(condition)由于没有返回值，
    # 所以返回符合条件的值的位置，为两个向量，第一个向量为横坐标，第二个向量为纵坐标。这里取[1]即表示取纵坐标向量，因为纵坐标的向量代表了其类型

    idx_train = torch.LongTensor(idx_train)
    idx_val = torch.LongTensor(idx_val)
    idx_test = torch.LongTensor(idx_test)

    return adj, features, labels, idx_train, idx_val, idx_test
    
def normalize_adj(mx):
    """Row-normalize sparse matrix"""
    rowsum = np.array(mx.sum(1))
    r_inv_sqrt = np.power(rowsum, -0.5).flatten()
    r_inv_sqrt[np.isinf(r_inv_sqrt)] = 0.
    r_mat_inv_sqrt = sp.diags(r_inv_sqrt)
    return mx.dot(r_mat_inv_sqrt).transpose().dot(r_mat_inv_sqrt)


def normalize_features(mx):
    """Row-normalize sparse matrix"""
    rowsum = np.array(mx.sum(1),dtype= np.float32) #按行求和
    r_inv = np.power(rowsum, -1).flatten()
    r_inv[np.isinf(r_inv)] = 0. #np.isinf()用于判断该数字是否是无穷大或者无穷小。代码表示如果是把对应位置便为其0
    r_mat_inv = sp.diags(r_inv) #构建对角矩阵，并利用矩阵相乘对每个数进行归一化处理
    mx = r_mat_inv.dot(mx)
    return mx
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

如上代码相对明确，并且在代码中加入了注释，值得注意的地方有这么三个：

• 矩阵的对称性，adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)的作用在于使得矩阵变对称，因为这是因为GCN是基于谱域的方法，需要保证矩阵对称，而这里GAT是空域方法，所以并不需要保证对称，可以删除。
• 拉普拉斯矩阵构建，adj = normalize_adj(adj + sp.eye(adj.shape[0]))的作用在于构建拉普拉斯矩阵，前提是矩阵对称，多数情况下，空域里我们并不需要构建。但是注意需要加入自连接，也即可改成adj = adj + sp.eye(adj.shape[0])
• 稀疏矩阵，因为论文的特征向量中大多数为0，少数为1，所以为稀疏矩阵，为了降低存储加入了稀疏矩阵的表示方法，若自己的特征不是稀疏矩阵，可以不用这样表示。

2.3 搭建注意力模型

模型的搭建是逐步完成的，先设计attention layer，然后搭建模型。
2.3.1 GraphAttentionLayer
源代码中设计了两种注意力层，包括正常注意力和针对稀疏矩阵，这里介绍正常的注意层。

class GraphAttentionLayer(nn.Module):
    """
    Simple GAT layer, similar to https://arxiv.org/abs/1710.10903
    """
    def __init__(self, in_features, out_features, dropout, alpha, concat=True):
        super(GraphAttentionLayer, self).__init__()
        self.dropout = dropout
        self.in_features = in_features
        self.out_features = out_features
        self.alpha = alpha
        self.concat = concat

        self.W = nn.Parameter(torch.empty(size=(in_features, out_features)))
        nn.init.xavier_uniform_(self.W.data, gain=1.414)
        self.a = nn.Parameter(torch.empty(size=(2*out_features, 1)))
        nn.init.xavier_uniform_(self.a.data, gain=1.414) #Xavier均匀分布初始化
        #xavier初始化方法中服从均匀分布U(−a,a) ，分布的参数a = gain * sqrt(6/fan_in+fan_out)，这里有一个gain，增益的大小是依据激活函数类型来设定

        self.leakyrelu = nn.LeakyReLU(self.alpha)

    def forward(self, h, adj):
        Wh = torch.mm(h, self.W) # h.shape: (N, in_features), Wh.shape: (N, out_features) #torch.mm矩阵相乘
        a_input = self._prepare_attentional_mechanism_input(Wh)
        e = self.leakyrelu(torch.matmul(a_input, self.a).squeeze(2)) #torch.matmul矩阵乘法，输入可以是高维；squeeze可以将维度为1的那个维度去掉，当输入值中
        #存在dim时，则只有当dim对应的维度为1时会实现降维
        #eij = a([Whi||Whj]),j属于Ni

        zero_vec = -9e15*torch.ones_like(e) #范围维度和e一样的全1的矩阵
        attention = torch.where(adj > 0, e, zero_vec)
        attention = F.softmax(attention, dim=1)
        attention = F.dropout(attention, self.dropout, training=self.training)
        h_prime = torch.matmul(attention, Wh)

        if self.concat:
            return F.elu(h_prime)
        else:
            return h_prime

    def _prepare_attentional_mechanism_input(self, Wh):
        N = Wh.size()[0] # number of nodes        
        Wh_repeated_in_chunks = Wh.repeat_interleave(N, dim=0)
        Wh_repeated_alternating = Wh.repeat(N, 1)
        #epeat_interleave()：在原有的tensor上，按每一个tensor复制。
        #repeat()：根据原有的tensor复制n个，然后拼接在一起。
        all_combinations_matrix = torch.cat([Wh_repeated_in_chunks, Wh_repeated_alternating], dim=1)
        # dim=0表示按行拼接，1表示按列拼接
        # all_combinations_matrix.shape == (N * N, 2 * out_features)
        return all_combinations_matrix.view(N, N, 2 * self.out_features)
        #torch中的view()的作用相当于numpy中的reshape，重新定义矩阵的形状。
    def __repr__(self):
        return self.__class__.__name__ + ' (' + str(self.in_features) + ' -> ' + str(self.out_features) + ')'
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

__init__函数
xavier初始化需要关注一下，“Xavier”初始化方法是一种很有效的神经网络初始化方法，方法来源于2010年的一篇论文《Understanding the difficulty of training deep feedforward neural networks》，其目标是为了使得网络中信息更好的流动，每一层输出的方差应该尽量相等。

映射函数$a$的维度也可以关注一下，为（2*out_features, 1），中心节点的输出特征和邻居节点的特征进行了拼接。最后输出的维度为1，从而得到相关系数。

forward函数
难点在于函数**_prepare_attentional_mechanism_input**的理解，对于通过$W$映射后的out_feature进行两两组合，

• 代码Wh.repeat_interleave(N, dim=0)：在原有的tensor上，按每一个tensor复制，dim为要复制的维度，0即表示行。为如下拼接符号左边部分。
• 代码Wh.repeat(N, 1)：根据原有的tensor复制n个，然后拼接在一起。为如下拼接符号右边部分。

        #e1 || e1
        # e1 || e2
        # e1 || e3
        # ...
        # e1 || eN
        # e2 || e1  ---------------从这里开始换了中心节点
        # e2 || e2
        # e2 || e3
        # ...
        # e2 || eN
        # ...
        # eN || e1 ---------------从这里开始换了中心节点
        # eN || e2
        # eN || e3
        # ...
        # eN || eN
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

这样通过torch.cat([Wh_repeated_in_chunks, Wh_repeated_alternating], dim=1)可将$N$个节点拼接成如上形式。最后通过all_combinations_matrix.view（）函数将拼接的矩阵重新reshape，这样变成(N, N, 2 * self.out_features)维度后，第一个维度的矩阵则为第一个节点及其邻居节点构成的特征向量。

后面则是通过映射函数$a(\ast )$（其实是一层全连接）对特征进行隐射求解注意力系数。大家可以发现，这里是任意一个节点都计算了全局的节点的相关性，并没有利用到结构信息，因此为了抛出掉那些没有在连接关系范围内的注意力系数，利用attention = torch.where(adj > 0, e, zero_vec)进行修正。

最后将注意力系数进行softmax后求解新的特征向量。

2.3.2 GAT模型搭建
考虑了多头注意力，代码如下

class GAT(nn.Module):
    def __init__(self, nfeat, nhid, nclass, dropout, alpha, nheads):
        """Dense version of GAT."""
        super(GAT, self).__init__()
        self.dropout = dropout
        self.attentions = [GraphAttentionLayer(nfeat, nhid, dropout=dropout, alpha=alpha, concat=True) for _ in range(nheads)]
        for i, attention in enumerate(self.attentions):
            self.add_module('attention_{}'.format(i), attention)
            #Module.add_module(name: str, module: Module)。功能为，为Module添加一个子module，对应名字为name。

        self.out_att = GraphAttentionLayer(nhid * nheads, nclass, dropout=dropout, alpha=alpha, concat=False)

    def forward(self, x, adj):
        x = F.dropout(x, self.dropout, training=self.training)
        x = torch.cat([att(x, adj) for att in self.attentions], dim=1)
        x = F.dropout(x, self.dropout, training=self.training)
        x = F.elu(self.out_att(x, adj))
        return F.log_softmax(x, dim=1)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

利用add_module（）函数添加了nheads个注意力层，输入输出维度为（nfeat, nhid）。在搭建模型时，对结果进行了按行拼接，从而得到（nhid*nheads）为的tensor。

代码在dropout后，将特征重新输给GraphAttentionLayer，输出维度定义为nclass即标签的类别数，这样将多个注意力的结果进行融合，并利用log_softmax输出结果向量。

2.3 模型训练

首先利用类生成模型 ，并如果有cuda可以用下gpu。

model = GAT(nfeat=features.shape[1], 
                nhid=args.hidden, 
                nclass=int(labels.max()) + 1, 
                dropout=args.dropout, 
                nheads=args.nb_heads, 
                alpha=args.alpha)
                
optimizer = optim.Adam(model.parameters(), 
                       lr=args.lr, 
                       weight_decay=args.weight_decay) #权重衰减，为L2正则化
               
if args.cuda:
    model.cuda()
    features = features.cuda()
    adj = adj.cuda()
    labels = labels.cuda()
    idx_train = idx_train.cuda()
    idx_val = idx_val.cuda()
    idx_test = idx_test.cuda()

features, adj, labels = Variable(features), Variable(adj), Variable(labels)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

定义epoch，训练模型，这里使用F.nll_loss，NLLLoss 的输入是一个对数概率向量和一个目标标签， 它不会为我们计算对数概率. 适合网络的最后一层是log_softmax。 损失函数 nn.CrossEntropyLoss() 与 NLLLoss() 相同, 唯一的不同是它为我们去做 softmax。

def train(epoch):
    t = time.time()
    model.train()
    optimizer.zero_grad()
    output = model(features, adj)
    loss_train = F.nll_loss(output[idx_train], labels[idx_train])
    acc_train = accuracy(output[idx_train], labels[idx_train])
    loss_train.backward()
    optimizer.step()

    if not args.fastmode:
        # Evaluate validation set performance separately,
        # deactivates dropout during validation run.
        model.eval() #训练完train样本后，生成的模型model要用来测试样本。在model(test)之前，需要加上model.eval()，否则的话，有输入数据，即使不训练，它也会改变权值。
        output = model(features, adj)

    loss_val = F.nll_loss(output[idx_val], labels[idx_val]) #得到batch内的均值
    acc_val = accuracy(output[idx_val], labels[idx_val])
    print('Epoch: {:04d}'.format(epoch+1),
          'loss_train: {:.4f}'.format(loss_train.data.item()),
          'acc_train: {:.4f}'.format(acc_train.data.item()),
          'loss_val: {:.4f}'.format(loss_val.data.item()),
          'acc_val: {:.4f}'.format(acc_val.data.item()),
          'time: {:.4f}s'.format(time.time() - t))

    return loss_val.data.item()
    #.data返回的是一个tensor,而.item()返回的是一个具体的数值。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

这里训练和验证时都用了output = model(features, adj)作为结果，因为训练用的是一些节点，验证是另外一些节点。train（）函数会在训练集上进行模型训练，然而会返回验证集的准确率作为返回值。
accuracy函数是自定义的，主要是计算分类对的比例。

def accuracy(output, labels):
    preds = output.max(1)[1].type_as(labels)
    correct = preds.eq(labels).double()
    correct = correct.sum()
    return correct / len(labels)
1
2
3
4
5

2.5 模型测试

测试用的结果和训练验证时的一样，只是标签不同而已。

def compute_test():
    model.eval()
    output = model(features, adj)
    loss_test = F.nll_loss(output[idx_test], labels[idx_test])
    acc_test = accuracy(output[idx_test], labels[idx_test])
    print("Test set results:",
          "loss= {:.4f}".format(loss_test.data[0]),
          "accuracy= {:.4f}".format(acc_test.data[0]))

# Train model
t_total = time.time()
loss_values = []
bad_counter = 0
best = args.epochs + 1
best_epoch = 0
for epoch in range(args.epochs):
    loss_values.append(train(epoch))

    torch.save(model.state_dict(), '{}.pkl'.format(epoch))
    if loss_values[-1] < best:
        best = loss_values[-1]
        best_epoch = epoch
        bad_counter = 0
    else:
        bad_counter += 1

    if bad_counter == args.patience:
        break

    files = glob.glob('*.pkl')
    for file in files:
        epoch_nb = int(file.split('.')[0])
        if epoch_nb < best_epoch:
            os.remove(file)

files = glob.glob('*.pkl')
for file in files:
    epoch_nb = int(file.split('.')[0])
    if epoch_nb > best_epoch:
        os.remove(file)

print("Optimization Finished!")
print("Total time elapsed: {:.4f}s".format(time.time() - t_total))

# Restore best model
print('Loading {}th epoch'.format(best_epoch))
model.load_state_dict(torch.load('{}.pkl'.format(best_epoch)))

# Testing
compute_test()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

如上是测试和记录数据/模型的过程，到这一步，模型就完成了。

备注
作者是个学习者，若有错误请包涵。