chore: import upstream snapshot with attribution

examples/pytorch/diffpool/README.md

@@ -0,0 +1,59 @@
+Hierarchical Graph Representation Learning with Differentiable Pooling
+============
+
+
+Paper link: [https://arxiv.org/abs/1806.08804](https://arxiv.org/abs/1806.08804)
+
+Author's code repo: [https://github.com/RexYing/diffpool](https://github.com/RexYing/diffpool)
+
+This folder contains a DGL implementation of the DiffPool model. The first pooling layer is computed with DGL, and following pooling layers are computed with tensorized operation since the pooled graphs are dense.
+
+Dependencies
+------------
+* PyTorch 1.0+
+
+How to run
+----------
+
+```bash
+python train.py --dataset ENZYMES --pool_ratio 0.10 --num_pool 1 --epochs 1000
+python train.py --dataset DD --pool_ratio 0.15 --num_pool 1  --batch-size 10
+```
+Performance
+-----------
+ENZYMES 63.33% (with early stopping)
+DD 79.31% (with early stopping)
+
+
+## Update (2021-03-09)
+
+**Changes:**
+
+* Fix bug in Diffpool: the wrong `assign_dim` parameter
+* Improve efficiency of DiffPool, make the model independent of batch size. Remove redundant computation.
+
+
+**Efficiency:**
+
+On V100-SXM2 16GB
+
+|                    | Train time/epoch (original) (s) | Train time/epoch (improved) (s) |
+| ------------------ | ------------------------------: | ------------------------------: |
+| DD (batch_size=10) |                          21.302 |                      **17.282** |
+| DD (batch_size=20) |                             OOM |                      **44.682** |
+| ENZYMES            |                           1.749 |                       **1.685** |
+
+|                    | Memory usage (original) (MB) | Memory usage (improved) (MB) |
+| ------------------ | ---------------------------: | ---------------------------: |
+| DD (batch_size=10) |                     5274.620 |                 **2928.568** |
+| DD (batch_size=20) |                          OOM |                **10088.889** |
+| ENZYMES            |                       25.685 |                   **21.909** |
+
+**Accuracy**
+
+Each experiment with improved model is only conducted once, thus the result may has noise.
+
+|         |   Original |   Improved |
+| ------- | ---------: | ---------: |
+| DD      | **79.31%** |     78.33% |
+| ENZYMES |     63.33% | **68.33%** |

examples/pytorch/diffpool/data_utils.py

@@ -0,0 +1,40 @@
+import numpy as np
+import torch
+
+
+def one_hotify(labels, pad=-1):
+    """
+    cast label to one hot vector
+    """
+    num_instances = len(labels)
+    if pad <= 0:
+        dim_embedding = np.max(labels) + 1  # zero-indexed assumed
+    else:
+        assert pad > 0, "result_dim for padding one hot embedding not set!"
+        dim_embedding = pad + 1
+    embeddings = np.zeros((num_instances, dim_embedding))
+    embeddings[np.arange(num_instances), labels] = 1
+
+    return embeddings
+
+
+def pre_process(dataset, prog_args):
+    """
+    diffpool specific data partition, pre-process and shuffling
+    """
+    if prog_args.data_mode != "default":
+        print("overwrite node attributes with DiffPool's preprocess setting")
+        if prog_args.data_mode == "id":
+            for g, _ in dataset:
+                id_list = np.arange(g.num_nodes())
+                g.ndata["feat"] = one_hotify(id_list, pad=dataset.max_num_node)
+
+        elif prog_args.data_mode == "deg-num":
+            for g, _ in dataset:
+                g.ndata["feat"] = np.expand_dims(g.in_degrees(), axis=1)
+
+        elif prog_args.data_mode == "deg":
+            for g in dataset:
+                degs = list(g.in_degrees())
+                degs_one_hot = one_hotify(degs, pad=dataset.max_degrees)
+                g.ndata["feat"] = degs_one_hot

examples/pytorch/diffpool/model/dgl_layers/__init__.py

@@ -0,0 +1 @@
+from .gnn import DiffPoolBatchedGraphLayer, GraphSage, GraphSageLayer

examples/pytorch/diffpool/model/dgl_layers/aggregator.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Aggregator(nn.Module):
+    """
+    Base Aggregator class. Adapting
+    from PR# 403
+
+    This class is not supposed to be called
+    """
+
+    def __init__(self):
+        super(Aggregator, self).__init__()
+
+    def forward(self, node):
+        neighbour = node.mailbox["m"]
+        c = self.aggre(neighbour)
+        return {"c": c}
+
+    def aggre(self, neighbour):
+        # N x F
+        raise NotImplementedError
+
+
+class MeanAggregator(Aggregator):
+    """
+    Mean Aggregator for graphsage
+    """
+
+    def __init__(self):
+        super(MeanAggregator, self).__init__()
+
+    def aggre(self, neighbour):
+        mean_neighbour = torch.mean(neighbour, dim=1)
+        return mean_neighbour
+
+
+class MaxPoolAggregator(Aggregator):
+    """
+    Maxpooling aggregator for graphsage
+    """
+
+    def __init__(self, in_feats, out_feats, activation, bias):
+        super(MaxPoolAggregator, self).__init__()
+        self.linear = nn.Linear(in_feats, out_feats, bias=bias)
+        self.activation = activation
+        # Xavier initialization of weight
+        nn.init.xavier_uniform_(
+            self.linear.weight, gain=nn.init.calculate_gain("relu")
+        )
+
+    def aggre(self, neighbour):
+        neighbour = self.linear(neighbour)
+        if self.activation:
+            neighbour = self.activation(neighbour)
+        maxpool_neighbour = torch.max(neighbour, dim=1)[0]
+        return maxpool_neighbour
+
+
+class LSTMAggregator(Aggregator):
+    """
+    LSTM aggregator for graphsage
+    """
+
+    def __init__(self, in_feats, hidden_feats):
+        super(LSTMAggregator, self).__init__()
+        self.lstm = nn.LSTM(in_feats, hidden_feats, batch_first=True)
+        self.hidden_dim = hidden_feats
+        self.hidden = self.init_hidden()
+
+        nn.init.xavier_uniform_(
+            self.lstm.weight, gain=nn.init.calculate_gain("relu")
+        )
+
+    def init_hidden(self):
+        """
+        Defaulted to initialite all zero
+        """
+        return (
+            torch.zeros(1, 1, self.hidden_dim),
+            torch.zeros(1, 1, self.hidden_dim),
+        )
+
+    def aggre(self, neighbours):
+        """
+        aggregation function
+        """
+        # N X F
+        rand_order = torch.randperm(neighbours.size()[1])
+        neighbours = neighbours[:, rand_order, :]
+
+        (lstm_out, self.hidden) = self.lstm(
+            neighbours.view(neighbours.size()[0], neighbours.size()[1], -1)
+        )
+        return lstm_out[:, -1, :]
+
+    def forward(self, node):
+        neighbour = node.mailbox["m"]
+        c = self.aggre(neighbour)
+        return {"c": c}

examples/pytorch/diffpool/model/dgl_layers/bundler.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Bundler(nn.Module):
+    """
+    Bundler, which will be the node_apply function in DGL paradigm
+    """
+
+    def __init__(self, in_feats, out_feats, activation, dropout, bias=True):
+        super(Bundler, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        self.linear = nn.Linear(in_feats * 2, out_feats, bias)
+        self.activation = activation
+
+        nn.init.xavier_uniform_(
+            self.linear.weight, gain=nn.init.calculate_gain("relu")
+        )
+
+    def concat(self, h, aggre_result):
+        bundle = torch.cat((h, aggre_result), 1)
+        bundle = self.linear(bundle)
+        return bundle
+
+    def forward(self, node):
+        h = node.data["h"]
+        c = node.data["c"]
+        bundle = self.concat(h, c)
+        bundle = F.normalize(bundle, p=2, dim=1)
+        if self.activation:
+            bundle = self.activation(bundle)
+        return {"h": bundle}

examples/pytorch/diffpool/model/dgl_layers/gnn.py

@@ -0,0 +1,158 @@
+import dgl.function as fn
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from scipy.linalg import block_diag
+
+from model.loss import EntropyLoss
+from ..model_utils import masked_softmax
+
+from .aggregator import LSTMAggregator, MaxPoolAggregator, MeanAggregator
+from .bundler import Bundler
+
+
+class GraphSageLayer(nn.Module):
+    """
+    GraphSage layer in Inductive learning paper by hamilton
+    Here, graphsage layer is a reduced function in DGL framework
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        activation,
+        dropout,
+        aggregator_type,
+        bn=False,
+        bias=True,
+    ):
+        super(GraphSageLayer, self).__init__()
+        self.use_bn = bn
+        self.bundler = Bundler(
+            in_feats, out_feats, activation, dropout, bias=bias
+        )
+        self.dropout = nn.Dropout(p=dropout)
+
+        if aggregator_type == "maxpool":
+            self.aggregator = MaxPoolAggregator(
+                in_feats, in_feats, activation, bias
+            )
+        elif aggregator_type == "lstm":
+            self.aggregator = LSTMAggregator(in_feats, in_feats)
+        else:
+            self.aggregator = MeanAggregator()
+
+    def forward(self, g, h):
+        h = self.dropout(h)
+        g.ndata["h"] = h
+        if self.use_bn and not hasattr(self, "bn"):
+            device = h.device
+            self.bn = nn.BatchNorm1d(h.size()[1]).to(device)
+        g.update_all(fn.copy_u(u="h", out="m"), self.aggregator, self.bundler)
+        if self.use_bn:
+            h = self.bn(h)
+        h = g.ndata.pop("h")
+        return h
+
+
+class GraphSage(nn.Module):
+    """
+    Grahpsage network that concatenate several graphsage layer
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        n_hidden,
+        n_classes,
+        n_layers,
+        activation,
+        dropout,
+        aggregator_type,
+    ):
+        super(GraphSage, self).__init__()
+        self.layers = nn.ModuleList()
+
+        # input layer
+        self.layers.append(
+            GraphSageLayer(
+                in_feats, n_hidden, activation, dropout, aggregator_type
+            )
+        )
+        # hidden layers
+        for _ in range(n_layers - 1):
+            self.layers.append(
+                GraphSageLayer(
+                    n_hidden, n_hidden, activation, dropout, aggregator_type
+                )
+            )
+        # output layer
+        self.layers.append(
+            GraphSageLayer(n_hidden, n_classes, None, dropout, aggregator_type)
+        )
+
+    def forward(self, g, features):
+        h = features
+        for layer in self.layers:
+            h = layer(g, h)
+        return h
+
+
+class DiffPoolBatchedGraphLayer(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        assign_dim,
+        output_feat_dim,
+        activation,
+        dropout,
+        aggregator_type,
+        link_pred,
+    ):
+        super(DiffPoolBatchedGraphLayer, self).__init__()
+        self.embedding_dim = input_dim
+        self.assign_dim = assign_dim
+        self.hidden_dim = output_feat_dim
+        self.link_pred = link_pred
+        self.feat_gc = GraphSageLayer(
+            input_dim, output_feat_dim, activation, dropout, aggregator_type
+        )
+        self.pool_gc = GraphSageLayer(
+            input_dim, assign_dim, activation, dropout, aggregator_type
+        )
+        self.reg_loss = nn.ModuleList([])
+        self.loss_log = {}
+        self.reg_loss.append(EntropyLoss())
+
+    def forward(self, g, h):
+        feat = self.feat_gc(
+            g, h
+        )  # size = (sum_N, F_out), sum_N is num of nodes in this batch
+        device = feat.device
+        assign_tensor = self.pool_gc(
+            g, h
+        )  # size = (sum_N, N_a), N_a is num of nodes in pooled graph.
+        assign_tensor = F.softmax(assign_tensor, dim=1)
+        assign_tensor = torch.split(assign_tensor, g.batch_num_nodes().tolist())
+        assign_tensor = torch.block_diag(
+            *assign_tensor
+        )  # size = (sum_N, batch_size * N_a)
+
+        h = torch.matmul(torch.t(assign_tensor), feat)
+        adj = g.adj_external(transpose=True, ctx=device)
+        adj_new = torch.sparse.mm(adj, assign_tensor)
+        adj_new = torch.mm(torch.t(assign_tensor), adj_new)
+
+        if self.link_pred:
+            current_lp_loss = torch.norm(
+                adj.to_dense() - torch.mm(assign_tensor, torch.t(assign_tensor))
+            ) / np.power(g.num_nodes(), 2)
+            self.loss_log["LinkPredLoss"] = current_lp_loss
+
+        for loss_layer in self.reg_loss:
+            loss_name = str(type(loss_layer).__name__)
+            self.loss_log[loss_name] = loss_layer(adj, adj_new, assign_tensor)
+
+        return adj_new, h

examples/pytorch/diffpool/model/encoder.py

@@ -0,0 +1,243 @@
+import time
+
+import dgl
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from scipy.linalg import block_diag
+from torch.nn import init
+
+from .dgl_layers import DiffPoolBatchedGraphLayer, GraphSage, GraphSageLayer
+from .model_utils import batch2tensor
+from .tensorized_layers import *
+
+
+class DiffPool(nn.Module):
+    """
+    DiffPool Fuse
+    """
+
+    def __init__(
+        self,
+        input_dim,
+        hidden_dim,
+        embedding_dim,
+        label_dim,
+        activation,
+        n_layers,
+        dropout,
+        n_pooling,
+        linkpred,
+        batch_size,
+        aggregator_type,
+        assign_dim,
+        pool_ratio,
+        cat=False,
+    ):
+        super(DiffPool, self).__init__()
+        self.link_pred = linkpred
+        self.concat = cat
+        self.n_pooling = n_pooling
+        self.batch_size = batch_size
+        self.link_pred_loss = []
+        self.entropy_loss = []
+
+        # list of GNN modules before the first diffpool operation
+        self.gc_before_pool = nn.ModuleList()
+        self.diffpool_layers = nn.ModuleList()
+
+        # list of list of GNN modules, each list after one diffpool operation
+        self.gc_after_pool = nn.ModuleList()
+        self.assign_dim = assign_dim
+        self.bn = True
+        self.num_aggs = 1
+
+        # constructing layers
+        # layers before diffpool
+        assert n_layers >= 3, "n_layers too few"
+        self.gc_before_pool.append(
+            GraphSageLayer(
+                input_dim,
+                hidden_dim,
+                activation,
+                dropout,
+                aggregator_type,
+                self.bn,
+            )
+        )
+        for _ in range(n_layers - 2):
+            self.gc_before_pool.append(
+                GraphSageLayer(
+                    hidden_dim,
+                    hidden_dim,
+                    activation,
+                    dropout,
+                    aggregator_type,
+                    self.bn,
+                )
+            )
+        self.gc_before_pool.append(
+            GraphSageLayer(
+                hidden_dim, embedding_dim, None, dropout, aggregator_type
+            )
+        )
+
+        assign_dims = []
+        assign_dims.append(self.assign_dim)
+        if self.concat:
+            # diffpool layer receive pool_emedding_dim node feature tensor
+            # and return pool_embedding_dim node embedding
+            pool_embedding_dim = hidden_dim * (n_layers - 1) + embedding_dim
+        else:
+            pool_embedding_dim = embedding_dim
+
+        self.first_diffpool_layer = DiffPoolBatchedGraphLayer(
+            pool_embedding_dim,
+            self.assign_dim,
+            hidden_dim,
+            activation,
+            dropout,
+            aggregator_type,
+            self.link_pred,
+        )
+        gc_after_per_pool = nn.ModuleList()
+
+        for _ in range(n_layers - 1):
+            gc_after_per_pool.append(BatchedGraphSAGE(hidden_dim, hidden_dim))
+        gc_after_per_pool.append(BatchedGraphSAGE(hidden_dim, embedding_dim))
+        self.gc_after_pool.append(gc_after_per_pool)
+
+        self.assign_dim = int(self.assign_dim * pool_ratio)
+        # each pooling module
+        for _ in range(n_pooling - 1):
+            self.diffpool_layers.append(
+                BatchedDiffPool(
+                    pool_embedding_dim,
+                    self.assign_dim,
+                    hidden_dim,
+                    self.link_pred,
+                )
+            )
+            gc_after_per_pool = nn.ModuleList()
+            for _ in range(n_layers - 1):
+                gc_after_per_pool.append(
+                    BatchedGraphSAGE(hidden_dim, hidden_dim)
+                )
+            gc_after_per_pool.append(
+                BatchedGraphSAGE(hidden_dim, embedding_dim)
+            )
+            self.gc_after_pool.append(gc_after_per_pool)
+            assign_dims.append(self.assign_dim)
+            self.assign_dim = int(self.assign_dim * pool_ratio)
+
+        # predicting layer
+        if self.concat:
+            self.pred_input_dim = (
+                pool_embedding_dim * self.num_aggs * (n_pooling + 1)
+            )
+        else:
+            self.pred_input_dim = embedding_dim * self.num_aggs
+        self.pred_layer = nn.Linear(self.pred_input_dim, label_dim)
+
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                m.weight.data = init.xavier_uniform_(
+                    m.weight.data, gain=nn.init.calculate_gain("relu")
+                )
+                if m.bias is not None:
+                    m.bias.data = init.constant_(m.bias.data, 0.0)
+
+    def gcn_forward(self, g, h, gc_layers, cat=False):
+        """
+        Return gc_layer embedding cat.
+        """
+        block_readout = []
+        for gc_layer in gc_layers[:-1]:
+            h = gc_layer(g, h)
+            block_readout.append(h)
+        h = gc_layers[-1](g, h)
+        block_readout.append(h)
+        if cat:
+            block = torch.cat(block_readout, dim=1)  # N x F, F = F1 + F2 + ...
+        else:
+            block = h
+        return block
+
+    def gcn_forward_tensorized(self, h, adj, gc_layers, cat=False):
+        block_readout = []
+        for gc_layer in gc_layers:
+            h = gc_layer(h, adj)
+            block_readout.append(h)
+        if cat:
+            block = torch.cat(block_readout, dim=2)  # N x F, F = F1 + F2 + ...
+        else:
+            block = h
+        return block
+
+    def forward(self, g):
+        self.link_pred_loss = []
+        self.entropy_loss = []
+        h = g.ndata["feat"]
+        # node feature for assignment matrix computation is the same as the
+        # original node feature
+        h_a = h
+
+        out_all = []
+
+        # we use GCN blocks to get an embedding first
+        g_embedding = self.gcn_forward(g, h, self.gc_before_pool, self.concat)
+
+        g.ndata["h"] = g_embedding
+
+        readout = dgl.sum_nodes(g, "h")
+        out_all.append(readout)
+        if self.num_aggs == 2:
+            readout = dgl.max_nodes(g, "h")
+            out_all.append(readout)
+
+        adj, h = self.first_diffpool_layer(g, g_embedding)
+        node_per_pool_graph = int(adj.size()[0] / len(g.batch_num_nodes()))
+
+        h, adj = batch2tensor(adj, h, node_per_pool_graph)
+        h = self.gcn_forward_tensorized(
+            h, adj, self.gc_after_pool[0], self.concat
+        )
+        readout = torch.sum(h, dim=1)
+        out_all.append(readout)
+        if self.num_aggs == 2:
+            readout, _ = torch.max(h, dim=1)
+            out_all.append(readout)
+
+        for i, diffpool_layer in enumerate(self.diffpool_layers):
+            h, adj = diffpool_layer(h, adj)
+            h = self.gcn_forward_tensorized(
+                h, adj, self.gc_after_pool[i + 1], self.concat
+            )
+            readout = torch.sum(h, dim=1)
+            out_all.append(readout)
+            if self.num_aggs == 2:
+                readout, _ = torch.max(h, dim=1)
+                out_all.append(readout)
+        if self.concat or self.num_aggs > 1:
+            final_readout = torch.cat(out_all, dim=1)
+        else:
+            final_readout = readout
+        ypred = self.pred_layer(final_readout)
+        return ypred
+
+    def loss(self, pred, label):
+        """
+        loss function
+        """
+        # softmax + CE
+        criterion = nn.CrossEntropyLoss()
+        loss = criterion(pred, label)
+        for key, value in self.first_diffpool_layer.loss_log.items():
+            loss += value
+        for diffpool_layer in self.diffpool_layers:
+            for key, value in diffpool_layer.loss_log.items():
+                loss += value
+        return loss

examples/pytorch/diffpool/model/loss.py

@@ -0,0 +1,23 @@
+import torch
+import torch.nn as nn
+
+
+class EntropyLoss(nn.Module):
+    # Return Scalar
+    def forward(self, adj, anext, s_l):
+        entropy = (
+            (torch.distributions.Categorical(probs=s_l).entropy())
+            .sum(-1)
+            .mean(-1)
+        )
+        assert not torch.isnan(entropy)
+        return entropy
+
+
+class LinkPredLoss(nn.Module):
+    def forward(self, adj, anext, s_l):
+        link_pred_loss = (adj - s_l.matmul(s_l.transpose(-1, -2))).norm(
+            dim=(1, 2)
+        )
+        link_pred_loss = link_pred_loss / (adj.size(1) * adj.size(2))
+        return link_pred_loss.mean()

examples/pytorch/diffpool/model/model_utils.py

@@ -0,0 +1,47 @@
+import torch as th
+from torch.autograd import Function
+
+
+def batch2tensor(batch_adj, batch_feat, node_per_pool_graph):
+    """
+    transform a batched graph to batched adjacency tensor and node feature tensor
+    """
+    batch_size = int(batch_adj.size()[0] / node_per_pool_graph)
+    adj_list = []
+    feat_list = []
+    for i in range(batch_size):
+        start = i * node_per_pool_graph
+        end = (i + 1) * node_per_pool_graph
+        adj_list.append(batch_adj[start:end, start:end])
+        feat_list.append(batch_feat[start:end, :])
+    adj_list = list(map(lambda x: th.unsqueeze(x, 0), adj_list))
+    feat_list = list(map(lambda x: th.unsqueeze(x, 0), feat_list))
+    adj = th.cat(adj_list, dim=0)
+    feat = th.cat(feat_list, dim=0)
+
+    return feat, adj
+
+
+def masked_softmax(
+    matrix, mask, dim=-1, memory_efficient=True, mask_fill_value=-1e32
+):
+    """
+    masked_softmax for dgl batch graph
+    code snippet contributed by AllenNLP (https://github.com/allenai/allennlp)
+    """
+    if mask is None:
+        result = th.nn.functional.softmax(matrix, dim=dim)
+    else:
+        mask = mask.float()
+        while mask.dim() < matrix.dim():
+            mask = mask.unsqueeze(1)
+        if not memory_efficient:
+            result = th.nn.functional.softmax(matrix * mask, dim=dim)
+            result = result * mask
+            result = result / (result.sum(dim=dim, keepdim=True) + 1e-13)
+        else:
+            masked_matrix = matrix.masked_fill(
+                (1 - mask).byte(), mask_fill_value
+            )
+            result = th.nn.functional.softmax(masked_matrix, dim=dim)
+    return result

examples/pytorch/diffpool/model/tensorized_layers/__init__.py

@@ -0,0 +1,2 @@
+from .diffpool import BatchedDiffPool
+from .graphsage import BatchedGraphSAGE

examples/pytorch/diffpool/model/tensorized_layers/assignment.py

@@ -0,0 +1,17 @@
+import torch
+from torch import nn as nn
+from torch.autograd import Variable
+from torch.nn import functional as F
+
+from model.tensorized_layers.graphsage import BatchedGraphSAGE
+
+
+class DiffPoolAssignment(nn.Module):
+    def __init__(self, nfeat, nnext):
+        super().__init__()
+        self.assign_mat = BatchedGraphSAGE(nfeat, nnext, use_bn=True)
+
+    def forward(self, x, adj, log=False):
+        s_l_init = self.assign_mat(x, adj)
+        s_l = F.softmax(s_l_init, dim=-1)
+        return s_l

examples/pytorch/diffpool/model/tensorized_layers/diffpool.py

@@ -0,0 +1,37 @@
+import torch
+from torch import nn as nn
+
+from model.loss import EntropyLoss, LinkPredLoss
+from model.tensorized_layers.assignment import DiffPoolAssignment
+from model.tensorized_layers.graphsage import BatchedGraphSAGE
+
+
+class BatchedDiffPool(nn.Module):
+    def __init__(self, nfeat, nnext, nhid, link_pred=False, entropy=True):
+        super(BatchedDiffPool, self).__init__()
+        self.link_pred = link_pred
+        self.log = {}
+        self.link_pred_layer = LinkPredLoss()
+        self.embed = BatchedGraphSAGE(nfeat, nhid, use_bn=True)
+        self.assign = DiffPoolAssignment(nfeat, nnext)
+        self.reg_loss = nn.ModuleList([])
+        self.loss_log = {}
+        if link_pred:
+            self.reg_loss.append(LinkPredLoss())
+        if entropy:
+            self.reg_loss.append(EntropyLoss())
+
+    def forward(self, x, adj, log=False):
+        z_l = self.embed(x, adj)
+        s_l = self.assign(x, adj)
+        if log:
+            self.log["s"] = s_l.cpu().numpy()
+        xnext = torch.matmul(s_l.transpose(-1, -2), z_l)
+        anext = (s_l.transpose(-1, -2)).matmul(adj).matmul(s_l)
+
+        for loss_layer in self.reg_loss:
+            loss_name = str(type(loss_layer).__name__)
+            self.loss_log[loss_name] = loss_layer(adj, anext, s_l)
+        if log:
+            self.log["a"] = anext.cpu().numpy()
+        return xnext, anext

examples/pytorch/diffpool/model/tensorized_layers/graphsage.py

@@ -0,0 +1,43 @@
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+class BatchedGraphSAGE(nn.Module):
+    def __init__(
+        self, infeat, outfeat, use_bn=True, mean=False, add_self=False
+    ):
+        super().__init__()
+        self.add_self = add_self
+        self.use_bn = use_bn
+        self.mean = mean
+        self.W = nn.Linear(infeat, outfeat, bias=True)
+
+        nn.init.xavier_uniform_(
+            self.W.weight, gain=nn.init.calculate_gain("relu")
+        )
+
+    def forward(self, x, adj):
+        num_node_per_graph = adj.size(1)
+        if self.use_bn and not hasattr(self, "bn"):
+            self.bn = nn.BatchNorm1d(num_node_per_graph).to(adj.device)
+
+        if self.add_self:
+            adj = adj + torch.eye(num_node_per_graph).to(adj.device)
+
+        if self.mean:
+            adj = adj / adj.sum(-1, keepdim=True)
+
+        h_k_N = torch.matmul(adj, x)
+        h_k = self.W(h_k_N)
+        h_k = F.normalize(h_k, dim=2, p=2)
+        h_k = F.relu(h_k)
+        if self.use_bn:
+            h_k = self.bn(h_k)
+        return h_k
+
+    def __repr__(self):
+        if self.use_bn:
+            return "BN" + super(BatchedGraphSAGE, self).__repr__()
+        else:
+            return super(BatchedGraphSAGE, self).__repr__()

examples/pytorch/diffpool/train.py

@@ -0,0 +1,381 @@
+import argparse
+import os
+import random
+import time
+
+import dgl
+import dgl.function as fn
+
+import networkx as nx
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data
+from data_utils import pre_process
+from dgl import DGLGraph
+from dgl.data import tu
+from model.encoder import DiffPool
+
+global_train_time_per_epoch = []
+
+
+def arg_parse():
+    """
+    argument parser
+    """
+    parser = argparse.ArgumentParser(description="DiffPool arguments")
+    parser.add_argument("--dataset", dest="dataset", help="Input Dataset")
+    parser.add_argument(
+        "--pool_ratio", dest="pool_ratio", type=float, help="pooling ratio"
+    )
+    parser.add_argument(
+        "--num_pool", dest="num_pool", type=int, help="num_pooling layer"
+    )
+    parser.add_argument(
+        "--no_link_pred",
+        dest="linkpred",
+        action="store_false",
+        help="switch of link prediction object",
+    )
+    parser.add_argument("--cuda", dest="cuda", type=int, help="switch cuda")
+    parser.add_argument("--lr", dest="lr", type=float, help="learning rate")
+    parser.add_argument(
+        "--clip", dest="clip", type=float, help="gradient clipping"
+    )
+    parser.add_argument(
+        "--batch-size", dest="batch_size", type=int, help="batch size"
+    )
+    parser.add_argument("--epochs", dest="epoch", type=int, help="num-of-epoch")
+    parser.add_argument(
+        "--train-ratio",
+        dest="train_ratio",
+        type=float,
+        help="ratio of trainning dataset split",
+    )
+    parser.add_argument(
+        "--test-ratio",
+        dest="test_ratio",
+        type=float,
+        help="ratio of testing dataset split",
+    )
+    parser.add_argument(
+        "--num_workers",
+        dest="n_worker",
+        type=int,
+        help="number of workers when dataloading",
+    )
+    parser.add_argument(
+        "--gc-per-block",
+        dest="gc_per_block",
+        type=int,
+        help="number of graph conv layer per block",
+    )
+    parser.add_argument(
+        "--bn",
+        dest="bn",
+        action="store_const",
+        const=True,
+        default=True,
+        help="switch for bn",
+    )
+    parser.add_argument(
+        "--dropout", dest="dropout", type=float, help="dropout rate"
+    )
+    parser.add_argument(
+        "--bias",
+        dest="bias",
+        action="store_const",
+        const=True,
+        default=True,
+        help="switch for bias",
+    )
+    parser.add_argument(
+        "--save_dir",
+        dest="save_dir",
+        help="model saving directory: SAVE_DICT/DATASET",
+    )
+    parser.add_argument(
+        "--load_epoch",
+        dest="load_epoch",
+        type=int,
+        help="load trained model params from\
+                         SAVE_DICT/DATASET/model-LOAD_EPOCH",
+    )
+    parser.add_argument(
+        "--data_mode",
+        dest="data_mode",
+        help="data\
+                        preprocessing mode: default, id, degree, or one-hot\
+                        vector of degree number",
+        choices=["default", "id", "deg", "deg_num"],
+    )
+
+    parser.set_defaults(
+        dataset="ENZYMES",
+        pool_ratio=0.15,
+        num_pool=1,
+        cuda=1,
+        lr=1e-3,
+        clip=2.0,
+        batch_size=20,
+        epoch=4000,
+        train_ratio=0.7,
+        test_ratio=0.1,
+        n_worker=1,
+        gc_per_block=3,
+        dropout=0.0,
+        method="diffpool",
+        bn=True,
+        bias=True,
+        save_dir="./model_param",
+        load_epoch=-1,
+        data_mode="default",
+    )
+    return parser.parse_args()
+
+
+def prepare_data(dataset, prog_args, train=False, pre_process=None):
+    """
+    preprocess TU dataset according to DiffPool's paper setting and load dataset into dataloader
+    """
+    if train:
+        shuffle = True
+    else:
+        shuffle = False
+
+    if pre_process:
+        pre_process(dataset, prog_args)
+
+    # dataset.set_fold(fold)
+    return dgl.dataloading.GraphDataLoader(
+        dataset,
+        batch_size=prog_args.batch_size,
+        shuffle=shuffle,
+        num_workers=prog_args.n_worker,
+    )
+
+
+def graph_classify_task(prog_args):
+    """
+    perform graph classification task
+    """
+
+    dataset = tu.LegacyTUDataset(name=prog_args.dataset)
+    train_size = int(prog_args.train_ratio * len(dataset))
+    test_size = int(prog_args.test_ratio * len(dataset))
+    val_size = int(len(dataset) - train_size - test_size)
+
+    dataset_train, dataset_val, dataset_test = torch.utils.data.random_split(
+        dataset, (train_size, val_size, test_size)
+    )
+    train_dataloader = prepare_data(
+        dataset_train, prog_args, train=True, pre_process=pre_process
+    )
+    val_dataloader = prepare_data(
+        dataset_val, prog_args, train=False, pre_process=pre_process
+    )
+    test_dataloader = prepare_data(
+        dataset_test, prog_args, train=False, pre_process=pre_process
+    )
+    input_dim, label_dim, max_num_node = dataset.statistics()
+    print("++++++++++STATISTICS ABOUT THE DATASET")
+    print("dataset feature dimension is", input_dim)
+    print("dataset label dimension is", label_dim)
+    print("the max num node is", max_num_node)
+    print("number of graphs is", len(dataset))
+    # assert len(dataset) % prog_args.batch_size == 0, "training set not divisible by batch size"
+
+    hidden_dim = 64  # used to be 64
+    embedding_dim = 64
+
+    # calculate assignment dimension: pool_ratio * largest graph's maximum
+    # number of nodes  in the dataset
+    assign_dim = int(max_num_node * prog_args.pool_ratio)
+    print("++++++++++MODEL STATISTICS++++++++")
+    print("model hidden dim is", hidden_dim)
+    print("model embedding dim for graph instance embedding", embedding_dim)
+    print("initial batched pool graph dim is", assign_dim)
+    activation = F.relu
+
+    # initialize model
+    # 'diffpool' : diffpool
+    model = DiffPool(
+        input_dim,
+        hidden_dim,
+        embedding_dim,
+        label_dim,
+        activation,
+        prog_args.gc_per_block,
+        prog_args.dropout,
+        prog_args.num_pool,
+        prog_args.linkpred,
+        prog_args.batch_size,
+        "meanpool",
+        assign_dim,
+        prog_args.pool_ratio,
+    )
+
+    if prog_args.load_epoch >= 0 and prog_args.save_dir is not None:
+        model.load_state_dict(
+            torch.load(
+                prog_args.save_dir
+                + "/"
+                + prog_args.dataset
+                + "/model.iter-"
+                + str(prog_args.load_epoch),
+                weights_only=False,
+            )
+        )
+
+    print("model init finished")
+    print("MODEL:::::::", prog_args.method)
+    if prog_args.cuda:
+        model = model.cuda()
+
+    logger = train(
+        train_dataloader, model, prog_args, val_dataset=val_dataloader
+    )
+    result = evaluate(test_dataloader, model, prog_args, logger)
+    print("test  accuracy {:.2f}%".format(result * 100))
+
+
+def train(dataset, model, prog_args, same_feat=True, val_dataset=None):
+    """
+    training function
+    """
+    dir = prog_args.save_dir + "/" + prog_args.dataset
+    if not os.path.exists(dir):
+        os.makedirs(dir)
+    dataloader = dataset
+    optimizer = torch.optim.Adam(
+        filter(lambda p: p.requires_grad, model.parameters()), lr=0.001
+    )
+    early_stopping_logger = {"best_epoch": -1, "val_acc": -1}
+
+    if prog_args.cuda > 0:
+        torch.cuda.set_device(0)
+    for epoch in range(prog_args.epoch):
+        begin_time = time.time()
+        model.train()
+        accum_correct = 0
+        total = 0
+        print("\nEPOCH ###### {} ######".format(epoch))
+        computation_time = 0.0
+        for batch_idx, (batch_graph, graph_labels) in enumerate(dataloader):
+            for key, value in batch_graph.ndata.items():
+                batch_graph.ndata[key] = value.float()
+            graph_labels = graph_labels.long()
+            if torch.cuda.is_available():
+                batch_graph = batch_graph.to(torch.cuda.current_device())
+                graph_labels = graph_labels.cuda()
+
+            model.zero_grad()
+            compute_start = time.time()
+            ypred = model(batch_graph)
+            indi = torch.argmax(ypred, dim=1)
+            correct = torch.sum(indi == graph_labels).item()
+            accum_correct += correct
+            total += graph_labels.size()[0]
+            loss = model.loss(ypred, graph_labels)
+            loss.backward()
+            batch_compute_time = time.time() - compute_start
+            computation_time += batch_compute_time
+            nn.utils.clip_grad_norm_(model.parameters(), prog_args.clip)
+            optimizer.step()
+
+        train_accu = accum_correct / total
+        print(
+            "train accuracy for this epoch {} is {:.2f}%".format(
+                epoch, train_accu * 100
+            )
+        )
+        elapsed_time = time.time() - begin_time
+        print(
+            "loss {:.4f} with epoch time {:.4f} s & computation time {:.4f} s ".format(
+                loss.item(), elapsed_time, computation_time
+            )
+        )
+        global_train_time_per_epoch.append(elapsed_time)
+        if val_dataset is not None:
+            result = evaluate(val_dataset, model, prog_args)
+            print("validation  accuracy {:.2f}%".format(result * 100))
+            if (
+                result >= early_stopping_logger["val_acc"]
+                and result <= train_accu
+            ):
+                early_stopping_logger.update(best_epoch=epoch, val_acc=result)
+                if prog_args.save_dir is not None:
+                    torch.save(
+                        model.state_dict(),
+                        prog_args.save_dir
+                        + "/"
+                        + prog_args.dataset
+                        + "/model.iter-"
+                        + str(early_stopping_logger["best_epoch"]),
+                    )
+            print(
+                "best epoch is EPOCH {}, val_acc is {:.2f}%".format(
+                    early_stopping_logger["best_epoch"],
+                    early_stopping_logger["val_acc"] * 100,
+                )
+            )
+        torch.cuda.empty_cache()
+    return early_stopping_logger
+
+
+def evaluate(dataloader, model, prog_args, logger=None):
+    """
+    evaluate function
+    """
+    if logger is not None and prog_args.save_dir is not None:
+        model.load_state_dict(
+            torch.load(
+                prog_args.save_dir
+                + "/"
+                + prog_args.dataset
+                + "/model.iter-"
+                + str(logger["best_epoch"]),
+                weights_only=False,
+            )
+        )
+    model.eval()
+    correct_label = 0
+    with torch.no_grad():
+        for batch_idx, (batch_graph, graph_labels) in enumerate(dataloader):
+            for key, value in batch_graph.ndata.items():
+                batch_graph.ndata[key] = value.float()
+            graph_labels = graph_labels.long()
+            if torch.cuda.is_available():
+                batch_graph = batch_graph.to(torch.cuda.current_device())
+                graph_labels = graph_labels.cuda()
+            ypred = model(batch_graph)
+            indi = torch.argmax(ypred, dim=1)
+            correct = torch.sum(indi == graph_labels)
+            correct_label += correct.item()
+    result = correct_label / (len(dataloader) * prog_args.batch_size)
+    return result
+
+
+def main():
+    """
+    main
+    """
+    prog_args = arg_parse()
+    print(prog_args)
+    graph_classify_task(prog_args)
+
+    print(
+        "Train time per epoch: {:.4f}".format(
+            sum(global_train_time_per_epoch) / len(global_train_time_per_epoch)
+        )
+    )
+    print(
+        "Max memory usage: {:.4f}".format(
+            torch.cuda.max_memory_allocated(0) / (1024 * 1024)
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()