GAT学习：PyG实现multi-head GAT（二）

接上篇学习笔记GAT学习：PyG实现GAT（图注意力神经网络）网络（一）为了使得Attention的效果更好，所以加入multi-head attention。画个图说明multi-head attention的工作原理。

其实就相当于并联了head_num个attention后，将每个attention层的输出特征拼接起来，然后再输入一个attenion层得到输出结果。

预备知识

关于GAT的原理等知识，参考我的上篇博客：PyG实现GAT（图注意力神经网络）网络（一）

代码分析

import torch
import math
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops,degree
from torch_geometric.datasets import Planetoid
import ssl
import torch.nn.functional as F


class GAL(MessagePassing):
    def __init__(self,in_features,out_featrues):
        super(GAL,self).__init__(aggr='add')
        self.a = torch.nn.Parameter(torch.zeros(size=(2*out_featrues, 1)))
        torch.nn.init.xavier_uniform_(self.a.data, gain=1.414)  # 初始化
        # 定义leakyrelu激活函数
        self.leakyrelu = torch.nn.LeakyReLU()
        self.linear=torch.nn.Linear(in_features,out_featrues)
    def forward(self,x,edge_index):
        x=self.linear(x)
        N=x.size()[0]
        row，col=edge_index
        a_input = torch.cat([x[row], x[col]], dim=1)
        # [N, N, 1] => [N, N] 图注意力的相关系数（未归一化）
        temp=torch.mm(a_input,self.a).squeeze()
        e = self.leakyrelu(temp)
        #e_all为同一个节点与其全部邻居的计算的分数的和，用于计算归一化softmax
        e_all=torch.zeros(x.size()[0])
        count = 0
        for i in col:
            e_all[i]+=e[count]
            count=count+1

        for i in range(len(e)):
            e[i]=math.exp(e[i])/math.exp(e_all[col[i]])

        return self.propagate(edge_index,x=x,norm=e)

    def message(self, x_j, norm):
        return norm.view(-1, 1) * x_j


class GAT(torch.nn.Module):
    def __init__(self, in_features, hid_features, out_features, n_heads):
        """
        n_heads 表示有几个GAL层，最后进行拼接在一起，类似self-attention
        从不同的子空间进行抽取特征。
        """
        super(GAT, self).__init__()

        # 定义multi-head的图注意力层
        self.attentions = [GAL(in_features, hid_features) for _ in
                           range(n_heads)]
        # 输出层，也通过图注意力层来实现，可实现分类、预测等功能
        self.out_att = GAL(hid_features * n_heads, out_features)

    def forward(self, x, edge_index):
        # 将每个head得到的x特征进行拼接
        x = torch.cat([att(x, edge_index) for att in self.attentions], dim=1)
        print('x.size after cat',x.size())
        x = F.elu(self.out_att(x,edge_index))  # 输出并激活
        print('x.size after elu',x.size())
        return F.log_softmax(x, dim=1)  # log_softmax速度变快，保持数值稳定



class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.gat = GAT(dataset.num_node_features,16,7,4)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = F.dropout(x, training=self.training)
        x = self.gat(x, edge_index)
        print('X_GAT',x.size())
        return F.log_softmax(x, dim=1)



ssl._create_default_https_context = ssl._create_unverified_context
dataset = Planetoid(root='Cora', name='Cora')
x=dataset[0].x
edge_index=dataset[0].edge_index
model=Net()
data=dataset[0]
out=Net()(data)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
model.train()
for epoch in range(2):
    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('Accuracy:{:.4f}'.format(acc))
>>>Accuracy:0.1930

GAT

class GAT(torch.nn.Module):
    def __init__(self, in_features, hid_features, out_features, n_heads):
        """
        n_heads 表示有几个GAL层，最后进行拼接在一起，类似self-attention
        从不同的子空间进行抽取特征。
        """
        super(GAT, self).__init__()

        # 定义multi-head的图注意力层
        self.attentions = [GAL(in_features, hid_features) for _ in
                           range(n_heads)]
        # 输出层，也通过图注意力层来实现，可实现分类、预测等功能
        self.out_att = GAL(hid_features * n_heads, out_features)

    def forward(self, x, edge_index):
        # 将每个head得到的x特征进行拼接
        x = torch.cat([att(x, edge_index) for att in self.attentions], dim=1)
        print('x.size after cat',x.size())
        x = F.elu(self.out_att(x,edge_index))  # 输出并激活
        print('x.size after elu',x.size())
        return F.log_softmax(x, dim=1)  # log_softmax速度变快，保持数值稳定
>>>x.size after cat torch.Size([2708, 64])
x.size after elu torch.Size([2708, 7])
x.size after cat torch.Size([2708, 64])
x.size after elu torch.Size([2708, 7])
x.size after cat torch.Size([2708, 64])
x.size after elu torch.Size([2708, 7])
x.size after cat torch.Size([2708, 64])
x.size after elu torch.Size([2708, 7])