model.py

import mxnet as mx
from mxnet import gluon


class BaseRGCN(gluon.Block):
    def __init__(
        self,
        num_nodes,
        h_dim,
        out_dim,
        num_rels,
        num_bases=-1,
        num_hidden_layers=1,
        dropout=0,
        use_self_loop=False,
        gpu_id=-1,
    ):
        super(BaseRGCN, 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
        self.dropout = dropout
        self.use_self_loop = use_self_loop
        self.gpu_id = gpu_id

        # create rgcn layers
        self.build_model()

    def build_model(self):
        self.layers = gluon.nn.Sequential()
        # i2h
        i2h = self.build_input_layer()
        if i2h is not None:
            self.layers.add(i2h)
        # h2h
        for idx in range(self.num_hidden_layers):
            h2h = self.build_hidden_layer(idx)
            self.layers.add(h2h)
        # h2o
        h2o = self.build_output_layer()
        if h2o is not None:
            self.layers.add(h2o)

    def build_input_layer(self):
        return None

    def build_hidden_layer(self):
        raise NotImplementedError

    def build_output_layer(self):
        return None

    def forward(self, g, h, r, norm):
        for layer in self.layers:
            h = layer(g, h, r, norm)
        return h