chore: import upstream snapshot with attribution

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__()