pytorch|图卷积神经网络(GCN)与图注意力(GAT)在cora数据集的应用
本文实现的两层GCN与GAT在cora数据集上效果，并可视化展示。
1.导入的相关包
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pandas as pd
import torch
from torch._C import parse_ir
import torch_geometric
from torch_geometric import datasets
from torch_geometric.data import Data, InMemoryDataset
from torch_geometric.datasets import Planetoid
import torch.nn.functional as F
from torch_geometric.nn import GCNConv,GATConv # GCN
from sklearn.manifold import TSNE
2.数据处理部分
path = "E:\\download\\cora\\cora\\"
cites = path + "cora.cites"
content = path + "cora.content"
# 索引字典，将原本的论文id转换到从0开始编码
index_dict = dict()
# 标签字典，将字符串标签转化为数值
label_to_index = dict()
features = []
labels = []
edge_index = []
with open(content,"r") as f:
nodes = f.readlines()
for node in nodes:
node_info = node.split()
index_dict[int(node_info[0])] = len(index_dict)
features.append([int(i) for i in node_info[1:-1]])
label_str = node_info[-1]
if(label_str not in label_to_index.keys()):
label_to_index[label_str] = len(label_to_index)
labels.append(label_to_index[label_str])
with open(cites,"r") as f:
edges = f.readlines()
for edge in edges:
start, end = edge.split()
# 训练时将边视为无向的，但原本的边是有向的，因此需要正反添加两次
edge_index.append([index_dict[int(start)],index_dict[int(end)]])
edge_index.append([index_dict[int(end)],index_dict[int(start)]])
# 转换为Tensor
labels = torch.LongTensor(labels)
features = torch.FloatTensor(features)
# 行归一化
features = torch.nn.functional.normalize(features, p=1, dim=1)
edge_index = torch.LongTensor(edge_index)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # 本电脑只有一个GPU
mask = torch.randperm(len(index_dict)) # 随机打乱顺序
train_mask = mask[:140]
val_mask = mask[140:640]
test_mask = mask[1708:2708]
cora = Data(x = features, edge_index = edge_index.t().contiguous(), y = labels).to(device)
3.GCN模块
class GCNNet(torch.nn.Module):
def __init__(self, num_feature, num_label):
super(GCNNet,self).__init__()
self.GCN1 = GCNConv(num_feature, 16)
self.GCN2 = GCNConv(16, num_label)
self.dropout = torch.nn.Dropout(p=0.5)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.GCN1(x, edge_index)
x = F.relu(x)
x = self.dropout(x)
x = self.GCN2(x, edge_index)
return F.log_softmax(x, dim=1)
4.GAT模块
class GATNet(torch.nn.Module):
def __init__(self, num_feature, num_label):
super(GATNet,self).__init__()
self.GAT1 = GATConv(num_feature, 8, heads = 8, concat = True, dropout = 0.6)
self.GAT2 = GATConv(8*8, num_label, dropout = 0.6)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.GAT1(x, edge_index)
x = F.relu(x)
x = self.GAT2(x, edge_index)
return F.log_softmax(x, dim=1)
5.画图及显示模块
def gcn_apply():
model = GCNNet(features.shape[1], len(label_to_index)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
acc_record = []
loss_record = []
for epoch in range(500):
optimizer.zero_grad()
out = model(cora)
loss = F.nll_loss(out[train_mask], cora.y[train_mask])
loss_record.append(loss.item())
print('epoch: %d loss: %.4f' %(epoch, loss))
loss.backward()
optimizer.step()
# if((epoch + 1)% 10 == 0):
model.eval()
_, pred = model(cora).max(dim=1)
correct = int(pred[test_mask].eq(cora.y[test_mask]).sum().item())
acc = correct / len(test_mask)
acc_record.append(acc)
print('Accuracy: {:.4f}'.format(acc))
model.train()
ts = TSNE(n_components=2)
ts.fit_transform(out[test_mask].to('cpu').detach().numpy())
x = ts.embedding_
y = cora.y[test_mask].to('cpu').detach().numpy()
xi = []
for i in range(7):
xi.append(x[np.where(y==i)])
colors = ['mediumblue','green','red','yellow','cyan','mediumvioletred','mediumspringgreen']
plt.figure(figsize=(8, 6))
for i in range(7):
plt.scatter(xi[i][:,0],xi[i][:,1],s=30,color=colors[i],marker='+',alpha=1)
plt.show()
plt.figure(figsize=(14, 9))
x = range(500)
plt.plot(x,acc_record,label='acc')
plt.plot(x,loss_record,label='loss')
plt.legend()
plt.grid()
plt.show()
def gat_apply():
seed = 1
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed) # Numpy module.
# random.seed(seed) # Python random module.
torch.manual_seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
model = GATNet(features.shape[1], len(label_to_index)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
acc_record = []
loss_record = []
for epoch in range(500):
optimizer.zero_grad()
out = model(cora)
loss = F.nll_loss(out[train_mask], cora.y[train_mask])
loss_record.append(loss.item())
print('epoch: %d loss: %.4f' %(epoch, loss))
loss.backward()
optimizer.step()
# if((epoch + 1)% 10 == 0):
model.eval()
_, pred = model(cora).max(dim=1)
correct = int(pred[test_mask].eq(cora.y[test_mask]).sum().item())
acc = correct / len(test_mask)
acc_record.append(acc)
print('Accuracy: {:.4f}'.format(acc))
model.train()
ts = TSNE(n_components=2)
ts.fit_transform(out[test_mask].to('cpu').detach().numpy())
x = ts.embedding_
y = cora.y[test_mask].to('cpu').detach().numpy()
xi = []
for i in range(7):
xi.append(x[np.where(y==i)])
colors = ['mediumblue','green','red','yellow','cyan','mediumvioletred','mediumspringgreen']
plt.figure(figsize=(8, 6))
for i in range(7):
plt.scatter(xi[i][:,0],xi[i][:,1],s=30,color=colors[i],marker='+',alpha=1)
plt.show()
plt.figure(figsize=(14, 9))
x = range(500)
plt.plot(x,acc_record,label='acc')
plt.plot(x,loss_record,label='loss')
plt.legend()
plt.grid()
plt.show()
6.主函数
if __name__ == "__main__":
gcn_apply()
gat_apply()
7.整个网络的代码
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pandas as pd
import torch
from torch._C import parse_ir
import torch_geometric
from torch_geometric import datasets
from torch_geometric.data import Data, InMemoryDataset
from torch_geometric.datasets import Planetoid
import torch.nn.functional as F
from torch_geometric.nn import GCNConv,GATConv # GCN
from sklearn.manifold import TSNE
path = "E:\\download\\cora\\cora\\"
cites = path + "cora.cites"
content = path + "cora.content"
# 索引字典，将原本的论文id转换到从0开始编码
index_dict = dict()
# 标签字典，将字符串标签转化为数值
label_to_index = dict()
features = []
labels = []
edge_index = []
with open(content,"r") as f:
nodes = f.readlines()
for node in nodes:
node_info = node.split()
index_dict[int(node_info[0])] = len(index_dict)
features.append([int(i) for i in node_info[1:-1]])
label_str = node_info[-1]
if(label_str not in label_to_index.keys()):
label_to_index[label_str] = len(label_to_index)
labels.append(label_to_index[label_str])
with open(cites,"r") as f:
edges = f.readlines()
for edge in edges:
start, end = edge.split()
# 训练时将边视为无向的，但原本的边是有向的，因此需要正反添加两次
edge_index.append([index_dict[int(start)],index_dict[int(end)]])
edge_index.append([index_dict[int(end)],index_dict[int(start)]])
class GCNNet(torch.nn.Module):
def __init__(self, num_feature, num_label):
super(GCNNet,self).__init__()
self.GCN1 = GCNConv(num_feature, 16)
self.GCN2 = GCNConv(16, num_label)
self.dropout = torch.nn.Dropout(p=0.5)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.GCN1(x, edge_index)
x = F.relu(x)
x = self.dropout(x)
x = self.GCN2(x, edge_index)
return F.log_softmax(x, dim=1)
class GATNet(torch.nn.Module):
def __init__(self, num_feature, num_label):
super(GATNet,self).__init__()
self.GAT1 = GATConv(num_feature, 8, heads = 8, concat = True, dropout = 0.6)
self.GAT2 = GATConv(8*8, num_label, dropout = 0.6)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.GAT1(x, edge_index)
x = F.relu(x)
x = self.GAT2(x, edge_index)
return F.log_softmax(x, dim=1)
# 为每个节点增加自环，但后续GCN层默认会添加自环，跳过即可
# for i in range(2708):
# edge_index.append([i,i])
# 转换为Tensor
labels = torch.LongTensor(labels)
features = torch.FloatTensor(features)
# 行归一化
features = torch.nn.functional.normalize(features, p=1, dim=1)
edge_index = torch.LongTensor(edge_index)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # 本电脑只有一个GPU
mask = torch.randperm(len(index_dict)) # 随机打乱顺序
train_mask = mask[:140]
val_mask = mask[140:640]
test_mask = mask[1708:2708]
cora = Data(x = features, edge_index = edge_index.t().contiguous(), y = labels).to(device)
def gcn_apply():
model = GCNNet(features.shape[1], len(label_to_index)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
acc_record = []
loss_record = []
for epoch in range(500):
optimizer.zero_grad()
out = model(cora)
loss = F.nll_loss(out[train_mask], cora.y[train_mask])
loss_record.append(loss.item())
print('epoch: %d loss: %.4f' %(epoch, loss))
loss.backward()
optimizer.step()
# if((epoch + 1)% 10 == 0):
model.eval()
_, pred = model(cora).max(dim=1)
correct = int(pred[test_mask].eq(cora.y[test_mask]).sum().item())
acc = correct / len(test_mask)
acc_record.append(acc)
print('Accuracy: {:.4f}'.format(acc))
model.train()
ts = TSNE(n_components=2)
ts.fit_transform(out[test_mask].to('cpu').detach().numpy())
x = ts.embedding_
y = cora.y[test_mask].to('cpu').detach().numpy()
xi = []
for i in range(7):
xi.append(x[np.where(y==i)])
colors = ['mediumblue','green','red','yellow','cyan','mediumvioletred','mediumspringgreen']
plt.figure(figsize=(8, 6))
for i in range(7):
plt.scatter(xi[i][:,0],xi[i][:,1],s=30,color=colors[i],marker='+',alpha=1)
plt.show()
plt.figure(figsize=(14, 9))
x = range(500)
plt.plot(x,acc_record,label='acc')
plt.plot(x,loss_record,label='loss')
plt.legend()
plt.grid()
plt.show()
def gat_apply():
seed = 1
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed) # Numpy module.
# random.seed(seed) # Python random module.
torch.manual_seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
model = GATNet(features.shape[1], len(label_to_index)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
acc_record = []
loss_record = []
for epoch in range(500):
optimizer.zero_grad()
out = model(cora)
loss = F.nll_loss(out[train_mask], cora.y[train_mask])
loss_record.append(loss.item())
print('epoch: %d loss: %.4f' %(epoch, loss))
loss.backward()
optimizer.step()
# if((epoch + 1)% 10 == 0):
model.eval()
_, pred = model(cora).max(dim=1)
correct = int(pred[test_mask].eq(cora.y[test_mask]).sum().item())
acc = correct / len(test_mask)
acc_record.append(acc)
print('Accuracy: {:.4f}'.format(acc))
model.train()
ts = TSNE(n_components=2)
ts.fit_transform(out[test_mask].to('cpu').detach().numpy())
x = ts.embedding_
y = cora.y[test_mask].to('cpu').detach().numpy()
xi = []
for i in range(7):
xi.append(x[np.where(y==i)])
colors = ['mediumblue','green','red','yellow','cyan','mediumvioletred','mediumspringgreen']
plt.figure(figsize=(8, 6))
for i in range(7):
plt.scatter(xi[i][:,0],xi[i][:,1],s=30,color=colors[i],marker='+',alpha=1)
plt.show()
plt.figure(figsize=(14, 9))
x = range(500)
plt.plot(x,acc_record,label='acc')
plt.plot(x,loss_record,label='loss')
plt.legend()
plt.grid()
plt.show()
if __name__ == "__main__":
gcn_apply()
gat_apply()
8.相关图片
图片(1)程序运行截图
图片(2)GCN分类图
图片(3)GCN LOSS与ACC
图片(4)GAT分类可视化
图片(5)GAT LOSS与ACC
