wavedl 1.5.7__py3-none-any.whl → 1.6.1__py3-none-any.whl

wavedl/models/__init__.py

@@ -6,10 +6,11 @@ This module provides a centralized registry for neural network architectures,
 enabling dynamic model selection via command-line arguments.
 
 **Dimensionality Coverage**:
-    - 1D (waveforms): TCN, CNN, ResNet, ConvNeXt, DenseNet, ViT
-    - 2D (images): CNN, ResNet, ConvNeXt, DenseNet, ViT, UNet,
-                   EfficientNet, MobileNetV3, RegNet, Swin
-    - 3D (volumes): ResNet3D, CNN, ResNet, ConvNeXt, DenseNet
+    - 1D (waveforms): TCN, CNN, ResNet, ConvNeXt, ConvNeXt V2, DenseNet, ViT, Mamba
+    - 2D (images): CNN, ResNet, ConvNeXt, ConvNeXt V2, DenseNet, ViT, UNet,
+                   EfficientNet, MobileNetV3, RegNet, Swin, MaxViT, FastViT,
+                   CAFormer, PoolFormer, Vision Mamba
+    - 3D (volumes): ResNet3D, CNN, ResNet, ConvNeXt, ConvNeXt V2, DenseNet
 
 Usage:
     from wavedl.models import get_model, list_models, MODEL_REGISTRY
@@ -46,9 +47,19 @@ from .base import BaseModel
 # Import model implementations (triggers registration via decorators)
 from .cnn import CNN
 from .convnext import ConvNeXtBase_, ConvNeXtSmall, ConvNeXtTiny
+
+# New models (v1.6+)
+from .convnext_v2 import (
+    ConvNeXtV2Base,
+    ConvNeXtV2BaseLarge,
+    ConvNeXtV2Small,
+    ConvNeXtV2Tiny,
+    ConvNeXtV2TinyPretrained,
+)
 from .densenet import DenseNet121, DenseNet169
 from .efficientnet import EfficientNetB0, EfficientNetB1, EfficientNetB2
 from .efficientnetv2 import EfficientNetV2L, EfficientNetV2M, EfficientNetV2S
+from .mamba import Mamba1D, VimBase, VimSmall, VimTiny
 from .mobilenetv3 import MobileNetV3Large, MobileNetV3Small
 from .registry import (
     MODEL_REGISTRY,
@@ -66,6 +77,33 @@ from .unet import UNetRegression
 from .vit import ViTBase_, ViTSmall, ViTTiny
 
 
+# Optional timm-based models (imported conditionally)
+try:
+    from .caformer import CaFormerS18, CaFormerS36, PoolFormerS12
+    from .efficientvit import (
+        EfficientViTB0,
+        EfficientViTB1,
+        EfficientViTB2,
+        EfficientViTB3,
+        EfficientViTL1,
+        EfficientViTL2,
+        EfficientViTM0,
+        EfficientViTM1,
+        EfficientViTM2,
+    )
+    from .fastvit import FastViTS12, FastViTSA12, FastViTT8, FastViTT12
+    from .maxvit import MaxViTBaseLarge, MaxViTSmall, MaxViTTiny
+    from .unireplknet import (
+        UniRepLKNetBaseLarge,
+        UniRepLKNetSmall,
+        UniRepLKNetTiny,
+    )
+
+    _HAS_TIMM_MODELS = True
+except ImportError:
+    _HAS_TIMM_MODELS = False
+
+
 # Export public API (sorted alphabetically per RUF022)
 # See module docstring for dimensionality support details
 __all__ = [
@@ -77,6 +115,11 @@ __all__ = [
     "ConvNeXtBase_",
     "ConvNeXtSmall",
     "ConvNeXtTiny",
+    "ConvNeXtV2Base",
+    "ConvNeXtV2BaseLarge",
+    "ConvNeXtV2Small",
+    "ConvNeXtV2Tiny",
+    "ConvNeXtV2TinyPretrained",
     "DenseNet121",
     "DenseNet169",
     "EfficientNetB0",
@@ -85,6 +128,7 @@ __all__ = [
     "EfficientNetV2L",
     "EfficientNetV2M",
     "EfficientNetV2S",
+    "Mamba1D",
     "MobileNetV3Large",
     "MobileNetV3Small",
     "RegNetY1_6GF",
@@ -105,8 +149,40 @@ __all__ = [
     "ViTBase_",
     "ViTSmall",
     "ViTTiny",
+    "VimBase",
+    "VimSmall",
+    "VimTiny",
     "build_model",
     "get_model",
     "list_models",
     "register_model",
 ]
+
+# Add timm-based models to __all__ if available
+if _HAS_TIMM_MODELS:
+    __all__.extend(
+        [
+            "CaFormerS18",
+            "CaFormerS36",
+            "EfficientViTB0",
+            "EfficientViTB1",
+            "EfficientViTB2",
+            "EfficientViTB3",
+            "EfficientViTL1",
+            "EfficientViTL2",
+            "EfficientViTM0",
+            "EfficientViTM1",
+            "EfficientViTM2",
+            "FastViTS12",
+            "FastViTSA12",
+            "FastViTT8",
+            "FastViTT12",
+            "MaxViTBaseLarge",
+            "MaxViTSmall",
+            "MaxViTTiny",
+            "PoolFormerS12",
+            "UniRepLKNetBaseLarge",
+            "UniRepLKNetSmall",
+            "UniRepLKNetTiny",
+        ]
+    )

wavedl/models/_pretrained_utils.py

@@ -0,0 +1,366 @@
+"""
+Shared Utilities for Model Architectures
+=========================================
+
+Common components used across multiple models:
+- GRN (Global Response Normalization) for ConvNeXt V2
+- Dimension-agnostic layer factories
+- Regression head builders
+- Input channel adaptation for pretrained models
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# =============================================================================
+# DIMENSION-AGNOSTIC LAYER FACTORIES
+# =============================================================================
+
+
+def get_conv_layer(dim: int) -> type[nn.Module]:
+    """Get dimension-appropriate Conv class."""
+    layers = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d}
+    if dim not in layers:
+        raise ValueError(f"Unsupported dimension: {dim}")
+    return layers[dim]
+
+
+def get_norm_layer(dim: int) -> type[nn.Module]:
+    """Get dimension-appropriate BatchNorm class."""
+    layers = {1: nn.BatchNorm1d, 2: nn.BatchNorm2d, 3: nn.BatchNorm3d}
+    if dim not in layers:
+        raise ValueError(f"Unsupported dimension: {dim}")
+    return layers[dim]
+
+
+def get_pool_layer(dim: int) -> type[nn.Module]:
+    """Get dimension-appropriate AdaptiveAvgPool class."""
+    layers = {1: nn.AdaptiveAvgPool1d, 2: nn.AdaptiveAvgPool2d, 3: nn.AdaptiveAvgPool3d}
+    if dim not in layers:
+        raise ValueError(f"Unsupported dimension: {dim}")
+    return layers[dim]
+
+
+# =============================================================================
+# GLOBAL RESPONSE NORMALIZATION (GRN) - ConvNeXt V2
+# =============================================================================
+
+
+class GRN1d(nn.Module):
+    """
+    Global Response Normalization for 1D inputs.
+
+    GRN enhances inter-channel feature competition and promotes diversity.
+    Replaces LayerScale in ConvNeXt V2.
+
+    Reference: ConvNeXt V2 (CVPR 2023)
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, dim, 1))
+        self.beta = nn.Parameter(torch.zeros(1, dim, 1))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, C, L)
+        Gx = torch.norm(x, p=2, dim=2, keepdim=True)  # (B, C, 1)
+        Nx = Gx / (Gx.mean(dim=1, keepdim=True) + self.eps)  # (B, C, 1)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+class GRN2d(nn.Module):
+    """
+    Global Response Normalization for 2D inputs.
+
+    GRN enhances inter-channel feature competition and promotes diversity.
+    Replaces LayerScale in ConvNeXt V2.
+
+    Reference: ConvNeXt V2 (CVPR 2023)
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1))
+        self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, C, H, W)
+        Gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True)  # (B, C, 1, 1)
+        Nx = Gx / (Gx.mean(dim=1, keepdim=True) + self.eps)  # (B, C, 1, 1)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+class GRN3d(nn.Module):
+    """
+    Global Response Normalization for 3D inputs.
+
+    GRN enhances inter-channel feature competition and promotes diversity.
+    Replaces LayerScale in ConvNeXt V2.
+
+    Reference: ConvNeXt V2 (CVPR 2023)
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1, 1))
+        self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1, 1))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, C, D, H, W)
+        Gx = torch.norm(x, p=2, dim=(2, 3, 4), keepdim=True)  # (B, C, 1, 1, 1)
+        Nx = Gx / (Gx.mean(dim=1, keepdim=True) + self.eps)  # (B, C, 1, 1, 1)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+def get_grn_layer(dim: int) -> type[nn.Module]:
+    """Get dimension-appropriate GRN class."""
+    layers = {1: GRN1d, 2: GRN2d, 3: GRN3d}
+    if dim not in layers:
+        raise ValueError(f"Unsupported dimension: {dim}")
+    return layers[dim]
+
+
+# =============================================================================
+# LAYER NORMALIZATION (Channels Last for CNNs)
+# =============================================================================
+
+
+class LayerNormNd(nn.Module):
+    """
+    LayerNorm that works with channels-first tensors of any dimension.
+    Applies normalization over the channel dimension.
+    """
+
+    def __init__(self, normalized_shape: int, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.dim = dim
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Move channels to last, apply LN, move back
+        if self.dim == 1:
+            # (B, C, L) -> (B, L, C) -> LN -> (B, C, L)
+            x = x.permute(0, 2, 1)
+            x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+            x = x.permute(0, 2, 1)
+        elif self.dim == 2:
+            # (B, C, H, W) -> (B, H, W, C) -> LN -> (B, C, H, W)
+            x = x.permute(0, 2, 3, 1)
+            x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+            x = x.permute(0, 3, 1, 2)
+        elif self.dim == 3:
+            # (B, C, D, H, W) -> (B, D, H, W, C) -> LN -> (B, C, D, H, W)
+            x = x.permute(0, 2, 3, 4, 1)
+            x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+            x = x.permute(0, 4, 1, 2, 3)
+        return x
+
+
+# =============================================================================
+# REGRESSION HEAD BUILDERS
+# =============================================================================
+
+
+def build_regression_head(
+    in_features: int,
+    out_size: int,
+    dropout_rate: float = 0.3,
+    hidden_dim: int = 512,
+) -> nn.Sequential:
+    """
+    Build a standard regression head for pretrained models.
+
+    Args:
+        in_features: Input feature dimension
+        out_size: Number of regression targets
+        dropout_rate: Dropout rate
+        hidden_dim: Hidden layer dimension
+
+    Returns:
+        nn.Sequential regression head
+    """
+    return nn.Sequential(
+        nn.Dropout(dropout_rate),
+        nn.Linear(in_features, hidden_dim),
+        nn.SiLU(inplace=True),
+        nn.Dropout(dropout_rate * 0.5),
+        nn.Linear(hidden_dim, hidden_dim // 2),
+        nn.SiLU(inplace=True),
+        nn.Linear(hidden_dim // 2, out_size),
+    )
+
+
+def adapt_input_channels(
+    conv_layer: nn.Module,
+    new_in_channels: int = 1,
+    pretrained: bool = True,
+) -> nn.Module:
+    """
+    Adapt a convolutional layer for different input channels.
+
+    For pretrained models, averages RGB weights to grayscale.
+
+    Args:
+        conv_layer: Original conv layer (expects 3 input channels)
+        new_in_channels: New number of input channels (default: 1)
+        pretrained: Whether to adapt pretrained weights
+
+    Returns:
+        New conv layer with adapted input channels
+    """
+    if isinstance(conv_layer, nn.Conv2d):
+        new_conv = nn.Conv2d(
+            new_in_channels,
+            conv_layer.out_channels,
+            kernel_size=conv_layer.kernel_size,
+            stride=conv_layer.stride,
+            padding=conv_layer.padding,
+            bias=conv_layer.bias is not None,
+        )
+        if pretrained and conv_layer.in_channels == 3:
+            with torch.no_grad():
+                # Average RGB weights
+                new_conv.weight.copy_(conv_layer.weight.mean(dim=1, keepdim=True))
+                if conv_layer.bias is not None:
+                    new_conv.bias.copy_(conv_layer.bias)
+        return new_conv
+    else:
+        raise NotImplementedError(f"Unsupported layer type: {type(conv_layer)}")
+
+
+def adapt_first_conv_for_single_channel(
+    module: nn.Module,
+    conv_path: str,
+    pretrained: bool = True,
+) -> None:
+    """
+    Adapt the first convolutional layer of a pretrained model for single-channel input.
+
+    This is a convenience function for torchvision-style models where the path
+    to the first conv layer is known. It modifies the model in-place.
+
+    For pretrained models, the RGB weights are averaged to create grayscale weights,
+    which provides a reasonable initialization for single-channel inputs.
+
+    Args:
+        module: The model or submodule containing the conv layer
+        conv_path: Dot-separated path to the conv layer (e.g., "conv1", "features.0.0")
+        pretrained: Whether to adapt pretrained weights by averaging RGB channels
+
+    Example:
+        >>> # For torchvision ResNet
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "conv1", pretrained=True
+        ... )
+        >>> # For torchvision ConvNeXt
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "features.0.0", pretrained=True
+        ... )
+        >>> # For torchvision DenseNet
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "features.conv0", pretrained=True
+        ... )
+    """
+    # Navigate to parent and get the conv layer
+    parts = conv_path.split(".")
+    parent = module
+    for part in parts[:-1]:
+        if part.isdigit():
+            parent = parent[int(part)]
+        else:
+            parent = getattr(parent, part)
+
+    # Get the final attribute name and the old conv
+    final_attr = parts[-1]
+    if final_attr.isdigit():
+        old_conv = parent[int(final_attr)]
+    else:
+        old_conv = getattr(parent, final_attr)
+
+    # Create and set the new conv
+    new_conv = adapt_input_channels(old_conv, new_in_channels=1, pretrained=pretrained)
+
+    if final_attr.isdigit():
+        parent[int(final_attr)] = new_conv
+    else:
+        setattr(parent, final_attr, new_conv)
+
+
+def find_and_adapt_input_convs(
+    backbone: nn.Module,
+    pretrained: bool = True,
+    adapt_all: bool = False,
+) -> int:
+    """
+    Find and adapt Conv2d layers with 3 input channels for single-channel input.
+
+    This is useful for timm-style models where the exact path to the first
+    conv layer may vary or where multiple layers need adaptation.
+
+    Args:
+        backbone: The backbone model to adapt
+        pretrained: Whether to adapt pretrained weights by averaging RGB channels
+        adapt_all: If True, adapt all Conv2d layers with 3 input channels.
+                   If False (default), only adapt the first one found.
+
+    Returns:
+        Number of layers adapted
+
+    Example:
+        >>> # For timm models (adapt first conv only)
+        >>> count = find_and_adapt_input_convs(model.backbone, pretrained=True)
+        >>> # For models with multiple input convs (e.g., FastViT)
+        >>> count = find_and_adapt_input_convs(
+        ...     model.backbone, pretrained=True, adapt_all=True
+        ... )
+    """
+    adapted_count = 0
+
+    for name, module in backbone.named_modules():
+        if not hasattr(module, "in_channels") or module.in_channels != 3:
+            continue
+
+        # Check if this is a wrapper with inner .conv attribute
+        if hasattr(module, "conv") and isinstance(module.conv, nn.Conv2d):
+            old_conv = module.conv
+            module.conv = adapt_input_channels(
+                old_conv, new_in_channels=1, pretrained=pretrained
+            )
+            adapted_count += 1
+
+        elif isinstance(module, nn.Conv2d):
+            # Direct Conv2d - need to replace it in parent
+            parts = name.split(".")
+            parent = backbone
+            for part in parts[:-1]:
+                if part.isdigit():
+                    parent = parent[int(part)]
+                else:
+                    parent = getattr(parent, part)
+
+            child_name = parts[-1]
+            new_conv = adapt_input_channels(
+                module, new_in_channels=1, pretrained=pretrained
+            )
+
+            if child_name.isdigit():
+                parent[int(child_name)] = new_conv
+            else:
+                setattr(parent, child_name, new_conv)
+
+            adapted_count += 1
+
+        if not adapt_all and adapted_count > 0:
+            break
+
+    return adapted_count

wavedl/models/base.py

@@ -15,6 +15,54 @@ import torch
 import torch.nn as nn
 
 
+# =============================================================================
+# TYPE ALIASES
+# =============================================================================
+
+# Spatial shape type aliases for model input dimensions
+SpatialShape1D = tuple[int]
+SpatialShape2D = tuple[int, int]
+SpatialShape3D = tuple[int, int, int]
+SpatialShape = SpatialShape1D | SpatialShape2D | SpatialShape3D
+
+
+# =============================================================================
+# UTILITY FUNCTIONS
+# =============================================================================
+
+
+def compute_num_groups(num_channels: int, preferred_groups: int = 32) -> int:
+    """
+    Compute valid num_groups for GroupNorm that divides num_channels evenly.
+
+    GroupNorm requires num_channels to be divisible by num_groups. This function
+    finds the largest valid divisor up to preferred_groups.
+
+    Args:
+        num_channels: Number of channels to normalize (must be positive)
+        preferred_groups: Preferred number of groups (default: 32)
+
+    Returns:
+        Valid num_groups that satisfies num_channels % num_groups == 0
+
+    Example:
+        >>> compute_num_groups(64)  # Returns 32
+        >>> compute_num_groups(48)  # Returns 16 (48 % 32 != 0)
+        >>> compute_num_groups(7)  # Returns 1 (prime number)
+    """
+    # Try preferred groups first, then common divisors
+    for groups in [preferred_groups, 16, 8, 4, 2, 1]:
+        if groups <= num_channels and num_channels % groups == 0:
+            return groups
+
+    # Fallback: find any valid divisor (always returns at least 1)
+    for groups in range(min(32, num_channels), 0, -1):
+        if num_channels % groups == 0:
+            return groups
+
+    return 1  # Always valid
+
+
 class BaseModel(nn.Module, ABC):
     """
     Abstract base class for all regression models.