utils.py(与gcn源码一致)
import numpy as np
import scipy.sparse as sp
import torch
'''
先将所有由字符串表示的标签数组用set保存,set的重要特征就是元素没有重复,
因此表示成set后可以直接得到所有标签的总数,随后为每个标签分配一个编号,创建一个单位矩阵,
单位矩阵的每一行对应一个one-hot向量,也就是np.identity(len(classes))[i, :],
再将每个数据对应的标签表示成的one-hot向量,类型为numpy数组
'''
def encode_onehot(labels):
classes = set(labels) # set() 函数创建一个无序不重复元素集
classes_dict = {c: np.identity(len(classes))[i, :] for i, c in # identity创建方矩阵
enumerate(classes)} # 字典 key为label的值,value为矩阵的每一行
# enumerate函数用于将一个可遍历的数据对象组合为一个索引序列
labels_onehot = np.array(list(map(classes_dict.get, labels)), # get函数得到字典key对应的value
dtype=np.int32)
return labels_onehot
# map() 会根据提供的函数对指定序列做映射
# 第一个参数 function 以参数序列中的每一个元素调用 function 函数,返回包含每次 function 函数返回值的新列表
# map(lambda x: x ** 2, [1, 2, 3, 4, 5])
# output:[1, 4, 9, 16, 25]
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) # 储存为csr型稀疏矩阵
labels = encode_onehot(idx_features_labels[:, -1])
# content file的每一行的格式为 : <paper_id> <word_attributes>+ <class_label>
# 分别对应 0, 1:-1, -1
# feature为第二列到倒数第二列,labels为最后一列
# build graph
# cites file的每一行格式为: <cited paper ID> <citing paper ID>
# 根据前面的contents与这里的cites创建图,算出edges矩阵与adj 矩阵
idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
idx_map = {j: i for i, j in enumerate(idx)}
# 由于文件中节点并非是按顺序排列的,因此建立一个编号为0-(node_size-1)的哈希表idx_map,
# 哈希表中每一项为id: number,即节点id对应的编号为number
edges_unordered = np.genfromtxt("{}{}.cites".format(path, dataset),
dtype=np.int32)
# edges_unordered为直接从边表文件中直接读取的结果,是一个(edge_num, 2)的数组,每一行表示一条边两个端点的idx
edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), # flatten:降维,返回一维数组
dtype=np.int32).reshape(edges_unordered.shape)
# 边的edges_unordered中存储的是端点id,要将每一项的id换成编号。
# 在idx_map中以idx作为键查找得到对应节点的编号,reshape成与edges_unordered形状一样的数组
adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), # coo型稀疏矩阵
shape=(labels.shape[0], labels.shape[0]),
dtype=np.float32)
# 根据coo矩阵性质,这一段的作用就是,网络有多少条边,邻接矩阵就有多少个1,
# 所以先创建一个长度为edge_num的全1数组,每个1的填充位置就是一条边中两个端点的编号,
# 即edges[:, 0], edges[:, 1],矩阵的形状为(node_size, node_size)。
# build symmetric adjacency matrix 论文里A^=(D~)^0.5 A~ (D~)^0.5这个公式
adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)
# 对于无向图,邻接矩阵是对称的。上一步得到的adj是按有向图构建的,转换成无向图的邻接矩阵需要扩充成对称矩阵
features = normalize(features)
adj = normalize(adj + sp.eye(adj.shape[0])) # eye创建单位矩阵,第一个参数为行数,第二个为列数
# 对应公式A~=A+IN
# 分别构建训练集、验证集、测试集,并创建特征矩阵、标签向量和邻接矩阵的tensor,用来做模型的输入
idx_train = range(140)
idx_val = range(200, 500)
idx_test = range(500, 1500)
features = torch.FloatTensor(np.array(features.todense())) # tensor为pytorch常用的数据结构
labels = torch.LongTensor(np.where(labels)[1])
adj = sparse_mx_to_torch_sparse_tensor(adj) # 邻接矩阵转为tensor处理
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(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. # 如果某一行全为0,则r_inv算出来会等于无穷大,将这些行的r_inv置为0
r_mat_inv = sp.diags(r_inv) # 构建对角元素为r_inv的对角矩阵
mx = r_mat_inv.dot(mx)
# 用对角矩阵与原始矩阵的点积起到标准化的作用,原始矩阵中每一行元素都会与对应的r_inv相乘,最终相当于除以了sum
return mx
def accuracy(output, labels):
preds = output.max(1)[1].type_as(labels) # 使用type_as(tesnor)将张量转换为给定类型的张量。
correct = preds.eq(labels).double() # 记录等于preds的label eq:equal
correct = correct.sum()
return correct / len(labels)
def sparse_mx_to_torch_sparse_tensor(sparse_mx): # 把一个sparse matrix转为torch稀疏张量
"""
numpy中的ndarray转化成pytorch中的tensor : torch.from_numpy()
pytorch中的tensor转化成numpy中的ndarray : numpy()
"""
sparse_mx = sparse_mx.tocoo().astype(np.float32)
indices = torch.from_numpy(np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
# 不懂的可以去看看COO性稀疏矩阵的结构
values = torch.from_numpy(sparse_mx.data)
shape = torch.Size(sparse_mx.shape)
return torch.sparse.FloatTensor(indices, values, shape)
layer.py
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
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.zeros(size=(in_features, out_features))) #建立都是0的矩阵,大小为(输入维度,输出维度)
nn.init.xavier_uniform_(self.W.data, gain=1.414)#xavier初始化
self.a = nn.Parameter(torch.zeros(size=(2*out_features, 1)))#见下图
#print(self.a.shape) torch.Size([16, 1])
nn.init.xavier_uniform_(self.a.data, gain=1.414)
self.leakyrelu = nn.LeakyReLU(self.alpha)
这里的self.a,对应的是论文里的向量a,故其维度大小应该为(2*out_features, 1)
def forward(self, input, adj):
h = torch.mm(input, self.W)
#print(h.shape) torch.Size([2708, 8]) 8是label的个数
N = h.size()[0]
#print(N) 2708 nodes的个数
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)#见下图
#print(a_input.shape) torch.Size([2708, 2708, 16])
e = self.leakyrelu(torch.matmul(a_input, self.a).squeeze(2)) #即论文里的eij
#squeeze除去维数为1的维度
#[2708, 2708, 16]与[16, 1]相乘再除去维数为1的维度,故其维度为[2708,2708],与领接矩阵adj的维度一样
zero_vec = -9e15*torch.ones_like(e)
#维度大小与e相同,所有元素都是-9*10的15次方
attention = torch.where(adj > 0, e, zero_vec)
'''这里我们回想一下在utils.py里adj怎么建成的:两个节点有边,则为1,否则为0。
故adj的领接矩阵的大小为[2708,2708]。(不熟的自己去复习一下图结构中的领接矩阵)。
print(adj)这里我们看其中一个adj
tensor([[0.1667, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.5000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.2000, ..., 0.0000, 0.0000, 0.0000],
...,
[0.0000, 0.0000, 0.0000, ..., 0.2000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.2000, 0.0000],
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.2500]])
不是1而是小数是因为进行了归一化处理
故当adj>0,即两结点有边,则用gat构建的矩阵e,若adj=0,则另其为一个很大的负数,这么做的原因是进行softmax时,这些数就会接近于0了。
'''
attention = F.softmax(attention, dim=1)
#对应论文公式3,attention就是公式里的αij
'''print(attention)
tensor([[0.1661, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.5060, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.2014, ..., 0.0000, 0.0000, 0.0000],
...,
[0.0000, 0.0000, 0.0000, ..., 0.1969, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.1998, 0.0000],
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.2548]]'''
attention = F.dropout(attention, self.dropout, training=self.training)
h_prime = torch.matmul(attention, h)
if self.concat:
return F.elu(h_prime)
else:
return h_prime
def __repr__(self):
return self.__class__.__name__ + ' (' + str(self.in_features) + ' -> ' + str(self.out_features) + ')'
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)这句有点复杂,我们先做个小实验看看函数是什么意思:
由此我们知道,a_input是由Whi和Whj concat得到,对应论文里的Whi||Whj
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
self.concat = concat
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]
edge = adj.nonzero().t()
h = torch.mm(input, self.W)
# h: N x out
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
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
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) + ')'
models.py
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)
#multi-head 隐藏层到输出
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)
这里的torch.cat即公式(5)中的||
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)
train.py
与gcn源码一致,在此不做讲解,如需了解请参照文章开头的GCN源码解析连接。
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 = argparse.ArgumentParser()
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.')
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).')
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.')
parser.add_argument('--dropout', type=float, default=0.6, help='Dropout rate (1 - keep probability).')
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
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)
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)
def train(epoch):
t = time.time()
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()
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[0]),
"accuracy= {:.4f}".format(acc_test.data[0]))
# 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()