第四周.02.RGCN代码讲解

本文内容整理自深度之眼《GNN核心能力培养计划》
参考案例：https://docs.dgl.ai/tutorials/models/1_gnn/4_rgcn.html
RGCN在DGL里面有集成一个RGCN的layer，这个例子里面我们用自定义的方式来实现。
再次强调整个图卷积模型的套路：
1.先自定义模型layer；
2.利用layer定义模型；
3.初始化后训练。

这里和原文一样，对于结点直接使用独热编码来进行初始化，然后在自定义模型layer中实现了两个步骤：
Compute outgoing message using node representation and weight matrix associated with the edge type (message function)
Aggregate incoming messages and generate new node representations (reduce and apply function)

导入包

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

套路步骤1

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

class RGCNLayer(nn.Module):
    # 参数说明：
    # in_feat 输入维度
    # out_feat 输出维度
    # num_rels 边类型数量
    # num_bases W_r分解的数量，对应原文公式3的B（求和符号的上界）
    # bias 偏置
    # activation 激活函数
    # is_input_layer 是否是输入层（第一层）
    def __init__(self, in_feat, out_feat, num_rels, num_bases=-1, bias=None,
                 activation=None, is_input_layer=False):
        super(RGCNLayer, self).__init__()
        self.in_feat = in_feat
        self.out_feat = out_feat
        self.num_rels = num_rels
        self.num_bases = num_bases
        self.bias = bias
        self.activation = activation
        self.is_input_layer = is_input_layer

        # sanity check
        # 矩阵分解的参数校验条件：不能小于0，不能比现有维度大（复杂度会变高，参数反而增加）
        if self.num_bases <= 0 or self.num_bases > self.num_rels:
            self.num_bases = self.num_rels

        # weight bases in equation (3)
        # 这里是根据公式3把W_r算出来，用V_b（weight）表示，共有num_bases个V_b累加得到
        # 得到的结果是Tensor，因此用 nn.Parameter将一个不可训练的类型Tensor
        # 转换成可以训练的类型parameter
        # 并将这个parameter绑定到这个module里面
        self.weight = nn.Parameter(torch.Tensor(self.num_bases, self.in_feat,
                                                self.out_feat))
        if self.num_bases < self.num_rels:
            # linear combination coefficients in equation (3)
            # 这里的w_comp是公式3里面的a_{rb}
            # 一个边类型对应一个W_r（那么就一共有num_rels种W_r），每个W_r分解为num_bases个组合
            # 因此w_comp这里的维度就是num_rels×num_bases
            self.w_comp = nn.Parameter(torch.Tensor(self.num_rels, self.num_bases))

        # add bias
        # 偏置要进行加法运算其维度要和输出维度大小一样
        if self.bias:
            self.bias = nn.Parameter(torch.Tensor(out_feat))

        # init trainable parameters
        # 这里用的是xavier初始化，可以查下为什么
        # 因为我们用的是线性激活函数，网上有说用ReLU和Leaky ReLU
        # 可以考虑用别的初始化方法：He init 。
        nn.init.xavier_uniform_(self.weight,
                                gain=nn.init.calculate_gain('relu'))
        if self.num_bases < self.num_rels:
            nn.init.xavier_uniform_(self.w_comp,
                                    gain=nn.init.calculate_gain('relu'))
        if self.bias:
            nn.init.xavier_uniform_(self.bias,
                                    gain=nn.init.calculate_gain('relu'))

    def forward(self, g):
        if self.num_bases < self.num_rels:#分解就走公式3
            # generate all weights from bases (equation (3))
            weight = self.weight.view(self.in_feat, self.num_bases, self.out_feat)
            weight = torch.matmul(self.w_comp, weight).view(self.num_rels,
                                                        self.in_feat, self.out_feat)
        else:#不分解就直接用weight算
            weight = self.weight

        if self.is_input_layer:#如果是第一层
            def message_func(edges):
                # for input layer, matrix multiply can be converted to be
                # an embedding lookup using source node id
                # 对于第一层，输入可以直接用独热编码进行aggregate
                # 信息的汇聚就可以直接写成矩阵相乘的形式
                embed = weight.view(-1, self.out_feat)# embed维度整成out_feat维度一样
                index = edges.data['rel_type'] * self.in_feat + edges.src['id']
                return {'msg': embed[index] * edges.data['norm']}
        else:
            def message_func(edges):
            	#根据边类型'rel_type'获取对应的w
                w = weight[edges.data['rel_type']]
                msg = torch.bmm(edges.src['h'].unsqueeze(1), w).squeeze()#消息汇聚，就是w乘以src['h']（输入节点特征）
                msg = msg * edges.data['norm']
                return {'msg': msg}

        def apply_func(nodes):
            h = nodes.data['h']
            if self.bias:#有偏置加偏置
                h = h + self.bias
            if self.activation:#经过激活函数
                h = self.activation(h)
            return {'h': h}

        g.update_all(message_func, fn.sum(msg='msg', out='h'), apply_func)

套路步骤2

class Model(nn.Module):
    def __init__(self, num_nodes, h_dim, out_dim, num_rels,
                 num_bases=-1, num_hidden_layers=1):
        # 先初始化参数
        super(Model, self).__init__()
        self.num_nodes = num_nodes
        self.h_dim = h_dim
        self.out_dim = out_dim
        self.num_rels = num_rels
        self.num_bases = num_bases
        self.num_hidden_layers = num_hidden_layers

        # create rgcn layers
        # 具体看下面函数
        self.build_model()

        # create initial features
        self.features = self.create_features()

    def build_model(self):
        self.layers = nn.ModuleList()
        # input to hidden
        i2h = self.build_input_layer()
        self.layers.append(i2h)
        # hidden to hidden
        for _ in range(self.num_hidden_layers):
            h2h = self.build_hidden_layer()
            self.layers.append(h2h)
        # hidden to output
        h2o = self.build_output_layer()
        self.layers.append(h2o)

    # initialize feature for each node
    def create_features(self):
        #torch.arange(start=0, end=5)的结果并不包含end，start默认是0
        features = torch.arange(self.num_nodes)
        return features

    # 输入num_nodes的独热编号
    # 输出h_dim
    # 激活函数是relu
    def build_input_layer(self):
        return RGCNLayer(self.num_nodes, self.h_dim, self.num_rels, self.num_bases,
                         activation=F.relu, is_input_layer=True)

    # 输入h_dim
    # 输出h_dim
    # 激活函数是relu
    def build_hidden_layer(self):
        return RGCNLayer(self.h_dim, self.h_dim, self.num_rels, self.num_bases,
                         activation=F.relu)

    # 输入h_dim
    # 输出out_dim
    # 激活函数是softmax后归一化
    def build_output_layer(self):
        return RGCNLayer(self.h_dim, self.out_dim, self.num_rels, self.num_bases,
                         activation=partial(F.softmax, dim=1))

    def forward(self, g):
        if self.features is not None:
            g.ndata['id'] = self.features
        for layer in self.layers:
            layer(g)
        return g.ndata.pop('h')

加载数据

这里第一次要先安装rdflib

pip install rdflib

# load graph data
from dgl.contrib.data import load_data
data = load_data(dataset='aifb')
num_nodes = data.num_nodes
num_rels = data.num_rels
num_classes = data.num_classes
labels = data.labels
train_idx = data.train_idx
# split training and validation set
val_idx = train_idx[:len(train_idx) // 5]#前20%做验证
train_idx = train_idx[len(train_idx) // 5:]#剩下做训练

# edge type and normalization factor
edge_type = torch.from_numpy(data.edge_type)
edge_norm = torch.from_numpy(data.edge_norm).unsqueeze(1)

labels = torch.from_numpy(labels).view(-1)

加载结果：

套路步骤3

初始化超参数，实例化模型

# configurations
n_hidden = 16 # number of hidden units
n_bases = -1 # use number of relations as number of bases这里设置-1相当于没进行分解
n_hidden_layers = 0 # use 1 input layer, 1 output layer, no hidden layer#输出接输出，没有中间层
n_epochs = 25 # epochs to train
lr = 0.01 # learning rate
l2norm = 0 # L2 norm coefficient

# create graph
g = DGLGraph((data.edge_src, data.edge_dst))
g.edata.update({'rel_type': edge_type, 'norm': edge_norm})

# create model
model = Model(len(g),
              n_hidden,
              num_classes,
              num_rels,
              num_bases=n_bases,
              num_hidden_layers=n_hidden_layers)

训练

# optimizer优化器
optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=l2norm)

print("start training...")
model.train()
for epoch in range(n_epochs):
    optimizer.zero_grad()
    logits = model.forward(g)
    loss = F.cross_entropy(logits[train_idx], labels[train_idx])
    loss.backward()

    optimizer.step()

    train_acc = torch.sum(logits[train_idx].argmax(dim=1) == labels[train_idx])
    train_acc = train_acc.item() / len(train_idx)
    val_loss = F.cross_entropy(logits[val_idx], labels[val_idx])
    val_acc = torch.sum(logits[val_idx].argmax(dim=1) == labels[val_idx])
    val_acc = val_acc.item() / len(val_idx)
    print("Epoch {:05d} | ".format(epoch) +
          "Train Accuracy: {:.4f} | Train Loss: {:.4f} | ".format(
              train_acc, loss.item()) +
          "Validation Accuracy: {:.4f} | Validation loss: {:.4f}".format(
              val_acc, val_loss.item()))

遇到报错：
tensor类型不对，没空改，坐等反馈