Move transformer models (DeepIce, GRIT) into models/transformer

examples/04_training/06_train_icemix_model.py

@@ -16,7 +16,7 @@
 from graphnet.data.constants import FEATURES, TRUTH
 from graphnet.models import StandardModel
 from graphnet.models.detector.prometheus import Prometheus
-from graphnet.models.gnn import DeepIce
+from graphnet.models.transformer import DeepIce
 from graphnet.models.graphs import KNNGraph
 from graphnet.models.graphs.nodes import IceMixNodes
 from graphnet.models.task.reconstruction import (

examples/04_training/08_train_grit_model.py

@@ -11,7 +11,7 @@
 from graphnet.data.constants import FEATURES, TRUTH
 from graphnet.models import StandardModel
 from graphnet.models.detector.prometheus import Prometheus
-from graphnet.models.gnn import GRIT
+from graphnet.models.transformer import GRIT
 from graphnet.models.graphs import KNNGraphRRWP
 from graphnet.models.task.reconstruction import EnergyReconstruction
 from graphnet.training.callbacks import PiecewiseLinearLR

src/graphnet/models/gnn/__init__.py

@@ -5,6 +5,4 @@
 from .dynedge_jinst import DynEdgeJINST
 from .dynedge_kaggle_tito import DynEdgeTITO
 from .RNN_tito import RNN_TITO
-from .icemix import DeepIce
 from .particlenet import ParticleNeT
-from .grit import GRIT

src/graphnet/models/gnn/grit.py

@@ -1,164 +1,19 @@
-"""Implementation of GRIT, a graph transformer model.
+"""Deprecated module path for `GRIT`.
 
-Original author: Liheng Ma
-Original code: https://github.com/LiamMa/GRIT
-Paper: "Graph Inductive Biases in Transformers without Message Passing",
-        https://arxiv.org/abs/2305.17589
-
-Adapted by: Philip Weigel
+`GRIT` is a graph-transformer model and now lives in
+`graphnet.models.transformer`. This module re-exports it so existing imports
+keep working, and warns on import that the path has moved.
 """
 
-import torch.nn as nn
-from torch import Tensor
-from torch_geometric.data import Data
+from warnings import warn
 
-from graphnet.models.gnn.gnn import GNN
+from graphnet.models.transformer.grit import GRIT
 
-from graphnet.models.components.grit_layers import (
-    GritTransformerLayer,
-    SANGraphHead,
-)
-from graphnet.models.components.embedding import (
-    RRWPLinearEdgeEncoder,
-    RRWPLinearNodeEncoder,
-    LinearNodeEncoder,
-    LinearEdgeEncoder,
-    RWSELinearNodeEncoder,
+warn(
+    "`graphnet.models.gnn.grit` is deprecated; import `GRIT` from "
+    "`graphnet.models.transformer` instead.",
+    DeprecationWarning,
+    stacklevel=2,
 )
 
-
-class GRIT(GNN):
-    """GRIT is a graph transformer model.
-
-    Original code:
-    https://github.com/LiamMa/GRIT/blob/main/grit/network/grit_model.py
-    """
-
-    def __init__(
-        self,
-        nb_inputs: int,
-        hidden_dim: int,
-        nb_outputs: int = 1,
-        ksteps: int = 21,
-        n_layers: int = 10,
-        n_heads: int = 8,
-        pad_to_full_graph: bool = True,
-        add_node_attr_as_self_loop: bool = False,
-        dropout: float = 0.0,
-        fill_value: float = 0.0,
-        norm: nn.Module = nn.BatchNorm1d,
-        attn_dropout: float = 0.2,
-        edge_enhance: bool = True,
-        update_edges: bool = True,
-        attn_clamp: float = 5.0,
-        activation: nn.Module = nn.ReLU,
-        attn_activation: nn.Module = nn.ReLU,
-        norm_edges: bool = True,
-        enable_edge_transform: bool = True,
-        pred_head_layers: int = 2,
-        pred_head_activation: nn.Module = nn.ReLU,
-        pred_head_pooling: str = "mean",
-        position_encoding: str = "NoPE",
-    ):
-        """Construct `GRIT` model.
-
-        Args:
-            nb_inputs: Number of inputs.
-            hidden_dim: Size of hidden dimension.
-            nb_outputs: Size of output dimension.
-            ksteps: Number of random walk steps.
-            n_layers: Number of GRIT layers.
-            n_heads: Number of heads in MHA.
-            pad_to_full_graph: Pad to form fully-connected graph.
-            add_node_attr_as_self_loop: Adds node attr as an self-edge.
-            dropout: Dropout probability.
-            fill_value: Padding value.
-            norm: Uninstantiated normalization layer.
-                Either `torch.nn.BatchNorm1d` or `torch.nn.LayerNorm`.
-            attn_dropout: Attention dropout probability.
-            edge_enhance: Applies learnable weight matrix with node-pair in
-                output node calculation for MHA.
-            update_edges: Update edge values after GRIT layer.
-            attn_clamp: Clamp absolute value of attention scores to a value.
-            activation: Uninstantiated activation function.
-                E.g. `torch.nn.ReLU`
-            attn_activation: Uninstantiated attention activation function.
-                E.g. `torch.nn.ReLU`
-            norm_edges: Apply normalization layer to edges.
-            enable_edge_transform: Apply transformation to edges.
-            pred_head_layers: Number of layers in the prediction head.
-            pred_head_activation: Uninstantiated prediction head activation
-                    function. E.g. `torch.nn.ReLU`
-            pred_head_pooling: Pooling function to use for the prediction head,
-                either "mean" (default) or "add".
-            position_encoding: Method of position encoding.
-        """
-        super().__init__(nb_inputs, nb_outputs)
-        self.position_encoding = position_encoding.lower()
-        if self.position_encoding == "nope":
-            encoders = [
-                LinearNodeEncoder(nb_inputs, hidden_dim),
-                LinearEdgeEncoder(hidden_dim),
-            ]
-        elif self.position_encoding == "rrwp":
-            encoders = [
-                LinearNodeEncoder(nb_inputs, hidden_dim),
-                LinearEdgeEncoder(hidden_dim),
-                RRWPLinearNodeEncoder(ksteps, hidden_dim),
-                RRWPLinearEdgeEncoder(
-                    ksteps,
-                    hidden_dim,
-                    pad_to_full_graph=pad_to_full_graph,
-                    add_node_attr_as_self_loop=add_node_attr_as_self_loop,
-                    fill_value=fill_value,
-                ),
-            ]
-        elif self.position_encoding == "rwse":
-            encoders = [
-                LinearNodeEncoder(nb_inputs, hidden_dim - (ksteps - 1)),
-                RWSELinearNodeEncoder(ksteps - 1, hidden_dim),
-            ]
-        self.encoders = nn.ModuleList(encoders)
-
-        layers = []
-        for _ in range(n_layers):
-            layers.append(
-                GritTransformerLayer(
-                    in_dim=hidden_dim,
-                    out_dim=hidden_dim,
-                    num_heads=n_heads,
-                    dropout=dropout,
-                    activation=activation,
-                    attn_dropout=attn_dropout,
-                    norm=norm,
-                    residual=True,
-                    norm_edges=norm_edges,
-                    enable_edge_transform=enable_edge_transform,
-                    update_edges=update_edges,
-                    attn_activation=attn_activation,
-                    attn_clamp=attn_clamp,
-                    attn_edge_enhance=edge_enhance,
-                )
-            )
-        self.layers = nn.ModuleList(layers)
-        self.head = SANGraphHead(
-            dim_in=hidden_dim,
-            dim_out=nb_outputs,
-            L=pred_head_layers,
-            activation=pred_head_activation,
-            pooling=pred_head_pooling,
-        )
-
-    def forward(self, x: Data) -> Tensor:
-        """Forward pass."""
-        for encoder in self.encoders:
-            x = encoder(x)
-
-        # Apply GRIT layers
-        for layer in self.layers:
-            x = layer(x)
-
-        # Graph head
-        x = self.head(x)
-
-        return x
+__all__ = ["GRIT"]

src/graphnet/models/gnn/icemix.py

@@ -1,167 +1,19 @@
-"""Implementation of IceMix.
+"""Deprecated module path for `DeepIce`.
 
-This method was a solution submitted to the IceCube - Neutrinos in Deep Ice
-Kaggle competition.
-
-Solution by DrHB: https://github.com/DrHB/icecube-2nd-place
+`DeepIce` is a transformer model and now lives in
+`graphnet.models.transformer`. This module re-exports it so existing imports
+keep working, and warns on import that the path has moved.
 """
 
-import torch
-import torch.nn as nn
-from typing import Set, Dict, Any, Optional
-
-from graphnet.models.components.attention_blocks import (
-    Block_rel,
-    Block,
-)
-from graphnet.models.components.embedding import (
-    FourierEncoderEPJC,
-    SpacetimeEncoderEPJC,
-)
-from graphnet.models.gnn.dynedge import DynEdge
-from graphnet.models.gnn.gnn import GNN
-from graphnet.models.utils import array_to_sequence
-
-from torch_geometric.utils import to_dense_batch
-from torch_geometric.data import Data
-from torch import Tensor
-
-
-class DeepIce(GNN):
-    """DeepIce model."""
-
-    def __init__(
-        self,
-        hidden_dim: int = 384,
-        mlp_ratio: int = 4,
-        seq_length: int = 192,
-        depth: int = 12,
-        head_size: int = 32,
-        depth_rel: int = 4,
-        n_rel: int = 1,
-        scaled_emb: bool = False,
-        include_dynedge: bool = False,
-        dynedge_args: Optional[Dict[str, Any]] = None,
-        n_features: int = 6,
-    ):
-        """Construct `DeepIce`.
+from warnings import warn
 
-        Args:
-            hidden_dim: The latent feature dimension.
-            mlp_ratio: Mlp expansion ratio of FourierEncoderEPJC and
-                Transformer.
-            seq_length: The base feature dimension.
-            depth: The depth of the transformer.
-            head_size: The size of the attention heads.
-            depth_rel: The depth of the relative transformer.
-            n_rel: The number of relative transformer layers to use.
-            scaled_emb: Whether to scale the sinusoidal positional embeddings.
-            include_dynedge: If True, pulse-level predictions from `DynEdge`
-                will be added as features to the model.
-            dynedge_args: Initialization arguments for DynEdge. If not
-                provided, DynEdge will be initialized with the original Kaggle
-                Competition settings. If `include_dynedge` is False, this
-                argument have no impact.
-            n_features: The number of features in the input data.
-        """
-        super().__init__(seq_length, hidden_dim)
-        fourier_out_dim = hidden_dim // 2 if include_dynedge else hidden_dim
-        self.fourier_ext = FourierEncoderEPJC(
-            seq_length=seq_length,
-            mlp_dim=None,
-            output_dim=fourier_out_dim,
-            scaled=scaled_emb,
-            n_features=n_features,
-        )
-        self.rel_pos = SpacetimeEncoderEPJC(head_size)
-        self.sandwich = nn.ModuleList(
-            [
-                Block_rel(
-                    input_dim=hidden_dim, num_heads=hidden_dim // head_size
-                )
-                for _ in range(depth_rel)
-            ]
-        )
-        self.cls_token = nn.Linear(hidden_dim, 1, bias=False)
-        self.blocks = nn.ModuleList(
-            [
-                Block(
-                    input_dim=hidden_dim,
-                    num_heads=hidden_dim // head_size,
-                    mlp_ratio=mlp_ratio,
-                    drop_path=0.0 * (i / (depth - 1)),
-                    init_values=1,
-                )
-                for i in range(depth)
-            ]
-        )
-        self.n_rel = n_rel
+from graphnet.models.transformer.deepice import DeepIce
 
-        if include_dynedge and dynedge_args is None:
-            self.warning_once("Running with default DynEdge settings")
-            self.dyn_edge = DynEdge(
-                nb_inputs=9,
-                nb_neighbours=9,
-                post_processing_layer_sizes=[336, hidden_dim // 2],
-                dynedge_layer_sizes=[
-                    (128, 256),
-                    (336, 256),
-                    (336, 256),
-                    (336, 256),
-                ],
-                global_pooling_schemes=None,
-                activation_layer="gelu",
-                add_norm_layer=True,
-                skip_readout=True,
-            )
-        elif include_dynedge and not (dynedge_args is None):
-            self.dyn_edge = DynEdge(**dynedge_args)
-
-        self.include_dynedge = include_dynedge
-
-    @torch.jit.ignore
-    def no_weight_decay(self) -> Set:
-        """cls_tocken should not be subject to weight decay during training."""
-        return {"cls_token"}
-
-    def forward(self, data: Data) -> Tensor:
-        """Apply learnable forward pass."""
-        x0, mask, seq_length = array_to_sequence(
-            data.x, data.batch, padding_value=0
-        )
-        x = self.fourier_ext(x0, seq_length)
-        rel_pos_bias = self.rel_pos(x0)
-        batch_size = mask.shape[0]
-        if self.include_dynedge:
-            graph = self.dyn_edge(data)
-            graph, _ = to_dense_batch(graph, data.batch)
-            x = torch.cat([x, graph], 2)
-
-        attn_mask = torch.zeros(mask.shape, device=mask.device)
-        attn_mask[~mask] = -torch.inf
-
-        for i, blk in enumerate(self.sandwich):
-            x = blk(x, attn_mask, rel_pos_bias)
-            if i + 1 == self.n_rel:
-                rel_pos_bias = None
-
-        mask = torch.cat(
-            [
-                torch.ones(
-                    batch_size, 1, dtype=mask.dtype, device=mask.device
-                ),
-                mask,
-            ],
-            1,
-        )
-        attn_mask = torch.zeros(mask.shape, device=mask.device)
-        attn_mask[~mask] = -torch.inf
-        cls_token = self.cls_token.weight.unsqueeze(0).expand(
-            batch_size, -1, -1
-        )
-        x = torch.cat([cls_token, x], 1)
-
-        for blk in self.blocks:
-            x = blk(x, None, attn_mask)
+warn(
+    "`graphnet.models.gnn.icemix` is deprecated; import `DeepIce` from "
+    "`graphnet.models.transformer` instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
 
-        return x[:, 0]
+__all__ = ["DeepIce"]

src/graphnet/models/transformer/__init__.py

@@ -1,3 +1,5 @@
 """Transformer-specific modules."""
 
+from .deepice import DeepIce
+from .grit import GRIT
 from .iseecube import ISeeCube