| """ |
| Model architectures for TopoHyper and baselines. |
| |
| Implements: |
| - TopoHyperNet: Full hybrid architecture |
| - GCNNet, GATNet: Standard GNN baselines |
| - HGNNNet: Hypergraph-only baseline |
| - SimplicialNet: Simplicial-only baseline |
| - SimpleHybrid: Naive concatenation hybrid |
| - get_model(): Factory function |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from .layers import (TopoHyperConv, TopoHyperPool, SimplicialConv, |
| HypergraphConv, GCNLayer, GATLayer) |
|
|
|
|
| class TopoHyperNet(nn.Module): |
| """ |
| Full TopoHyper architecture. |
| |
| Architecture: |
| 1. Input projection |
| 2. N x TopoHyperConv layers (3-phase: simplicial + hypergraph + fusion) |
| 3. TopoHyperPool (mean + max + attention readout) |
| 4. Classification head |
| |
| Args: |
| in_dim: Input feature dimension |
| hidden_dim: Hidden layer dimension |
| num_classes: Number of output classes |
| num_layers: Number of TopoHyperConv layers |
| dropout: Dropout rate |
| use_bridge: Whether to use bridge matrix in fusion |
| use_attention: Whether to use attention-gated fusion |
| """ |
| |
| def __init__(self, in_dim, hidden_dim, num_classes, num_layers=2, |
| dropout=0.3, use_bridge=True, use_attention=True): |
| super().__init__() |
| self.input_proj = nn.Linear(in_dim, hidden_dim) |
| |
| self.convs = nn.ModuleList() |
| for i in range(num_layers): |
| self.convs.append(TopoHyperConv( |
| hidden_dim, hidden_dim, |
| use_bridge=use_bridge, |
| use_attention=use_attention |
| )) |
| |
| self.pool = TopoHyperPool(hidden_dim) |
| self.classifier = nn.Sequential( |
| nn.Linear(hidden_dim * 3, hidden_dim), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, num_classes) |
| ) |
| self.dropout = dropout |
| |
| def forward(self, x, sc, hg, edge_index=None): |
| x = F.relu(self.input_proj(x)) |
| x = F.dropout(x, p=self.dropout, training=self.training) |
| |
| for conv in self.convs: |
| x = conv(x, sc, hg) |
| x = F.dropout(x, p=self.dropout, training=self.training) |
| |
| g = self.pool(x) |
| return self.classifier(g) |
|
|
|
|
| class GCNNet(nn.Module): |
| """GCN baseline.""" |
| |
| def __init__(self, in_dim, hidden_dim, num_classes, num_layers=2, dropout=0.3): |
| super().__init__() |
| self.layers = nn.ModuleList() |
| self.layers.append(GCNLayer(in_dim, hidden_dim)) |
| for _ in range(num_layers - 1): |
| self.layers.append(GCNLayer(hidden_dim, hidden_dim)) |
| |
| self.classifier = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, num_classes) |
| ) |
| self.dropout = dropout |
| |
| def forward(self, x, sc=None, hg=None, edge_index=None): |
| for layer in self.layers: |
| x = F.relu(layer(x, edge_index)) |
| x = F.dropout(x, p=self.dropout, training=self.training) |
| g = x.mean(dim=0) |
| return self.classifier(g) |
|
|
|
|
| class GATNet(nn.Module): |
| """GAT baseline.""" |
| |
| def __init__(self, in_dim, hidden_dim, num_classes, num_layers=2, dropout=0.3): |
| super().__init__() |
| self.layers = nn.ModuleList() |
| self.layers.append(GATLayer(in_dim, hidden_dim)) |
| for _ in range(num_layers - 1): |
| self.layers.append(GATLayer(hidden_dim, hidden_dim)) |
| |
| self.classifier = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, num_classes) |
| ) |
| self.dropout = dropout |
| |
| def forward(self, x, sc=None, hg=None, edge_index=None): |
| for layer in self.layers: |
| x = F.relu(layer(x, edge_index)) |
| x = F.dropout(x, p=self.dropout, training=self.training) |
| g = x.mean(dim=0) |
| return self.classifier(g) |
|
|
|
|
| class HGNNNet(nn.Module): |
| """Hypergraph-only baseline.""" |
| |
| def __init__(self, in_dim, hidden_dim, num_classes, num_layers=2, dropout=0.3): |
| super().__init__() |
| self.layers = nn.ModuleList() |
| self.layers.append(HypergraphConv(in_dim, hidden_dim)) |
| for _ in range(num_layers - 1): |
| self.layers.append(HypergraphConv(hidden_dim, hidden_dim)) |
| |
| self.classifier = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, num_classes) |
| ) |
| self.dropout = dropout |
| |
| def forward(self, x, sc=None, hg=None, edge_index=None): |
| for layer in self.layers: |
| x = F.relu(layer(x, hg)) |
| x = F.dropout(x, p=self.dropout, training=self.training) |
| g = x.mean(dim=0) |
| return self.classifier(g) |
|
|
|
|
| class SimplicialNet(nn.Module): |
| """Simplicial-only baseline.""" |
| |
| def __init__(self, in_dim, hidden_dim, num_classes, num_layers=2, dropout=0.3): |
| super().__init__() |
| self.layers = nn.ModuleList() |
| self.layers.append(SimplicialConv(in_dim, hidden_dim)) |
| for _ in range(num_layers - 1): |
| self.layers.append(SimplicialConv(hidden_dim, hidden_dim)) |
| |
| self.classifier = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, num_classes) |
| ) |
| self.dropout = dropout |
| |
| def forward(self, x, sc=None, hg=None, edge_index=None): |
| for layer in self.layers: |
| x = F.relu(layer(x, sc)) |
| x = F.dropout(x, p=self.dropout, training=self.training) |
| g = x.mean(dim=0) |
| return self.classifier(g) |
|
|
|
|
| class SimpleHybrid(nn.Module): |
| """ |
| Naive hybrid baseline: separate simplicial + hypergraph branches |
| concatenated (no bridge, no attention). |
| """ |
| |
| def __init__(self, in_dim, hidden_dim, num_classes, num_layers=2, dropout=0.3): |
| super().__init__() |
| half_dim = hidden_dim // 2 |
| |
| self.sc_layers = nn.ModuleList() |
| self.sc_layers.append(SimplicialConv(in_dim, half_dim)) |
| for _ in range(num_layers - 1): |
| self.sc_layers.append(SimplicialConv(half_dim, half_dim)) |
| |
| self.hg_layers = nn.ModuleList() |
| self.hg_layers.append(HypergraphConv(in_dim, half_dim)) |
| for _ in range(num_layers - 1): |
| self.hg_layers.append(HypergraphConv(half_dim, half_dim)) |
| |
| self.classifier = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, num_classes) |
| ) |
| self.dropout = dropout |
| |
| def forward(self, x, sc=None, hg=None, edge_index=None): |
| h_sc = x |
| for layer in self.sc_layers: |
| h_sc = F.relu(layer(h_sc, sc)) |
| h_sc = F.dropout(h_sc, p=self.dropout, training=self.training) |
| |
| h_hg = x |
| for layer in self.hg_layers: |
| h_hg = F.relu(layer(h_hg, hg)) |
| h_hg = F.dropout(h_hg, p=self.dropout, training=self.training) |
| |
| h = torch.cat([h_sc, h_hg], dim=-1) |
| g = h.mean(dim=0) |
| return self.classifier(g) |
|
|
|
|
| def get_model(name, in_dim, hidden_dim, num_classes, **kwargs): |
| """ |
| Factory function to create models by name. |
| |
| Args: |
| name: One of 'topohyper', 'gcn', 'gat', 'hgnn', 'simplicial', 'simple_hybrid' |
| in_dim: Input feature dimension |
| hidden_dim: Hidden dimension |
| num_classes: Number of classes |
| **kwargs: Additional model arguments (use_bridge, use_attention, etc.) |
| """ |
| models = { |
| 'topohyper': TopoHyperNet, |
| 'gcn': GCNNet, |
| 'gat': GATNet, |
| 'hgnn': HGNNNet, |
| 'simplicial': SimplicialNet, |
| 'simple_hybrid': SimpleHybrid, |
| } |
| |
| if name not in models: |
| raise ValueError(f"Unknown model: {name}. Choose from {list(models.keys())}") |
| |
| return models[name](in_dim, hidden_dim, num_classes, **kwargs) |
|
|
