CompGCN.py

import torch as th
import torch.nn as nn
import dgl.nn as dglnn
import torch.nn.functional as F
from . import BaseModel, register_model
from openhgnn.layers.micro_layer import CompConv
from ..utils.dgl_graph import edata_in_out_mask
from ..utils import get_nodes_dict
from ..utils.utils import ccorr

@register_model('CompGCN')
class CompGCN(BaseModel):
    """
    The models of the simplified CompGCN, without using basis vector, for a heterogeneous graph.

    Here, we present the implementation details for each task used for evaluation in the paper. 
    For all the tasks, we used COMPGCN build on PyTorch geometric framework (Fey & Lenssen, 2019).

    Link Prediction: For evaluation, 200-dimensional embeddings for node and relation embeddings are used. 
    For selecting the best model we perform a hyperparameter search using the validation data over the values listed in Table 8. 
    For training link prediction models, we use the standard binary cross entropy loss with label smoothing Dettmers et al. (2018).

    Node Classification: Following Schlichtkrull et al. (2017), we use 10% training data as validation for selecting the best model for both the datasets. 
    We restrict the number of hidden units to 32. We use cross-entropy loss for training our model.

    For all the experiments, training is done using Adam optimizer (Kingma & Ba, 2014) and Xavier initialization (Glorot & Bengio, 2010) is used for initializing parameters.
    """

    @classmethod
    def build_model_from_args(cls, args, hg):
        return cls(args.hidden_dim, args.hidden_dim, args.out_dim,
                   hg.etypes,
                   get_nodes_dict(hg), len(hg.etypes),
                   args.num_layers, comp_fn=args.comp_fn, dropout=args.dropout
                   )

    def __init__(self, in_dim, hid_dim, out_dim, etypes, n_nodes, n_rels, num_layers=2, comp_fn='sub', dropout=0.0,
                 activation=F.relu, batchnorm=True):
        super(CompGCN, self).__init__()
        self.in_dim = in_dim
        self.hid_dim = hid_dim
        self.out_dim = out_dim
        self.rel_names = list(set(etypes))
        self.rel_names.sort()
        self.n_rels = n_rels
        self.n_nodes = n_nodes
        self.num_layer = num_layers
        self.comp_fn = comp_fn
        self.dropout = dropout
        self.activation = activation
        self.batchnorm = batchnorm

        self.layers = nn.ModuleList()

        self.r_embedding = nn.Parameter(th.FloatTensor(self.n_rels + 1, self.in_dim))

        self.layers.append(CompGraphConvLayer(self.in_dim,
                                              self.hid_dim,
                                              self.rel_names,
                                              comp_fn=self.comp_fn,
                                              activation=self.activation,
                                              batchnorm=self.batchnorm,
                                              dropout=self.dropout))

        # Hidden layers with n - 1 CompGraphConv layers
        for i in range(self.num_layer - 2):
            self.layers.append(CompGraphConvLayer(self.hid_dim,
                                                  self.hid_dim,
                                                  self.rel_names,
                                                  comp_fn=self.comp_fn,
                                                  activation=self.activation,
                                                  batchnorm=self.batchnorm,
                                                  dropout=self.dropout))

        # Output layer with the output class
        self.layers.append(CompGraphConvLayer(self.hid_dim,
                                              self.out_dim,
                                              self.rel_names,
                                              comp_fn=self.comp_fn))
        nn.init.xavier_uniform_(self.r_embedding)

    def forward(self, hg, n_feats):

        # For full graph training, directly use the graph
        # Forward of n layers of CompGraphConv
        r_feats = self.r_embedding
        if hasattr(hg, 'ntypes'):
            # full graph training
            for layer in self.layers:
                n_feats, r_feats = layer(hg, n_feats, r_feats)
        else:
            # mini-batch training
            for layer, block in zip(self.layers, hg):
                n_feats, r_feats = layer(block, n_feats, r_feats)

        return n_feats

    def preprocess(self, hg):
        edata_in_out_mask(hg)


class CompGraphConvLayer(nn.Module):
    """One layer of simplified CompGCN."""

    def __init__(self,
                 in_dim,
                 out_dim,
                 rel_names,
                 comp_fn='sub',
                 activation=None,
                 batchnorm=False,
                 dropout=0):
        super(CompGraphConvLayer, self).__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.comp_fn = comp_fn
        self.actvation = activation
        self.batchnorm = batchnorm
        self.rel_names = rel_names
        # define dropout layer
        self.dropout = nn.Dropout(dropout)

        # define batch norm layer
        if self.batchnorm:
            self.bn = nn.BatchNorm1d(out_dim)
        # define weights of 3 node matrices
        self.W_O = nn.Linear(self.in_dim, self.out_dim)
        self.W_I = nn.Linear(self.in_dim, self.out_dim)
        self.W_S = nn.Linear(self.in_dim, self.out_dim)

        # define weights of the 1 relation matrix
        self.W_R = nn.Linear(self.in_dim, self.out_dim)
        self.conv = dglnn.HeteroGraphConv({
            rel: CompConv(comp_fn=comp_fn,
                          norm='right',
                          _allow_zero_in_degree=True)
            for rel in rel_names
        })

    def forward(self, hg, n_in_feats, r_feats):
        """
        Compute one layer of composition transfer for one relation only in a
        homogeneous graph with bidirectional edges.
        """
        with hg.local_scope():
            # Assign values to source nodes. In a homogeneous graph, this is equal to
            # assigning them to all nodes.

            # Assign feature to all edges with the same value, the r_feats.

            wdict = {}
            for i, etype in enumerate(self.rel_names):
                if etype[:4] == 'rev-' or etype[-4:] == '-rev':
                    W = self.W_I
                else:
                    W = self.W_O
                wdict[etype] = {'Linear': W, 'h_e': r_feats[i + 1]}

            if hg.is_block:
                inputs_src = n_in_feats
                inputs_dst = {k: v[:hg.number_of_dst_nodes(k)] for k, v in n_in_feats.items()}
            else:
                inputs_src = inputs_dst = n_in_feats
            outputs = self.conv(hg, inputs_src, mod_kwargs=wdict)

            for n, emd in outputs.items():
                # Step 4: add results of self-loop
                if self.comp_fn == 'sub':
                    h_self = self.W_S(inputs_dst[n] - r_feats[-1])
                elif self.comp_fn == 'mul':
                    h_self = self.W_S(inputs_dst[n] * r_feats[-1])
                elif self.comp_fn == 'ccorr':
                    h_self = self.W_S(ccorr(inputs_dst[n], r_feats[-1]))
                else:
                    raise Exception('Only supports sub, mul, and ccorr')
                h_self.add_(emd)
                # Compute relation output

                # Use batch norm
                if self.batchnorm:
                    if h_self.shape[0] > 1:
                        h_self = self.bn(h_self)

                # Use drop out
                n_out_feats = self.dropout(h_self)
                # Use activation function
                if self.actvation is not None:
                    n_out_feats = self.actvation(n_out_feats)

                outputs[n] = n_out_feats

        r_out_feats = self.W_R(r_feats)
        r_out_feats = self.dropout(r_out_feats)
        if self.actvation is not None:
            r_out_feats = self.actvation(r_out_feats)
        return outputs, r_out_feats