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

本文主要是介绍GAT学习：PyG实现multi-head GAT（二），希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

PyG实现GAT网络

预备知识
代码分析
- 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 Fclass 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_indexa_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为同一个节点与其全部邻居的计算的分数的和，用于计算归一化softmaxe_all=torch.zeros(x.size()[0])count = 0for i in col:e_all[i]+=e[count]count=count+1for 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_jclass 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 _ inrange(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_indexx = 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 _ inrange(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])