DGL系列之(二)：使用DGL實作GCN

import dgl
import dgl.function as fn
import torch as th
import torch.nn as nn
import torch.nn.functional as F
from dgl import DGLGraph

gcn_msg=fn.copy_src(src="h",out="m")
gcn_reduce=fn.sum(msg="m",out="h")#聚合鄰居節點的特征

#定義節點的UDF apply_nodes  他是一個完全連接配接層
class NodeApplyModule(nn.Module):
    #初始化
    def __init__(self,in_feats,out_feats,activation):
        super(NodeApplyModule,self).__init__()
        self.linear=nn.Linear(in_feats,out_feats)
        self.activation=activation
    #前向傳播
    def forward(self,node):
        h=self.linear(node.data["h"])
        if self.activation is not None:
            h=self.activation(h)
        return {"h":h}

#定義GCN子產品  GCN子產品的本質是在所有節點上執行消息傳遞  然後再調用NOdeApplyModule全連接配接層
class GCN(nn.Module):
    #初始化
    def __init__(self,in_feats,out_feats,activation):
        super(GCN,self).__init__()
        #調用全連接配接層子產品
        self.apply_mod=NodeApplyModule(in_feats,out_feats,activation)
    #前向傳播
    def forward(self,g,feature):
        g.ndata["h"]=feature#feature應該對應的整個圖的特征矩陣
        g.update_all(gcn_msg,gcn_reduce)
        g.apply_nodes(func=self.apply_mod)#将更新操作應用到節點上
        
        return g.ndata.pop("h")

#利用cora資料集搭建網絡然後訓練
class Net(nn.Module):
    #初始化網絡參數
    def __init__(self):
        super(Net,self).__init__()
        self.gcn1=GCN(1433,16,F.relu)#第一層GCN
        self.gcn2=GCN(16,7,None)
    #前向傳播
    def forward(self,g,features):
        x=self.gcn1(g,features)
        x=self.gcn2(g,x)
        return x
net=Net()
net

#使用DGL内置子產品加載cora資料集
from dgl.data import citation_graph as citegrh
import networkx as nx
def load_cora_data():
    data = citegrh.load_cora()#加載資料集
    features=th. FloatTensor(data.features)#特征向量  張量的形式
    labels=th.LongTensor(data.labels)#所屬類别
    train_mask=th.BoolTensor(data.train_mask)#那些參與訓練
    test_mask=th.BoolTensor(data.test_mask)#哪些是測試集
    g=data.graph
    g.remove_edges_from(nx.selfloop_edges(g))#删除自循環的邊
    g = DGLGraph(g)
    g.add_edges(g.nodes(), g.nodes())
    return g, features, labels, train_mask, test_mask

g, features, labels, train_mask, test_mask=load_cora_data()


import matplotlib.pyplot as plt
nx.draw(g.to_networkx(),node_size=50,with_labels=True)
plt.show()

#測試模型
def evaluate(model, g, features, labels, mask):
    model.eval()#會通知所有圖層您處于評估模式
    with th.no_grad():
        logits = model(g, features)
        logits = logits[mask]
        labels = labels[mask]
        _, indices = th.max(logits, dim=1)
        correct = th.sum(indices == labels)
        return correct.item() * 1.0 / len(labels)

#訓練網絡
import time
import numpy as np
g, features, labels, train_mask, test_mask = load_cora_data()

#定義優化器
optimizer=th.optim.Adam(net.parameters(),lr=1e-3)
dur=[]#時間
for epoch in range(100):
    print(epoch)
    if epoch>=3:
        t0=time.time()
    net.train()
    logits = net(g, features)
    logp = F.log_softmax(logits, 1)
    loss = F.nll_loss(logp[train_mask], labels[train_mask])

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch >=3:
        dur.append(time.time() - t0)

    acc = evaluate(net, g, features, labels, test_mask)
    print("Epoch {:05d} | Loss {:.4f} | Test Acc {:.4f} | Time(s) {:.4f}".format(
            epoch, loss.item(), acc, np.mean(dur)))