pyGAT(graph attention networks)图自注意网络pytorch版本代码

以下都是自己读的,有错误理解的地方请大家留言指出!

【毕竟学术无错 (_) 逃~】

看懂这个代码逻辑还是需要了解pytorch的经验,我就没有苟了好久/(ㄒoㄒ)/~~

train层

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
parser = argparse.ArgumentParser()
# 添加一系列默认的超参
# 辨识没有GPU的情况
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.')
# 默认epoch轮
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).')
# 8个隐藏层
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.')
# 抓包每次0.6
parser.add_argument('--dropout', type=float, default=0.6, help='Dropout rate (1 - keep probability).')
# 激活函数里面的alpha
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)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

# Load data,加载和预处理数据
adj, features, labels, idx_train, idx_val, idx_test = load_data()

# Model and optimizer
# 建立模型和优化,这里是adam优化。对于特征是否稀疏给与了两种模型
if args.sparse:
    # 调用init函数,初始化搭建好模型
    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:
    # 调用init函数,初始化搭建好模型
    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)

# 是否可以用显卡驱动
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()

# Variable()把向量转换为可变换的变量,即可以用来反向传播
features, adj, labels = Variable(features), Variable(adj), Variable(labels)


def train(epoch):
    # 记录训练开始时间
    t = time.time()
    # 训练模型
    model.train()
    # 优化
    optimizer.zero_grad()
    # Attention!在初始化结束后得到的model进行调用模型,此时调用的是forward函数,输入是零阶矩阵和特征编码
    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()
        output = model(features, adj)
    # 获得损失和准确性
    loss_val = F.nll_loss(output[idx_val], labels[idx_val])
    acc_val = accuracy(output[idx_val], labels[idx_val])
    # output the information
    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()))

# Train model
t_total = time.time()
loss_values = []
bad_counter = 0
# 记录每次的损失值
best = args.epochs + 1
# 记录下最好的一次epoch
best_epoch = 0
for epoch in range(args.epochs):
    # 记录每轮损失值
    loss_values.append(train(epoch))
    # 永久化存储model
    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
    # 连续100次没有优化直接停。
    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)))
torch.save(model, '{}.pth'.format(epoch))

# 测试
compute_test()

utils层 主要是数据处理

import numpy as np
import scipy.sparse as sp
import torch


def encode_onehot(labels):
    classes = set(labels)
    classes_dict = {c: np.identity(len(classes))[i, :] for i, c in enumerate(classes)}
    labels_onehot = np.array(list(map(classes_dict.get, labels)), dtype=np.int32)
    return labels_onehot


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

    # genfromtxt 从txt文件中读取数据。Cora数据是文献之间的引用网络。
    idx_features_labels = np.genfromtxt("{}{}.content".format(path, dataset), dtype=np.dtype(str))
    # read data,上面的数据是【id features label】,所以这里直接读的全是features
    features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)
    print(features)
    print("hi")
    # 对标志进行one0hot编码
    labels = encode_onehot(idx_features_labels[:, -1])
    
    # build graph 存图,首先是将数据第一列paper_id读取出来,然后最后一段是论文分类标志
    idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
    # 枚举所有的idx,将i,j组成map,对所有的文献id进行编号
    idx_map = {j: i for i, j in enumerate(idx)}
    # 从cite读取所有的无序边信息,该文件里存储的是全部的引用关系,前者是被引用,后者是引用
    edges_unordered = np.genfromtxt("{}{}.cites".format(path, dataset), dtype=np.int32)
    # 将cite文件里面未排序边的文献id对应为map里面的id,转换为在有序编号id之间建立的边的关系
    # map以前者的函数作用于后者的迭代数据
    edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), dtype=np.int32).reshape(edges_unordered.shape)
    
    # print(adj.multiply)的out为
    # <bound method spmatrix.multiply of <2708x2708 sparse matrix of type '<class 'numpy.float32'>'
    # with 5429 stored elements in COOrdinate format>>
    # 由信息可见它的含义是建立了2708*2708大小的矩阵但是只有5429个有效数值,故是稀疏矩阵。 
    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)
    print(features)
    # 特征标准化
    features = normalize_features(features)
    print("hello")
    print(features)
    adj = normalize_adj(adj + sp.eye(adj.shape[0]))
    
    # 把前140条数据作为训练集,200到500为验证数据集,500到1500作为测试集
    idx_train = range(140)
    idx_val = range(200, 500)
    idx_test = range(500, 1500)

    # 把矩阵和one-hot编码变为向量
    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)


  #  print(adj[1],features)
    # adj是邻接矩阵,features指的是把数值除以该文献为1的个数,比如0号有20个1每个1给0.05,labels是对7种分类的编码
    # 三种数据分类只是给了索引
    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"""
    # attentionn 这里其实不太好懂,它的本质意思其实是在读出的features里面统计一行的1的个数。
    print(mx)
    rowsum = np.array(mx.sum(1))
    print(rowsum)
    # 取数的倒数,变为0-1之间的数值
    r_inv = np.power(rowsum, -1).flatten()
    # 矩阵对角化并且将INF的数命为0
    r_inv[np.isinf(r_inv)] = 0.
    r_mat_inv = sp.diags(r_inv)
    # 迷惑行为 再将两个矩阵乘起来
    mx = r_mat_inv.dot(mx)
    return mx


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

if __name__ == '__main__':
    load_data()

models层

import torch
import torch.nn as nn
import torch.nn.functional as F
from layers import GraphAttentionLayer, SpGraphAttentionLayer


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的dropout,以便于避免过拟合
        x = F.dropout(x, self.dropout, training=self.training)
        # 对输入的值在每一层attention层中进行处理。把激活函数处理后的值进行拼接
        x = torch.cat([att(x, adj) for att in self.attentions], dim=1)
        # 还是定义在输出层需要的dropout
        x = F.dropout(x, self.dropout, training=self.training)
        # 把最后的cat完的特征值输入到输出层
        x = F.elu(self.out_att(x, adj))
        return F.log_softmax(x, dim=1)

# 集成pytorchnn.module,重写forward方法。
class SpGAT(nn.Module):
    def __init__(self, nfeat, nhid, nclass, dropout, alpha, nheads):
        """Sparse version of GAT."""
        super(SpGAT, self).__init__()
        self.dropout = dropout

        self.attentions = [SpGraphAttentionLayer(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 = SpGraphAttentionLayer(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)

layers层

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


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__()
        # dropout参数
        self.dropout = dropout
        # 输入特征
        self.in_features = in_features
        self.out_features = out_features
        self.alpha = alpha
        self.concat = concat
        
        # w和a就是paper里面的w和a
        self.W = nn.Parameter(torch.zeros(size=(in_features, out_features)))
        nn.init.xavier_uniform_(self.W.data, gain=1.414)
        self.a = nn.Parameter(torch.zeros(size=(2*out_features, 1)))
        nn.init.xavier_uniform_(self.a.data, gain=1.414)
        
        # 定义leakyrelu的参数α
        self.leakyrelu = nn.LeakyReLU(self.alpha)

    def forward(self, input, adj):
        # 输入参数内容:in表示输入的节点的特征向量,adj即一个邻接矩阵
        # 从邻接矩阵得到关注里系数,然后再关注系数的前提下获得h_prime
        
        # 对应公式1,将特征与矩阵w相乘
        h = torch.mm(input, self.W)
        # N最后其实也等于所有的节点数量
        N = h.size()[0]

        # 这里看起来有点费解
        a_input = torch.cat([h.repeat(1, N).view(N * N, -1), h.repeat(N, 1)], dim=1).view(N, -1, 2 * self.out_features)
        # 对应公式里面的e的求解,即求点与点之间的关注系数
        e = self.leakyrelu(torch.matmul(a_input, self.a).squeeze(2))

        # 表示如果邻接矩阵元素大于0时,则两个节点有连接,该位置的注意力系数保留,
        # 否则需要mask并置为非常小的值,原因是softmax的时候这个最小值会不考虑。
        zero_vec = -9e15*torch.ones_like(e)
        attention = torch.where(adj > 0, e, zero_vec)
        # 对应公式2
        attention = F.softmax(attention, dim=1)
        # 实现dropout
        attention = F.dropout(attention, self.dropout, training=self.training)
        # 张量的乘法
        h_prime = torch.matmul(attention, h)

        if self.concat:
            # 对应公式5的激活函数已经内部内容,此处需要先进行激活再cat
            return F.elu(h_prime)
        else:
            # 最后一层需要先进行平均再用激活函数
            return h_prime

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


class SpecialSpmmFunction(torch.autograd.Function):
    """Special function for only sparse region backpropataion layer."""
    @staticmethod
    def forward(ctx, indices, values, shape, b):
        assert indices.requires_grad == False
        a = torch.sparse_coo_tensor(indices, values, shape)
        ctx.save_for_backward(a, b)
        ctx.N = shape[0]
        return torch.matmul(a, b)

    @staticmethod
    def backward(ctx, grad_output):
        a, b = ctx.saved_tensors
        grad_values = grad_b = None
        if ctx.needs_input_grad[1]:
            grad_a_dense = grad_output.matmul(b.t())
            edge_idx = a._indices()[0, :] * ctx.N + a._indices()[1, :]
            grad_values = grad_a_dense.view(-1)[edge_idx]
        if ctx.needs_input_grad[3]:
            grad_b = a.t().matmul(grad_output)
        return None, grad_values, None, grad_b


class SpecialSpmm(nn.Module):
    def forward(self, indices, values, shape, b):
        return SpecialSpmmFunction.apply(indices, values, shape, b)

    
class SpGraphAttentionLayer(nn.Module):
    """
    Sparse version GAT layer, similar to https://arxiv.org/abs/1710.10903
    """

    def __init__(self, in_features, out_features, dropout, alpha, concat=True):
        super(SpGraphAttentionLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.alpha = alpha
        # concat对应的论文的5式
        self.concat = concat

        # init.xavier_normal_按照正态正态分布去填充向量里的数据
        # 这里的w和a就是paper里面的w和a
        self.W = nn.Parameter(torch.zeros(size=(in_features, out_features)))
        nn.init.xavier_normal_(self.W.data, gain=1.414)
        self.a = nn.Parameter(torch.zeros(size=(1, 2*out_features)))
        nn.init.xavier_normal_(self.a.data, gain=1.414)

        self.dropout = nn.Dropout(dropout)
        self.leakyrelu = nn.LeakyReLU(self.alpha)
        self.special_spmm = SpecialSpmm()

    def forward(self, input, adj):
        dv = 'cuda' if input.is_cuda else 'cpu'
        # 取出输入的行数
        N = input.size()[0]
        # nonzero函数可以获得numpy数组里面的非零元素的索引值
        edge = adj.nonzero().t()
        
        # h: N x out,即把输入和获得的W相乘
        h = torch.mm(input, self.W)
        
        # 断言h里面不能有任何一个无穷值
        assert not torch.isnan(h).any()

        # Self-attention on the nodes - Shared attention mechanism
        edge_h = torch.cat((h[edge[0, :], :], h[edge[1, :], :]), dim=1).t()
        # edge: 2*D x E

        # leakyrelu对应于公式3,leakyrelu(a * h)
        # squeeze可以把维度为1的维去掉
        edge_e = torch.exp(-self.leakyrelu(self.a.mm(edge_h).squeeze()))
        assert not torch.isnan(edge_e).any()
        # edge_e: E

        e_rowsum = self.special_spmm(edge, edge_e, torch.Size([N, N]), torch.ones(size=(N,1), device=dv))
        # e_rowsum: N x 1

        # 对边进行dropout
        edge_e = self.dropout(edge_e)
        # edge_e: E


        h_prime = self.special_spmm(edge, edge_e, torch.Size([N, N]), h)
        assert not torch.isnan(h_prime).any()
        # h_prime: N x out
        
        h_prime = h_prime.div(e_rowsum)
        # h_prime: N x out
        assert not torch.isnan(h_prime).any()

        if self.concat:
            # if this layer is not last layer,
            return F.elu(h_prime)
        else:
            # if this layer is last layer,
            return h_prime

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

以下是使用PyTorch实现GAT的代码示例: ``` python import torch import torch.nn as nn import torch.nn.functional as F class GATLayer(nn.Module): def __init__(self, in_dim, out_dim): super(GATLayer, self).__init__() self.in_dim = in_dim self.out_dim = out_dim self.W = nn.Parameter(torch.zeros(size=(in_dim, out_dim))) self.a = nn.Parameter(torch.zeros(size=(2*out_dim, 1))) nn.init.xavier_uniform_(self.W.data, gain=1.414) nn.init.xavier_uniform_(self.a.data, gain=1.414) def forward(self, h, adj): Wh = torch.mm(h, self.W) a_input = self._prepare_attentional_mechanism_input(Wh) e = F.leaky_relu(torch.matmul(a_input, self.a).squeeze(2)) zero_vec = -9e15*torch.ones_like(e) attention = torch.where(adj > 0, e, zero_vec) attention = F.softmax(attention, dim=1) h_prime = torch.matmul(attention, Wh) return h_prime def _prepare_attentional_mechanism_input(self, Wh): N = Wh.size()[0] Wh_repeated_in_chunks = Wh.repeat_interleave(N, dim=0) Wh_repeated_alternating = Wh.repeat(N, 1) all_combinations_matrix = torch.cat([Wh_repeated_in_chunks, Wh_repeated_alternating], dim=1) return all_combinations_matrix.view(N, N, 2*self.out_dim) class GAT(nn.Module): def __init__(self, n_feat, n_hid, n_class, dropout, alpha, n_heads): super(GAT, self).__init__() self.dropout = dropout self.attentions = [GATLayer(n_feat, n_hid) for _ in range(n_heads)] for i, attention in enumerate(self.attentions): self.add_module('attention_{}'.format(i), attention) self.out_att = GATLayer(n_hid*n_heads, n_class) self.alpha = alpha 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) ``` 在此示例中,我们实现了一个包含多头注意力机制的GAT模型。其中,GATLayer是GAT的核心组件,每个GATLayer都包含一个注意力头。在GAT模型中,我们将多个注意力头的输出连接在一起,再通过一个输出层进行分类。在forward函数中,我们首先对输入进行dropout,然后通过多个GATLayer进行特征提取,最后通过输出层进行分类并使用log_softmax进行预测。
评论 6
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值