今天周六,不是很想看论文,于是自己动手把几种GNN手撸一遍,GCN的代码如下:
1. 得到数据
这里我使用的是经典的Cora数据集,用于进行节点分类,我是直接使用的PyG中提供的数据集,节点特征、节点标签、训练集测试集都可以直接调用:
from torch_geometric.datasets import Planetoid
import torch
# 使用PyG提供的数据集,懒得自己写解析规则
# 这一步如果下载不下来亲测开个VPN就行了
dataset = Planetoid(root='../data/Cora', name='Cora')
data = dataset[0]
# 得到边的index,shape:[2,E]
edge_index=data.edge_index
# 这几步是为邻接矩阵加上自连接,并且进行行归一化
adj_coo=torch.sparse_coo_tensor(edge_index,torch.ones(edge_index.shape[1]),size=(data.num_nodes,data.num_nodes))
adj=adj_coo.to_dense()
D=torch.diag(torch.pow(adj.sum(dim=1),-1))
adj=torch.mm(D,adj)
# 将邻接矩阵表示为torch sparse矩阵,加快计算,减少内存
adj=adj.to_sparse()
2.定义GCN中的一层
import torch
import torch.nn as nn
import torch.nn.functional as F
# 定义GCN中的一层
class GCNLayer(nn.Module):
def __init__(self, in_features, out_features):
super(GCNLayer, self).__init__()
self.W = nn.Parameter(torch.zeros([in_features, out_features]))
self.init_parameters()
def init_parameters(self):
self.W.data.uniform_(0, 1)
def forward(self,H:torch.Tensor,adj:torch.Tensor):
# H^(l+1)=AH^(l)W^(l)
H=torch.mm(H,self.W)
H=torch.mm(adj,H)
return F.relu(H)
3. 定义两层GCN模型
import torch.nn as nn
import torch.nn.functional as F
from layers import GCNLayer
# 定义一个两层的GCN
class GCN_NET(nn.Module):
def __init__(self, num_input, num_hidden, num_classes, dropout=0.5):
super(GCN_NET, self).__init__()
self.gcn1 = GCNLayer(num_input, num_hidden)
self.gcn2 = GCNLayer(num_hidden, num_classes)
self.dropout = dropout
def forward(self, X, adj):
H = self.gcn1(X, adj)
H = F.dropout(H, self.dropout,training=self.training)
H = self.gcn2(H, adj)
return F.log_softmax(H, dim=1)
4.训练模型
from data import adj, data, dataset
from model import GCN_NET
from torch.optim import Adam
import torch.nn.functional as F
# 初始特征,节点标签
features, labels = data.x, data.y
# 初始化模型
model = GCN_NET(data.num_features, 16, dataset.num_classes,dropout=0.2)
# 定义优化器
optimizer = Adam(model.parameters(),lr=1e-2,weight_decay=5e-4)
# 最大迭代次数
epoches = 200
# 将self.training设置为True,可以dropout
model.train()
for epoch in range(epoches):
predict_label = model(features, adj)
loss = F.cross_entropy(predict_label[data.train_mask], labels[data.train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 将self.training设置为false,不允许dropout
model.eval()
predict_label = model(features, adj).max(dim=1).indices
acc_num = predict_label[data.test_mask].eq(labels[data.test_mask]).sum()
print('accuracy:', acc_num / data.test_mask.sum())
训练完成,精确度0.803