NARS.py

import dgl
import torch as th
import torch.nn as nn
from . import BaseModel, register_model
import dgl.function as fn


@register_model('NARS')
class NARS(BaseModel):
    r"""
        `SCALABLE GRAPH NEURAL NETWORKS FOR HETEROGENEOUS GRAPHS <https://arxiv.org/pdf/2011.09679.pdf>`_.

        Given a heterogeneous graph :math:`G` and its edge relation type set :math:`\mathcal{R}`,
        our proposed method first samples :math:`K` unique subsets from :math:`\mathcal{R}`.
        Then for each sampled subset :math:`R_i \subseteq \mathcal{R}`, we generate a relation subgraph
        :math:`G_i` from :math:`G` in which only edges whose type belongs to :math:`R_i` are kept.
        We treat :math:`G_i` as a homogeneous graph or a bipartite graph, and perform neighbor aggregation to generate :math:`L`-hop neighbor features for each node.
        Let :math:`H_{v,0}` be the input features (of dimension :math:`D`) for node :math:`v`. For each subgraph :math:`G_i`
        , the :math:`l`-th hop
        features :math:`H_{v,l}^{i}` are computed as

        .. math::
            H_{v, l}^{i}=\sum_{u \in N_{i}(v)} \frac{1}{\left|N_{i}(v)\right|} H_{u, l-1}^{i}


        where :math:`N_i(v)` is the set of neighbors of node :math:`v` in :math:`G_i`.

        For each layer :math:`l`, we let the model adaptively learn which relation-subgraph features to use by aggregating
        features from different subgraphs :math:`G_i` with learnable 1-D convolution. The aggregated :math:`l`-hop
        features across all subgraphs are calculated as


        .. math::
            H_{v, l}^{a g g}=\sum_{i=1}^{K} a_{i, l} \cdot H_{v, l}^{i}


        where :math:`H^i` is the neighbor averaging features on subgraph :math:`G_i` and :math:`a_{i,l}` is a learned vector of length equal
        to the feature dimension :math:`D`.

        Parameters
        ----------
        num_hops : int
            Number of hops.
        category : str
            Type of predicted nodes.
        hidden_dim : int
            The dimention of hidden layer.
        num_feats : int
            The number of relation subsets.

        Note
        ----
        We do not support the dataset without feature, (e.g. HGBn-Freebase
        because the model performs neighbor aggregation to generate :math:`L`-hop neighbor features at once.

        """

    @classmethod
    def build_model_from_args(cls, args, hg):
        num_hops = args.num_hops + 1

        return cls(num_hops=num_hops,
                   args = args,
                   hg = hg,
                   )

    def __init__(self, num_hops, args, hg):
        super(NARS, self).__init__()
        self.category = args.category
        self.dropout = args.dropout
        self.input_dropout = args.input_dropout
        # self.feats = feats
        self.device = args.device
        self.num_hops = num_hops
        self.args = args

        in_size = hg.nodes[args.category].data["h"].shape[1]

        etypes = hg.canonical_etypes
        mps = []
        for etype in etypes:
            if etype[0] == args.category:
                for dst_e in etypes:
                    if etype[0] == dst_e[2] and etype[2] == dst_e[0] and etype[0] != etype[2]:
                        mps.append([etype, dst_e])
        self.mps = mps
        self.num_feats = len(mps)

        with th.no_grad():
            self.feats = preprocess_features(hg, mps, args, args.device, self.args.category)
            print("Done preprocessing")

        self.seq = nn.Sequential(
            WeightedAggregator(self.num_feats, in_size, num_hops),
            SIGN(in_size, args.hidden_dim, args.out_dim, num_hops,
                 args.ff_layer, args.dropout, args.input_dropout)
        )

    def forward(self, hg, h_dict):

        #ffeats = [x.to(self.device) for x in self.feats]
        ffeats = [x.to(self.device) for x in self.feats]
        return {self.category: self.seq.forward(ffeats)}

    def reset_parameters(self):
        self.seq.register_parameter()


def preprocess_features(g, mps, args, device, predict):
    """
        pre-process heterogeneous graph g to generate neighbor-averaged features
        for each relation subsets

        Parameters
        -----------
        g :
            heterogeneous graph
        rel_subsets :
            relations of subsets
        args :
            arguments
        device :
            device

        Return
        ------
            new features of each relation subsets

    """
    category_dim = g.nodes[predict].data["h"].shape[1]
    for ntype in g.ntypes:
        ntype_dim = g.nodes[ntype].data["h"].shape[1]
        if category_dim != ntype_dim:
            rand_weight = th.Tensor(ntype_dim, category_dim).uniform_(-0.5, 0.5).to(device)
            g.nodes[ntype].data["h"] = th.matmul(g.nodes[ntype].data["h"], rand_weight)

    num_paper, feat_size = g.nodes[predict].data["h"].shape

    new_feats = [th.zeros(num_paper, len(mps), feat_size) for _ in range(args.num_hops + 1)]

    for subset_id, subset in enumerate(mps):
        # print(subset)
        feats = gen_rel_subset_feature(g, subset, args, device, predict)
        for i in range(args.num_hops + 1):
            feat = feats[i]
            new_feats[i][:feat.shape[0], subset_id, :] = feat
        feats = None
    return new_feats

def gen_rel_subset_feature(g, rel_subset, args, device, predict):
    """
        Build relation subgraph given relation subset and generate multi-hop
        neighbor-averaged feature on this subgraph

        Parameters
        ----------
        g :
            Heterogeneous graph
        rel_subset :
            relation of subsets
        args :
            arguments
        device :
            device

        Returns
        ------
        new features of a relation subsets
    """

    new_g = g.edge_type_subgraph(rel_subset).to(g.device)

    ntypes = new_g.ntypes
    # set node feature and calc deg
    for ntype in ntypes:
        num_nodes = new_g.number_of_nodes(ntype)
        if num_nodes < g.nodes[ntype].data["h"].shape[0]:
            new_g.nodes[ntype].data["hop_0"] = g.nodes[ntype].data["h"][:num_nodes, :]
        else:
            new_g.nodes[ntype].data["hop_0"] = g.nodes[ntype].data["h"]
        deg = 0
        for etype in new_g.etypes:
            _, _, dtype = new_g.to_canonical_etype(etype)
            if ntype == dtype:
                deg = deg + new_g.in_degrees(etype=etype)
        if th.is_tensor(deg):
            norm = 1.0 / deg.float()
            norm[th.isinf(norm)] = 0
            new_g.nodes[ntype].data["norm"] = norm.view(-1, 1).to(device)

    res = []

    # compute k-hop feature
    for hop in range(1, args.num_hops + 1):
        ntype2feat = {}
        for etype in new_g.etypes:
            stype, _, dtype = new_g.to_canonical_etype(etype)
            new_g[etype].update_all(fn.copy_u(f'hop_{hop-1}', 'm'), fn.sum('m', 'new_feat'))
            new_feat = new_g.nodes[dtype].data.pop("new_feat")
            assert("new_feat" not in new_g.nodes[stype].data)
            if dtype in ntype2feat:
                ntype2feat[dtype] += new_feat
            else:
                ntype2feat[dtype] = new_feat
        for ntype in new_g.ntypes:
            assert ntype in ntype2feat  # because subgraph is not directional
            feat_dict = new_g.nodes[ntype].data
            old_feat = feat_dict.pop(f"hop_{hop-1}")
            if ntype == predict:
                res.append(old_feat.cpu())
            feat_dict[f"hop_{hop}"] = ntype2feat.pop(ntype).mul_(feat_dict["norm"])

    res.append(new_g.nodes[predict].data.pop(f"hop_{args.num_hops}").cpu())
    return res



class FeedForwardNet(nn.Module):
    """
        A feedforward net.

        Input
        ------
        in_feats :
            input feature dimention
        hidden :
            hidden layer dimention
        out_feats :
            output feature dimention
        num_layers :
            number of layers
        dropout :
            dropout rate
    """
    def __init__(self, in_feats, hidden, out_feats, num_layers, dropout):
        super(FeedForwardNet, self).__init__()
        self.layers = nn.ModuleList()
        self.num_layers = num_layers
        if num_layers == 1:
            self.layers.append(nn.Linear(in_feats, out_feats))
        else:
            self.layers.append(nn.Linear(in_feats, hidden))
            for i in range(num_layers - 2):
                self.layers.append(nn.Linear(hidden, hidden))
            self.layers.append(nn.Linear(hidden, out_feats))
        if self.num_layers > 1:
            self.prelu = nn.PReLU()
            self.dropout = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self):
        gain = nn.init.calculate_gain("relu")
        for layer in self.layers:
            nn.init.xavier_uniform_(layer.weight, gain=gain)
            nn.init.zeros_(layer.bias)

    def forward(self, x):
        for layer_id, layer in enumerate(self.layers):
            x = layer(x)
            if layer_id < self.num_layers - 1:
                x = self.dropout(self.prelu(x))
        return x


class SIGN(nn.Module):
    """
        The SIGN model.

        Parameters
        ------------
        in_feats :
            input feature dimention
        hidden :
            hidden layer dimention
        out_feats :
            output feature dimention
        num_hops :
            number of hops
        num_layers :
            number of layers
        dropout :
            dropout rate
        input_drop :
            whether or not to dropout when inputting features

    """
    def __init__(
        self, in_feats, hidden, out_feats, num_hops, num_layers, dropout, input_drop
    ):
        super(SIGN, self).__init__()
        self.dropout = nn.Dropout(dropout)
        self.prelu = nn.PReLU()
        self.inception_ffs = nn.ModuleList()
        self.input_drop = input_drop
        for i in range(num_hops):
            self.inception_ffs.append(
                FeedForwardNet(in_feats, hidden, hidden, num_layers, dropout)
            )
        self.project = FeedForwardNet(
            num_hops * hidden, hidden, out_feats, num_layers, dropout
        )

    def forward(self, feats):
        hidden = []
        for feat, ff in zip(feats, self.inception_ffs):
            if self.input_drop:
                feat = self.dropout(feat)
            hidden.append(ff(feat))
        out = self.project(self.dropout(self.prelu(th.cat(hidden, dim=-1))))
        return out


class WeightedAggregator(nn.Module):
    """
        Get new features by multiplying the old features by the weight matrix.

        Parameters
        -------------
        num_feats :
            number of subsets
        in_feats :
            input feature dimention
        num_hops :
            number of hops


    """
    def __init__(self, num_feats, in_feats, num_hops):
        super(WeightedAggregator, self).__init__()
        self.agg_feats = nn.ParameterList()
        for _ in range(num_hops):
            self.agg_feats.append(nn.Parameter(th.Tensor(num_feats, in_feats)))
            nn.init.xavier_uniform_(self.agg_feats[-1])

    def forward(self, feats):
        new_feats = []
        for feat, weight in zip(feats, self.agg_feats):
            new_feats.append((feat * weight.unsqueeze(0)).sum(dim=1).squeeze())
        return new_feats