【图表示学习】pytorch实现图注意力网络GAT_pytorch实现gat

本文的实现只是对论文的原始代码进行了简单的修改、简化和说明，原始代码点这里！

一、实现前的一些说明

在实现前，先对代码中比较难理解的部分进行说明。

1. 任意两节点的$\text{W}\vec{h}$的拼接

计算注意力系数的公式如下：
$${\alpha }_{ij}=\frac{exp(LeakyReLU({\vec{\text{a}}}^T[\text{W}{\vec{h}}_i||\text{W}{\vec{h}}_j]))}{{\sum }_{k\in {\mathcal{N}}_i}exp(LeakyReLU({\vec{\text{a}}}^T[\text{W}{\vec{h}}_i||\text{W}{\vec{h}}_k]))}$$
可以看到在计算时，需要对任意两个节点进行拼接操作，$\text{W}{\vec{h}}_i||\text{W}{\vec{h}}_j$。

这里假设有3个节点，每个节点的特征向量维度为5，即假设${\vec{h}}_i\in {\mathbb{R}}^5,1\le i\le 3$，且令$\text{H}=[{\vec{h}}_1,\dots ,{\vec{h}}_3]\in {\mathbb{R}}^{5\times 3}$,$\text{W}\in {\mathbb{R}}^{5\times 5}$。

则任意两个节点进行拼接操作相当于对矩阵$(\text{WH}{)}^T\in {\mathbb{R}}^{3\times 5}$的行向量进行两两拼接，这里先定义$(\text{WH}{)}^T$并展示如何进行两两拼接。

WH = torch.arange(0,3).repeat(5,1).T # 3个节点
# 两种不同的repeat方式
Wh_repeated_in_chunks = WH.repeat_interleave(3, dim=0)
Wh_repeated_alternating = WH.repeat(3,1)
# Wh两两拼接
all_combinations_matrix = torch.cat([Wh_repeated_in_chunks, Wh_repeated_alternating], dim=1)
result = all_combinations_matrix.view(3, 3, 2 * 5)
print(result)
1
2
3
4
5
6
7
8

输出：

tensor([[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
         [0, 0, 0, 0, 0, 2, 2, 2, 2, 2]],

        [[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 2, 2, 2, 2, 2]],

        [[2, 2, 2, 2, 2, 0, 0, 0, 0, 0],
         [2, 2, 2, 2, 2, 1, 1, 1, 1, 1],
         [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]]])
1
2
3
4
5
6
7
8
9
10
11

2.如何聚合相邻节点

$$\overrightarrow{h_i^{\prime }}=\sigma \left(\sum _{j\in {\mathcal{N}}_i}{\alpha }_{ij}\text{W}\overrightarrow{h_j}\right)$$

这里假设是3个节点，其两两对应的注意力系数组成的矩阵为$A\in {\mathbb{R}}^{3\times 3}$，其中$A_{i,j}$表示节点$i$与$j$的注意力系数。邻居矩阵则同样是一个$3\times 3$的矩阵。

Wh = torch.randn(3,5) # 3个节点，每个节点5个特征
A = torch.randn(3,3) # 注意力系数矩阵
# 邻接矩阵
adj = torch.tensor([[0,1,1],
                    [1,0,0],
                    [1,0,0]])
zero_vec = -9e15*torch.ones_like(A)
# 使用adj作为掩码，将没有边连接的点对的注意力系数置为0
attention = torch.where(adj>0, A, zero_vec)
attention = F.softmax(attention, dim=1)
# h_prime.shape=(3,5)，得到了每个节点的聚合新特征
h_prime = torch.matmul(attention, Wh)
print(h_prime)
1
2
3
4
5
6
7
8
9
10
11
12
13

输出：

tensor([[-1.0747, -0.0669,  0.9770,  0.0211, -1.6639],
        [-0.0577, -0.1590, -0.0546,  2.2421, -2.0541],
        [-0.0577, -0.1590, -0.0546,  2.2421, -2.0541]])
1
2
3

3.节点特征向量的标准化

假设有3个节点，每个节点的特征向量维度为5，即假设${\vec{h}}_i\in {\mathbb{R}}^5,1\le i\le 3$，且令$\text{H}=[{\vec{h}}_1,\dots ,{\vec{h}}_3{]}^T\in {\mathbb{R}}^{3\times 5}$。对于单个节点的标准化就是$\frac{{\vec{h}}_1}{sum({\vec{h}}_1)}$，那么以矩阵运算的方法进行标准化。

H = torch.ones(3,5)
# 特征求和
rowsum = np.array(H.sum(1))
# 倒数
r_inv = np.power(rowsum,-1).flatten()
# 解决除0问题
r_inv[np.isinf(r_inv)] = 0.
# 转换为对角阵
r_mat_inv = np.diag(r_inv)
# 对角阵乘以H，得到标准化矩阵
H = r_mat_inv.dot(H)
print(H)
1
2
3
4
5
6
7
8
9
10
11
12

输出：

[[0.2 0.2 0.2 0.2 0.2]
 [0.2 0.2 0.2 0.2 0.2]
 [0.2 0.2 0.2 0.2 0.2]]
1
2
3

4.邻接矩阵的标准化

假设有4个节点，邻居矩阵$A\in {\mathbb{R}}^{4\times 4}$，那么邻接矩阵标准化

A = np.ones((4,4))
rowsum = A.sum(1)
r_inv_sqrt = np.power(rowsum, -0.5).flatten()
r_inv_sqrt[np.isinf(r_inv_sqrt)] = 0.
r_mat_inv_sqrt = np.diag(r_inv_sqrt)
A = A.dot(r_mat_inv_sqrt).transpose().dot(r_mat_inv_sqrt)
print(A)
1
2
3
4
5
6
7

输出：

[[0.25 0.25 0.25 0.25]
 [0.25 0.25 0.25 0.25]
 [0.25 0.25 0.25 0.25]
 [0.25 0.25 0.25 0.25]]
1
2
3
4

二、完整实现

1.数据加载和预处理

def load_data(path="./data/cora/", dataset="cora"):
    """Load citation network dataset (cora only for now)"""
    print('Loading {} dataset...'.format(dataset))
    
    # idx_features_labels.shape = (2708,1435)；
    # 第一列是节点编号，最后一列是节点类别，中间列是节点的特征
    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)
    # 将标签转换为one-hot编码
    labels = pd.get_dummies(idx_features_labels[:,-1]).values

    # build graph
    # 节点编号
    idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
    # (节点编号：出现顺序)
    idx_map = {j: i for i, j in enumerate(idx)}
    # 边表，shape = (5429,2)
    edges_unordered = np.genfromtxt("{}{}.cites".format(path, dataset), dtype=np.int32)
    # 使用idx_map映射edges_unordered中节点的编号
    edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), dtype=np.int32).reshape(edges_unordered.shape)
    # 将edges转换为邻接矩阵
    adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), shape=(labels.shape[0], labels.shape[0]), dtype=np.float32)

    # 转换为对称矩阵
    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


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))
    r_inv = np.power(rowsum, -1).flatten()
    r_inv[np.isinf(r_inv)] = 0.
    r_mat_inv = sp.diags(r_inv)
    mx = r_mat_inv.dot(mx)
    return mx

adj, features, labels, idx_train, idx_val, idx_test = load_data()
print(adj.shape)
print(features.shape)
print(labels.shape)
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
57
58
59
60
61
62
63
64
65
66
67

输出：

torch.Size([2708, 2708])
torch.Size([2708, 1433])
torch.Size([2708])
1
2
3

2.定义模型结构

2.1 定义图注意力层(Graph Attention Layer)

class GraphAttentionLayer(nn.Module):
    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)
        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)
        a_input = self._prepare_attentional_mechanism_input(Wh) # (N, N, 2 * out_features)
        e = self.leakyrelu(torch.matmul(a_input, self.a).squeeze(2))
        zero_vec = -9e15*torch.ones_like(e)
        # mask注意力系数
        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] # 节点数N
        Wh_repeated_in_chunks = Wh.repeat_interleave(N, dim=0) #(N*N, out_features)
        Wh_repeated_alternating = Wh.repeat(N, 1) #(N*N, out_features)
        all_combinations_matrix = torch.cat([Wh_repeated_in_chunks, Wh_repeated_alternating], dim=1)
        # all_combinations_matrix.shape == (N * N, 2 * out_features)
        return all_combinations_matrix.view(N, N, 2 * self.out_features)
    
    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

2.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)
        # 输出层
        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

3.模型训练

3.1参数

hidden = 8 
dropout = 0.6
nb_heads = 8 
alpha = 0.2
lr = 0.005
weight_decay = 5e-4
epochs = 10000
patience = 100
cuda = torch.cuda.is_available()
1
2
3
4
5
6
7
8
9

3.2 实例化模型和优化器

# 实例化模型 
model = GAT(nfeat=features.shape[1], 
            nhid=hidden, 
            nclass=int(labels.max()) + 1, 
            dropout=dropout, 
            nheads=nb_heads, 
            alpha=alpha)
# 优化器
optimizer = optim.Adam(model.parameters(), 
                       lr=lr, 
                       weight_decay=weight_decay)
if 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

3.3 训练

def train(epoch):
    t = time.time()
    # trian
    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()
        
    # eval
    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.item()),
          "accuracy= {:.4f}".format(acc_test.item()))
    
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
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

t_total = time.time()
loss_values = []
bad_counter = 0
best = epochs + 1
best_epoch = 0
for epoch in range(epochs):
    # 训练模型并保存loss
    loss_values.append(train(epoch))
    # 保存模型
    torch.save(model.state_dict(), '{}.pkl'.format(epoch))
    # 记录loss最小的epoch
    if loss_values[-1] < best:
        best = loss_values[-1]
        best_epoch = epoch
        bad_counter = 0
    else:
        bad_counter += 1
    
    # 如果连续patience个epoch，最小Loss都没有变则终止模型训练
    if bad_counter == 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))

# 加载最优模型
print('Loading {}th epoch'.format(best_epoch))
model.load_state_dict(torch.load('{}.pkl'.format(best_epoch)))

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

输出：

Epoch: 0001 loss_train: 1.9447 acc_train: 0.1643 loss_val: 1.9360 acc_val: 0.3867 time: 1.8232s
Epoch: 0002 loss_train: 1.9351 acc_train: 0.2500 loss_val: 1.9255 acc_val: 0.5100 time: 1.2726s
Epoch: 0003 loss_train: 1.9181 acc_train: 0.3286 loss_val: 1.9152 acc_val: 0.5200 time: 1.2716s
Epoch: 0004 loss_train: 1.9063 acc_train: 0.3857 loss_val: 1.9049 acc_val: 0.5033 time: 1.3065s
Epoch: 0005 loss_train: 1.8864 acc_train: 0.4929 loss_val: 1.8943 acc_val: 0.5000 time: 1.2736s
...
Epoch: 0755 loss_train: 0.5768 acc_train: 0.8357 loss_val: 0.6613 acc_val: 0.8133 time: 1.3564s
Epoch: 0756 loss_train: 0.5965 acc_train: 0.8143 loss_val: 0.6615 acc_val: 0.8133 time: 1.4003s
Epoch: 0757 loss_train: 0.5388 acc_train: 0.8357 loss_val: 0.6619 acc_val: 0.8133 time: 1.4023s
Optimization Finished!
Total time elapsed: 1054.6370s
Loading 656th epoch
Test set results: loss= 0.6534 accuracy= 0.8460
1
2
3
4
5
6
7
8
9
10
11
12
13