图神经网络（2）：基于DGL实现GCN算法

通过DGL框架，自定义图卷积层（GCN）。需要掌握“如何使用DGL”实现图卷积算法。

（一）加载依赖包

import numpy as np
import torch
import torch.nn as nn
import dgl

（二）定义消息传递、消息聚合、更新函数

通过DGL框架的自定义方式，定义消息传递函数gcn_msg。其中，edge.src[‘h’] 表示source节点的特征； edge.src['norm'] 表示source节点的 ${度}^{-0.5}$ 。

def gcn_msg(edge):
    msg = edge.src['h'] * edge.src['norm']
    return {'m': msg}

通过DGL框架的自定义方式，定义消息聚合函数gcn_reduce。node.data['norm'] 表示source节点的 ${度}^{-0.5}$ 。

def gcn_reduce(node):
    accum = torch.sum(node.mailbox['m'], 1) * node.data['norm']
    return {'h': accum}

更新函数定义如下。我们通过使用一个全连接层和一个Tanh激活函数实现，这里激活函数可以替换。

更新函数的主要目的是为了改变特征维度。

class NodeApplyModule(nn.Module):
    def __init__(self, in_feats, out_feats):
    	"""
    	in_feats: 特征输入的维度
    	out_feats：特征输出的维度
		"""
        super(NodeApplyModule, self).__init__()
        self.linear = nn.Linear(in_features=in_feats, out_features=out_feats)
        self.activation = nn.Tanh()
    
    def forward(self, node):
    	"""
    	node：图的节点
    	返回字典形式
    	"""
        feats = node.data['h']
        h = self.linear(feats)
        h = self.activation(h)
        return {'h': h}

（三）定义GCNLayer

在定义好消息传递、消息聚合以及更新函数后，我们进一步构建GCNLayer层。

class GCNLayer(nn.Module):
    def __init__(self, in_feats, out_feats):
    	"""
    	in_feats: 特征输入的维度
    	out_feats：特征输出的维度
    	"""
        super(GCNLayer, self).__init__()
        self.apply_node = NodeApplyModule(in_feats, out_feats)  # 引入更新函数
        
    def forward(self, g, features):
    	"""
    	g: 通过DGL定义的graph
    	features：图的节点特征
    	"""
        g.ndata['h'] = features  # 将节点特征，传入图的属性
        g.update_all(gcn_msg, gcn_reduce)  # 消息传递和聚合
        g.apply_nodes(func=self.apply_node)  # 更新
        h = g.ndata.pop('h')  # 提取计算后的节点特征，并删除Graph中的对应属性
        return h

（四）预处理数据集

选取Cora数据集，进行实验。通过以下代码进行下载、预处理数据。

from dgl.data import CoraGraphDataset
data = CoraGraphDataset()  # 采用Cora数据进行实验
g = data[0]  # 只用一个graph

features = g.ndata['feat']
labels = g.ndata['label']
train_mask = g.ndata['train_mask']
val_mask = g.ndata['val_mask']
test_mask = g.ndata['test_mask']
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()

g = dgl.remove_self_loop(g)
g = dgl.add_self_loop(g)
n_edges = g.number_of_edges()

# normalization
degs = g.in_degrees().float()  # 获取节点的度
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
g.ndata['norm'] = norm.unsqueeze(1)

（五）构建模型

这里，我们只构建两层GCN的模型。

class Net(nn.Module):
    def __init__(self, in_feats, out_feats):
        super(Net, self).__init__()
        self.gcn1 = GCNLayer(in_feats=in_feats, out_feats=16)
        self.gcn2 = GCNLayer(in_feats=16, out_feats=out_feats)
    
    def forward(self, g, features):
        h = features
        h = self.gcn1(g, h)
        h = self.gcn2(g, h)
        return h

训练模型：

# 实例化模型
model = Net(in_feats=1433, out_feats=7)

# loss
loss_fcn = torch.nn.CrossEntropyLoss()

# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
                             lr=0.001
                             )
# 训练模型
for epoch in range(50):
    model.train()
    # forward
    logits = model(g, features)
    loss = loss_fcn(logits[train_mask], labels[train_mask])
    
    # backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print('Epoch: {:5d} | Loss: {:.3f}'.format(epoch, loss.item()))

输出结果如下：

Epoch: 0 | Loss: 1.950
Epoch: 1 | Loss: 1.950
Epoch: 2 | Loss: 1.949
。。。
Epoch: 46 | Loss: 1.903
Epoch: 47 | Loss: 1.902
Epoch: 48 | Loss: 1.900
Epoch: 49 | Loss: 1.899