DGL节点分类简单实现

说明：这个系列来自于DGL上面的A Blitz Introduction to DGL。如果看英文习惯的小伙伴还是建议直接看官网文档。

Node Classification with DGL

本节结束您能够完成以下任务

• 加载DGL提供的数据集
• 通过DGL提供的模块构建GNN模型
• 能够在CPU和GPU上，训练或评估用于节点分类任务的GNN模块

说明：本节任务假定你已经具备了pytorch的一些先验知识。

开心学习每一天 ~~~

导入相关包（备注：为方便理解，这里导入了networkx和plt两个包，以便后面进行可视化，加深理解。）

import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F
import networkx as nx
import matplotlib.pyplot as plt

GNN中节点分类任务之概览

节点分类是GNN中应用最广泛同时也是最流行的任务，在节点分类任务中，GNN模型需要能够准确的预测出每一个节点的所属类别。在GNN流行之前也有一些技术能够做节点分类任务，类似于DeepWalk、Node2vec。GNN是通过聚合邻居特征表示来进行节点分类的。

加载数据集（这里使用的是Cora数据集，这是一个paper引用数据集）

import dgl.data
dataset=dgl.data.CoraGraphDataset()

取得数据集中的图（Cora数据集中只有一个图，如果一个数据集中含有多个图，则是按照列表的形式按照索引拿出不同的图）

g=dataset[0]

DGL通过ndata/edata分别用来存储节点特征和边特征

本示例中Cora数据集一共有四个节点特征

train_mask：用来标志当前节点是否属于训练集，bool类型

val_mask：用来表示当前节点是否属于验证集，bool类型

test_mask：用来表示当前节点是否属于测试集，bool类型

label：当前节点的准确类别

feat：节点特征

print("Node features")
print(g.ndata)
print('Edge features')
print(g.edata)

定义GCN；本次节点分类任务使用一个两层的GCN来聚合邻居信息，当然了你也可以很轻松的使用其他的领域聚类手段替换掉GCN。

# 定义一个两层的GCN，每一层都会生成新的节点表示通过汇聚其邻居节点
from dgl.nn import GraphConv
class GCN(nn.Module):
    def __init__(self,in_feats,h_feats,num_classes):
        '''
        in_feats:输入节点特征
        h_feats：隐藏特征维度
        num_classes：最后的分类数量
        '''
        super(GCN,self).__init__()
        self.conv1=GraphConv(in_feats,h_feats)
        self.conv2=GraphConv(h_feats,num_classes)
    def forward(self,g,in_feats):
        h=self.conv1(g,in_feats)
        h=F.relu(h)
        h=self.conv2(g,h)
        return h

实例化一个模型

model=GCN(g.ndata['feat'].shape[1],16,dataset.num_classes)
#g.ndata['feat'].shape[1]   节点特征维度
#16    隐藏层维度
#dataset.num_classes   实际种类数量，即最后输出层维度

训练模型（就和训练一个普通的NN模型一样的，如果你有pytorch的基本知识，看起来应该很轻松的吧）

def train(g,model):
    optimizer=torch.optim.Adam(model.parameters(),lr=0.01)   #优化器
    best_val_acc=0     #验证ACC，ACC可以简单理解是一种模型评价指标就行，这玩意越高越好
    best_test_acc=0
    
    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']     #测试数据集
    for e in range(100):
        logits=model(g,features)
        pred=logits.argmax(1)
        loss=F.cross_entropy(logits[train_mask],labels[train_mask])
        train_acc=(pred[train_mask]==labels[train_mask]).float().mean()
        val_acc=(pred[val_mask]==labels[val_mask]).float().mean()
        test_acc=(pred[test_mask]==labels[test_mask]).float().mean()
        if best_val_acc<val_acc:
            best_val_acc=val_acc
            best_test_acc=test_acc
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if e%5==0:
            print('In epoch {}, loss: {:.3f}, val acc: {:.3f} (best {:.3f}), test acc: {:.3f} (best {:.3f})'.format(
                e, loss, val_acc, best_val_acc, test_acc, best_test_acc))
model=GCN(g.ndata['feat'].shape[1],16,dataset.num_classes)
train(g,model)

在GPU上面训练（GPU和CPU上面的训练基本上都是一样的，只是小改一点点就可以了）

g=g.to('cuda')   #还记得g是什么吧
model=GCN(g.ndata['feat'].shape[1],16,dataset.num_classes).to('cuda')
train(g,model)    

nice~ 至此，你已经学会了用DGL进行简单的节点分类任务了，试着自己替换一下节点聚合方式吧。