/examples/pytorch/caregnn/model.py

import torch as th
import numpy as np
import torch.nn as nn
import dgl.function as fn


class CAREConv(nn.Module):
    """One layer of CARE-GNN."""

    def __init__(self, in_dim, out_dim, num_classes, edges, activation=None, step_size=0.02):
        super(CAREConv, self).__init__()

        self.activation = activation
        self.step_size = step_size
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.num_classes = num_classes
        self.edges = edges
        self.dist = {}

        self.linear = nn.Linear(self.in_dim, self.out_dim)
        self.MLP = nn.Linear(self.in_dim, self.num_classes)

        self.p = {}
        self.last_avg_dist = {}
        self.f = {}
        self.cvg = {}
        for etype in edges:
            self.p[etype] = 0.5
            self.last_avg_dist[etype] = 0
            self.f[etype] = []
            self.cvg[etype] = False

    def _calc_distance(self, edges):
        # formula 2
        d = th.norm(th.tanh(self.MLP(edges.src['h'])) - th.tanh(self.MLP(edges.dst['h'])), 1, 1)
        return {'d': d}

    def _top_p_sampling(self, g, p):
        # this implementation is low efficient
        # optimization requires dgl.sampling.select_top_p requested in issue #3100
        dist = g.edata['d']
        neigh_list = []
        for node in g.nodes():
            edges = g.in_edges(node, form='eid')
            num_neigh = int(g.in_degrees(node) * p)
            neigh_dist = dist[edges]
            neigh_index = np.argpartition(neigh_dist.cpu().detach(), num_neigh)[:num_neigh]
            neigh_list.append(edges[neigh_index])
        return th.cat(neigh_list)

    def forward(self, g, feat):
        with g.local_scope():
            g.ndata['h'] = feat

            hr = {}
            for i, etype in enumerate(g.canonical_etypes):
                g.apply_edges(self._calc_distance, etype=etype)
                self.dist[etype] = g.edges[etype].data['d']
                sampled_edges = self._top_p_sampling(g[etype], self.p[etype])

                # formula 8
                g.send_and_recv(sampled_edges, fn.copy_u('h', 'm'), fn.mean('m', 'h_%s' % etype[1]), etype=etype)
                hr[etype] = g.ndata['h_%s' % etype[1]]
                if self.activation is not None:
                    hr[etype] = self.activation(hr[etype])

            # formula 9 using mean as inter-relation aggregator
            p_tensor = th.Tensor(list(self.p.values())).view(-1, 1, 1).to(g.device)
            h_homo = th.sum(th.stack(list(hr.values())) * p_tensor, dim=0)
            h_homo += feat
            if self.activation is not None:
                h_homo = self.activation(h_homo)

            return self.linear(h_homo)


class CAREGNN(nn.Module):
    def __init__(self,
                 in_dim,
                 num_classes,
                 hid_dim=64,
                 edges=None,
                 num_layers=2,
                 activation=None,
                 step_size=0.02):
        super(CAREGNN, self).__init__()
        self.in_dim = in_dim
        self.hid_dim = hid_dim
        self.num_classes = num_classes
        self.edges = edges
        self.activation = activation
        self.step_size = step_size
        self.num_layers = num_layers

        self.layers = nn.ModuleList()

        if self.num_layers == 1:
            # Single layer
            self.layers.append(CAREConv(self.in_dim,
                                        self.num_classes,
                                        self.num_classes,
                                        self.edges,
                                        activation=self.activation,
                                        step_size=self.step_size))

        else:
            # Input layer
            self.layers.append(CAREConv(self.in_dim,
                                        self.hid_dim,
                                        self.num_classes,
                                        self.edges,
                                        activation=self.activation,
                                        step_size=self.step_size))

            # Hidden layers with n - 2 layers
            for i in range(self.num_layers - 2):
                self.layers.append(CAREConv(self.hid_dim,
                                            self.hid_dim,
                                            self.num_classes,
                                            self.edges,
                                            activation=self.activation,
                                            step_size=self.step_size))

            # Output layer
            self.layers.append(CAREConv(self.hid_dim,
                                        self.num_classes,
                                        self.num_classes,
                                        self.edges,
                                        activation=self.activation,
                                        step_size=self.step_size))

    def forward(self, graph, feat):
        # For full graph training, directly use the graph
        # formula 4
        sim = th.tanh(self.layers[0].MLP(feat))

        # Forward of n layers of CARE-GNN
        for layer in self.layers:
            feat = layer(graph, feat)

        return feat, sim

    def RLModule(self, graph, epoch, idx):
        for layer in self.layers:
            for etype in self.edges:
                if not layer.cvg[etype]:
                    # formula 5
                    eid = graph.in_edges(idx, form='eid', etype=etype)
                    avg_dist = th.mean(layer.dist[etype][eid])

                    # formula 6
                    if layer.last_avg_dist[etype] < avg_dist:
                        if layer.p[etype] - self.step_size > 0:
                            layer.p[etype] -= self.step_size
                        layer.f[etype].append(-1)
                    else:
                        if layer.p[etype] + self.step_size <= 1:
                            layer.p[etype] += self.step_size
                        layer.f[etype].append(+1)
                    layer.last_avg_dist[etype] = avg_dist

                    # formula 7
                    if epoch >= 9 and abs(sum(layer.f[etype][-10:])) <= 2:
                        layer.cvg[etype] = True