Rsage.py

import dgl
import torch as th
import torch.nn as nn
import torch.nn.functional as F
import dgl.nn as dglnn
from . import BaseModel, register_model


@register_model('Rsage')
class Rsage(BaseModel):
    @classmethod
    def build_model_from_args(cls, args, hg):
        return cls(in_dim=args.in_dim,
                   out_dim=args.out_dim,
                   h_dim=args.hidden_dim,
                   etypes=hg.etypes,
                   aggregator_type=args.aggregator_type,
                   num_hidden_layers=args.num_layers - 2,
                   dropout=args.dropout)

    def __init__(self, in_dim,
                 out_dim, 
                 h_dim, 
                 etypes, 
                 aggregator_type,
                 num_hidden_layers=1, 
                 dropout=0):
        super(Rsage, self).__init__()
        self.rel_names = etypes
        self.layers = nn.ModuleList()
        # input 2 hidden
        self.layers.append(RsageLayer(
            in_dim, h_dim, aggregator_type, self.rel_names, activation=F.relu, dropout=dropout))
        for i in range(num_hidden_layers):
            self.layers.append(RsageLayer(
                h_dim, h_dim, aggregator_type, self.rel_names, activation=F.relu, dropout=dropout
            ))
        self.layers.append(RsageLayer(
            h_dim, out_dim, aggregator_type, self.rel_names, activation=None))
        return

    def forward(self, hg, h_dict=None):
        if hasattr(hg, 'ntypes'):
            # full graph training,
            for layer in self.layers:
                h_dict = layer(hg, h_dict)
        else:
            # minibatch training, block
            for layer, block in zip(self.layers, hg):
                h_dict = layer(block, h_dict)
        return h_dict


class RsageLayer(nn.Module):

    def __init__(self,
                 in_feat,
                 out_feat,
                 aggregator_type,
                 rel_names,
                 activation=None,
                 dropout=0.0,
                 bias=True):
        super(RsageLayer, self).__init__()
        self.in_feat = in_feat
        self.out_feat = out_feat
        self.aggregator_type = aggregator_type
        self.activation = activation
        self.dropout = nn.Dropout(dropout)
        self.conv = dglnn.HeteroGraphConv({
            rel: dgl.nn.pytorch.SAGEConv(in_feat, out_feat, aggregator_type, bias=bias)
            for rel in rel_names
        })

    def forward(self, g, h_dict):
        h_dict = self.conv(g, h_dict)
        out_put = {}
        for n_type, h in h_dict.items():
            out_put[n_type] = h.squeeze()
        return out_put