手撸GCN

今天周六，不是很想看论文，于是自己动手把几种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