add BGNN_Adv and BGNN_MLP

Author: yifanfeng97

dhg/models/__init__.py

@@ -1,4 +1,4 @@
-from .graphs import GCN, GraphSAGE, GAT, GIN, NGCF, LightGCN, BGNN_Adv
+from .graphs import GCN, GraphSAGE, GAT, GIN, NGCF, LightGCN, BGNN_Adv, BGNN_MLP
 from .hypergraphs import HGNN, HGNNP, HNHN, HyperGCN, DHCF, UniGCN, UniGAT, UniSAGE, UniGIN
 
 __all__ = [
@@ -8,6 +8,8 @@
     "GIN",
     "NGCF",
     "LightGCN",
+    "BGNN_Adv",
+    "BGNN_MLP",
     "HGNN",
     "HGNNP",
     "HNHN",

dhg/models/graphs/__init__.py

@@ -4,4 +4,4 @@
 from .gin import GIN
 from .ngcf import NGCF
 from .lightgcn import LightGCN
-from .bgnn import BGNN_Adv
+from .bgnn import BGNN_Adv, BGNN_MLP

dhg/models/graphs/bgnn.py

@@ -10,48 +10,89 @@ class BGNN_Adv(nn.Module):
     r"""The BGNN-Adv model proposed in `Cascade-BGNN: Toward Efficient Self-supervised Representation Learning on Large-scale Bipartite Graphs <https://arxiv.org/pdf/1906.11994.pdf>`_ paper (TNNLS 2020).
 
     Args:
-        ``num_users`` (``int``): The Number of users.
-        ``num_items`` (``int``): The Number of items.
-        ``emb_dim`` (``int``): Embedding dimension.
-        ``num_layers`` (``int``): The Number of layers. Defaults to ``3``.
-        ``drop_rate`` (``float``): Dropout rate. Randomly dropout the connections in training stage with probability ``drop_rate``. Default: ``0.0``.
+        ``u_dim`` (``int``): The dimension of the vertex feature in set :math:`U`.
+        ``v_dim`` (``int``): The dimension of the vertex feature in set :math:`V`.
+        ``layer_depth`` (``int``): The depth of layers.
     """
 
-    def __init__(
-        self, num_users: int, num_items: int, emb_dim: int, num_layers: int = 3, drop_rate: float = 0.0
-    ) -> None:
+    def __init__(self, u_dim: int, v_dim: int, layer_depth: int = 3,) -> None:
 
         super().__init__()
-        self.num_users, self.num_items = num_users, num_items
-        self.num_layers = num_layers
-        self.drop_rate = drop_rate
-        self.u_embedding = nn.Embedding(num_users, emb_dim)
-        self.i_embedding = nn.Embedding(num_items, emb_dim)
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        r"""Initialize learnable parameters.
+        self.layer_depth = layer_depth
+        self.layers = nn.ModuleList()
+
+        for _idx in range(layer_depth):
+            if _idx % 2 == 0:
+                self.layers.append(nn.Linear(v_dim, u_dim))
+            else:
+                self.layers.append(nn.Linear(u_dim, v_dim))
+
+    def forward(self, X_u: torch.Tensor, X_v: torch.Tensor, g: BiGraph) -> Tuple[torch.Tensor, torch.Tensor]:
+        r"""The forward function.
+
+        Args:
+            ``X_u`` (``torch.Tensor``): The feature matrix of vertices in set :math:`U`.
+            ``X_v`` (``torch.Tensor``): The feature matrix of vertices in set :math:`V`.
+            ``g`` (``BiGraph``): The bipartite graph.
         """
-        nn.init.normal_(self.u_embedding.weight, 0, 0.1)
-        nn.init.normal_(self.i_embedding.weight, 0, 0.1)
+        last_X_u, last_X_v = X_u, X_v
+        for _idx in range(self.layer_depth):
+            if _idx % 2 == 0:
+                _tmp = self.layers[_idx](last_X_v)
+                last_X_u = g.v2u(_tmp, aggr="sum")
+            else:
+                _tmp = self.layers[_idx](last_X_u)
+                last_X_v = g.u2v(_tmp, aggr="sum")
+        return last_X_u
+
+    def train_with_cascaded(self):
+        pass
+
+    def train_with_end2end(self):
+        pass
+
+
+class BGNN_MLP(nn.Module):
+    r"""The BGNN-MLP model proposed in `Cascade-BGNN: Toward Efficient Self-supervised Representation Learning on Large-scale Bipartite Graphs <https://arxiv.org/pdf/1906.11994.pdf>`_ paper (TNNLS 2020).
 
-    def forward(self, ui_bigraph: BiGraph) -> Tuple[torch.Tensor, torch.Tensor]:
+    Args:
+        ``u_dim`` (``int``): The dimension of the vertex feature in set :math:`U`.
+        ``v_dim`` (``int``): The dimension of the vertex feature in set :math:`V`.
+        ``layer_depth`` (``int``): The depth of layers.
+    """
+
+    def __init__(self, u_dim: int, v_dim: int, layer_depth: int = 3,) -> None:
+
+        super().__init__()
+        self.layer_depth = layer_depth
+        self.layers = nn.ModuleList()
+
+        for _idx in range(layer_depth):
+            if _idx % 2 == 0:
+                self.layers.append(nn.Linear(v_dim, u_dim))
+            else:
+                self.layers.append(nn.Linear(u_dim, v_dim))
+
+    def forward(self, X_u: torch.Tensor, X_v: torch.Tensor, g: BiGraph) -> Tuple[torch.Tensor, torch.Tensor]:
         r"""The forward function.
 
         Args:
-            ``ui_bigraph`` (``dhg.BiGraph``): The user-item bipartite graph.
+            ``X_u`` (``torch.Tensor``): The feature matrix of vertices in set :math:`U`.
+            ``X_v`` (``torch.Tensor``): The feature matrix of vertices in set :math:`V`.
+            ``g`` (``BiGraph``): The bipartite graph.
         """
-        drop_rate = self.drop_rate if self.training else 0.0
-        u_embs = self.u_embedding.weight
-        i_embs = self.i_embedding.weight
-        all_embs = torch.cat([u_embs, i_embs], dim=0)
-
-        embs_list = [all_embs]
-        for _ in range(self.num_layers):
-            all_embs = ui_bigraph.smoothing_with_GCN(all_embs, drop_rate=drop_rate)
-            embs_list.append(all_embs)
-        embs = torch.stack(embs_list, dim=1)
-        embs = torch.mean(embs, dim=1)
-
-        u_embs, i_embs = torch.split(embs, [self.num_users, self.num_items], dim=0)
-        return u_embs, i_embs
+        last_X_u, last_X_v = X_u, X_v
+        for _idx in range(self.layer_depth):
+            if _idx % 2 == 0:
+                _tmp = self.layers[_idx](last_X_v)
+                last_X_u = g.v2u(_tmp, aggr="sum")
+            else:
+                _tmp = self.layers[_idx](last_X_u)
+                last_X_v = g.u2v(_tmp, aggr="sum")
+        return last_X_u
+
+    def train_with_cascaded(self):
+        pass
+
+    def train_with_end2end(self):
+        pass

dhg/structure/graphs/bipartite_graph.py

@@ -714,12 +714,12 @@ def v2u(
                 P = self.B
             # message passing
             if aggr == "mean":
-                X = torch.sparse.mm(self.B, X)
+                X = torch.sparse.mm(P, X)
                 X = torch.sparse.mm(self.D_u_neg_1, X)
             elif aggr == "sum":
-                X = torch.sparse.mm(self.B, X)
+                X = torch.sparse.mm(P, X)
             elif aggr == "softmax_then_sum":
-                P = torch.sparse.softmax(self.B, dim=1)
+                P = torch.sparse.softmax(P, dim=1)
                 X = torch.sparse.mm(P, X)
             else:
                 pass

docs/source/api/models.rst

@@ -25,6 +25,7 @@ Models on Bipartite Graph
    dhg.models.NGCF
    dhg.models.LightGCN
    dhg.models.BGNN_Adv
+   dhg.models.BGNN_MLP
 
 
 Models on Hypergraph