【GAT】图注意力网络 - 简单的源码记录

由于和GCN代码比较相似，所以部分内容从GCN那篇博客中截取。

1 - cora数据集

GNN常用数据集之Cora数据集

2 - 源码含义记录

首先我们来整体看一下代码的组成

截图中的这一大坨为命令行传递参数，含义参考命令行传递参数 argparse.ArgumentParser解析简单点说，就是想在不改动代码的情况下，使用命令行去改参数。

2.1 加载数据集

在代码中，加载数据集通过这个函数实现

# Load data
adj, features, labels, idx_train, idx_val, idx_test = load_data()

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)

    # 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()))
    features = torch.FloatTensor(np.array(features.todense()))
    labels = torch.LongTensor(np.where(labels)[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

np.genfromtxt函数为numpy加载数据集，当然还有其它几种加载数据集的方式，例如pandas等等，详情见Python加载数据的5种不同方式(收藏)。

2.1.1 加载节点数据

首先是加载节点的数据，即.content。

通过断点+命令行调试，可以看到输出的idx_features_labels的结果

因此我们想把中间节点的特征给取出来，即这一句

features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)

中间[1：-1]是从第二列到倒数第二列（因为python为左闭右开）

然后再取label，使用

labels = encode_onehot(idx_features_labels[:, -1])

只取最后一列，并且采用one-hot编码。

分别看看上述两行代码的输出
解释一下特征的含义

最后开始取节点的索引(index)，即第一列，并且构建节点的索引字典。

idx = np.array(idx_features_labels[:, 0], dtype=np.int32)  # 取节点
idx_map = {j: i for i, j in enumerate(idx)}  # 构建节点的索引字典

看看输出，字典就构建好啦

2.1.2 加载边的数据

节点的数据导入完了，我们再来导入边的数据

 edges_unordered = np.genfromtxt("{}{}.cites".format(path, dataset), dtype=np.int32)  # 导入edge的数据

将之前的转换成字典编号后的边

2.1.3 构造邻接矩阵

adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), 
                        shape=(labels.shape[0], labels.shape[0]),
                        dtype=np.float32)  # 构造边的邻接矩阵（稀疏矩阵）

看看邻接矩阵的大小和样子，因为有2708个节点，所以大小为(2708×2708)

由于Cora数据集是一个有向图，而GCN本身使用的为无向图，其邻接矩阵是对称的，因此我们需要构造一个对称矩阵，如下：

# build symmetric adjacency matrix  计算邻接矩阵的对称矩阵，因为cora是有向图，而GCN是无向图
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)

其中adj.T > adj的含义为，如果转置后为1，而原矩阵为0，则该位置为1
Python将非对称邻接矩阵转变为对称邻接矩阵(有向图转无向图)

2.1.4 归一化

features = normalize(features)
    adj = normalize(adj + sp.eye(adj.shape[0]))

我们可以看一下normalize函数

def normalize(mx):
    """Row-normalize sparse matrix"""
    rowsum = np.array(mx.sum(1))  # 矩阵行求和
    r_inv = np.power(rowsum, -1).flatten()  # 求和的-1次方，即倒数
    r_inv[np.isinf(r_inv)] = 0.  # 求倒数的时候有可能是无穷，因此如果是无穷的话，就转换成0
    r_mat_inv = sp.diags(r_inv)  # 构造对角矩阵
    mx = r_mat_inv.dot(mx)  # 构造D-1 * A，特征归一化，非对称方式，简化方式
    return mx

分别看看rowsum和r_inv的输出

然后我们在把r_inv中的无穷值变为0后，构造对角矩阵，可以看到结果

特征矩阵归一化（非必须）

做了上述这些后，我们对特征矩阵进行归一化
可以发现特征每一行加起来都是1，已经被归一化好了。

邻接矩阵归一化（必须）

再对邻接矩阵+单位阵进行归一化（必须操作）

首先看起来邻接矩阵+单位阵的值：

然后再进行归一化，这里用的是简便方法，和上面一样啦~

2.1.5 加载数据集/数据格式转换

    # 训练、验证、测试的样本
    idx_train = range(140)
    idx_val = range(200, 500)
    idx_test = range(500, 1500)

    # 将numpy的数据转换成torch格式
    features = torch.FloatTensor(np.array(features.todense()))
    labels = torch.LongTensor(np.where(labels)[1])
    adj = sparse_mx_to_torch_sparse_tensor(adj)

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

最后返回的格式

    return adj, features, labels, idx_train, idx_val, idx_test

2.2 构造模型

我们构造一个GAT模型和优化器

# Model and optimizer
if args.sparse:
    model = SpGAT(nfeat=features.shape[1], 
                nhid=args.hidden, 
                nclass=int(labels.max()) + 1, 
                dropout=args.dropout, 
                nheads=args.nb_heads, 
                alpha=args.alpha)
else:
    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)

看一下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)  # 这里就有8个attention layer，第一层的attention layer

        self.out_att = GraphAttentionLayer(nhid * nheads, nclass, dropout=dropout, alpha=alpha, concat=False)  # 第二层（最后一层）的attention layer

    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)

首先设置dropout，然后比较重要的就是接下来的

[GraphAttentionLayer(nfeat, nhid, dropout=dropout, alpha=alpha, concat=True) for _ in range(nheads)]

通过断点调试，可以看到上述函数中几个参数的值以及含义

• nfeat：输入特征=1433
• nhid：隐藏特征=8
• nheads：多头向量=8，代表有8个attentionlayer

多头操作如下：
了解了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)))  # in_features=1433, out_features=8
        nn.init.xavier_uniform_(self.W.data, gain=1.414)  # 初始化W向量
        self.a = nn.Parameter(torch.empty(size=(2*out_features, 1)))  # 建立一个需要训练的α向量，因为是将两个向量Whi和Whj拼接到一起，所以维度是2×out_features   self.a.shape = torch.Size([16, 1])
        nn.init.xavier_uniform_(self.a.data, gain=1.414) # 初始化α向量，赋予随机的值

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

    def forward(self, h, adj):
        Wh = torch.mm(h, self.W) # h.shape: (2708, 1433)  self.W.shape (1433.8)  W*h = Wh
        e = self._prepare_attentional_mechanism_input(Wh)  # 每一个节点和所有节点，特征拼接？
#  之前计算的是一个节点和所有节点的attention，其实需要的是连接的节点的attention系数
        zero_vec = -9e15*torch.ones_like(e)  # 建立了一个极小值矩阵
        attention = torch.where(adj > 0, e, zero_vec)  # 将邻接矩阵中小于0的变成负无穷
        attention = F.softmax(attention, dim=1)  # 按行求softmax，sum(axis=1) === 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):
        # Wh.shape (N, out_feature)
        # self.a.shape (2 * out_feature, 1)
        # Wh1&2.shape (N, 1)
        # e.shape (N, N)
        Wh1 = torch.matmul(Wh, self.a[:self.out_features, :])
        Wh2 = torch.matmul(Wh, self.a[self.out_features:, :])
        # broadcast add
        e = Wh1 + Wh2.T
        return self.leakyrelu(e)

    def __repr__(self):
        return self.__class__.__name__ + ' (' + str(self.in_features) + ' -> ' + str(self.out_features) + ')'

2.3 使用GPU

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)

2.4 定义训练和测试函数

def train(epoch):
    t = time.time()
    model.train()
    optimizer.zero_grad()
    output = model(features, adj)  # features [2708, 1433]
    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()
        output = model(features, adj)

    loss_val = F.nll_loss(output[idx_val], labels[idx_val])
    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()

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.item()),
          "accuracy= {:.4f}".format(acc_test.data.item()))

2.5 模型训练与测试

# 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()