wavedl 1.5.6__py3-none-any.whl → 1.6.0__py3-none-any.whl

wavedl/__init__.py

@@ -18,7 +18,7 @@ For inference:
     # or: python -m wavedl.test --checkpoint best_checkpoint --data_path test.npz
 """
 
-__version__ = "1.5.6"
+__version__ = "1.6.0"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

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,17 @@ from .unet import UNetRegression
 from .vit import ViTBase_, ViTSmall, ViTTiny
 
 
+# Optional timm-based models (imported conditionally)
+try:
+    from .caformer import CaFormerS18, CaFormerS36, PoolFormerS12
+    from .fastvit import FastViTS12, FastViTSA12, FastViTT8, FastViTT12
+    from .maxvit import MaxViTBaseLarge, MaxViTSmall, MaxViTTiny
+
+    _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 +99,11 @@ __all__ = [
     "ConvNeXtBase_",
     "ConvNeXtSmall",
     "ConvNeXtTiny",
+    "ConvNeXtV2Base",
+    "ConvNeXtV2BaseLarge",
+    "ConvNeXtV2Small",
+    "ConvNeXtV2Tiny",
+    "ConvNeXtV2TinyPretrained",
     "DenseNet121",
     "DenseNet169",
     "EfficientNetB0",
@@ -85,6 +112,7 @@ __all__ = [
     "EfficientNetV2L",
     "EfficientNetV2M",
     "EfficientNetV2S",
+    "Mamba1D",
     "MobileNetV3Large",
     "MobileNetV3Small",
     "RegNetY1_6GF",
@@ -105,8 +133,28 @@ __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",
+            "FastViTS12",
+            "FastViTSA12",
+            "FastViTT8",
+            "FastViTT12",
+            "MaxViTBaseLarge",
+            "MaxViTSmall",
+            "MaxViTTiny",
+            "PoolFormerS12",
+        ]
+    )

wavedl/models/_timm_utils.py

@@ -0,0 +1,238 @@
+"""
+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)}")

wavedl/models/caformer.py

@@ -0,0 +1,270 @@
+"""
+CaFormer: MetaFormer with Convolution and Attention
+====================================================
+
+CaFormer implements the MetaFormer architecture using depthwise separable
+convolutions in early stages and vanilla self-attention in later stages.
+
+**Key Features**:
+    - MetaFormer principle: architecture > token mixer
+    - Hybrid: Conv (early) + Attention (late)
+    - StarReLU activation for efficiency
+    - State-of-the-art on ImageNet (85.5%)
+
+**Variants**:
+    - caformer_s18: 26M params
+    - caformer_s36: 39M params
+    - caformer_m36: 56M params
+
+**Related Models**:
+    - PoolFormer: Uses pooling instead of attention
+    - ConvFormer: Uses only convolutions
+
+**Requirements**:
+    - timm >= 0.9.0 (for CaFormer models)
+
+Reference:
+    Yu, W., et al. (2023). MetaFormer Baselines for Vision.
+    TPAMI 2023. https://arxiv.org/abs/2210.13452
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+import torch
+import torch.nn as nn
+
+from wavedl.models._timm_utils import build_regression_head
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+__all__ = [
+    "CaFormerBase",
+    "CaFormerM36",
+    "CaFormerS18",
+    "CaFormerS36",
+    "PoolFormerS12",
+]
+
+
+# =============================================================================
+# CAFORMER BASE CLASS
+# =============================================================================
+
+
+class CaFormerBase(BaseModel):
+    """
+    CaFormer base class wrapping timm implementation.
+
+    MetaFormer with conv (early) + attention (late) token mixing.
+    2D only.
+    """
+
+    def __init__(
+        self,
+        in_shape: tuple[int, int],
+        out_size: int,
+        model_name: str = "caformer_s18",
+        pretrained: bool = True,
+        freeze_backbone: bool = False,
+        dropout_rate: float = 0.3,
+        **kwargs,
+    ):
+        super().__init__(in_shape, out_size)
+
+        if len(in_shape) != 2:
+            raise ValueError(f"CaFormer requires 2D input (H, W), got {len(in_shape)}D")
+
+        self.pretrained = pretrained
+        self.freeze_backbone = freeze_backbone
+        self.model_name = model_name
+
+        # Try to load from timm
+        try:
+            import timm
+
+            self.backbone = timm.create_model(
+                model_name,
+                pretrained=pretrained,
+                num_classes=0,  # Remove classifier
+            )
+
+            # Get feature dimension
+            with torch.no_grad():
+                dummy = torch.zeros(1, 3, *in_shape)
+                features = self.backbone(dummy)
+                in_features = features.shape[-1]
+
+        except ImportError:
+            raise ImportError(
+                "timm >= 0.9.0 is required for CaFormer. "
+                "Install with: pip install timm>=0.9.0"
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to load CaFormer model '{model_name}': {e}")
+
+        # Adapt input channels (3 -> 1)
+        self._adapt_input_channels()
+
+        # Regression head
+        self.head = build_regression_head(in_features, out_size, dropout_rate)
+
+        if freeze_backbone:
+            self._freeze_backbone()
+
+    def _adapt_input_channels(self):
+        """Adapt first conv layer for single-channel input."""
+        # CaFormer uses stem for first layer
+        if hasattr(self.backbone, "stem"):
+            first_conv = None
+            # Find first conv in stem
+            for name, module in self.backbone.stem.named_modules():
+                if isinstance(module, nn.Conv2d):
+                    first_conv = (name, module)
+                    break
+
+            if first_conv is not None:
+                name, old_conv = first_conv
+                new_conv = self._make_new_conv(old_conv)
+                # Set the new conv (handle nested structure)
+                self._set_module(self.backbone.stem, name, new_conv)
+
+    def _make_new_conv(self, old_conv: nn.Conv2d) -> nn.Conv2d:
+        """Create new conv layer with 1 input channel."""
+        new_conv = nn.Conv2d(
+            1,
+            old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+            bias=old_conv.bias is not None,
+        )
+        if self.pretrained:
+            with torch.no_grad():
+                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
+                if old_conv.bias is not None:
+                    new_conv.bias.copy_(old_conv.bias)
+        return new_conv
+
+    def _set_module(self, parent: nn.Module, name: str, module: nn.Module):
+        """Set a nested module by name."""
+        parts = name.split(".")
+        for part in parts[:-1]:
+            parent = getattr(parent, part)
+        setattr(parent, parts[-1], module)
+
+    def _freeze_backbone(self):
+        """Freeze backbone parameters."""
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        features = self.backbone(x)
+        return self.head(features)
+
+
+# =============================================================================
+# REGISTERED VARIANTS
+# =============================================================================
+
+
+@register_model("caformer_s18")
+class CaFormerS18(CaFormerBase):
+    """
+    CaFormer-S18: ~23.2M backbone parameters.
+
+    MetaFormer with conv + attention.
+    2D only.
+
+    Example:
+        >>> model = CaFormerS18(in_shape=(224, 224), out_size=3)
+        >>> x = torch.randn(4, 1, 224, 224)
+        >>> out = model(x)  # (4, 3)
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_name="caformer_s18",
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"CaFormer_S18(in_shape={self.in_shape}, out_size={self.out_size}, "
+            f"pretrained={self.pretrained})"
+        )
+
+
+@register_model("caformer_s36")
+class CaFormerS36(CaFormerBase):
+    """
+    CaFormer-S36: ~36.2M backbone parameters.
+
+    Deeper MetaFormer variant.
+    2D only.
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_name="caformer_s36",
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"CaFormer_S36(in_shape={self.in_shape}, out_size={self.out_size}, "
+            f"pretrained={self.pretrained})"
+        )
+
+
+@register_model("caformer_m36")
+class CaFormerM36(CaFormerBase):
+    """
+    CaFormer-M36: ~52.6M backbone parameters.
+
+    Medium-size MetaFormer variant.
+    2D only.
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_name="caformer_m36",
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"CaFormer_M36(in_shape={self.in_shape}, out_size={self.out_size}, "
+            f"pretrained={self.pretrained})"
+        )
+
+
+@register_model("poolformer_s12")
+class PoolFormerS12(CaFormerBase):
+    """
+    PoolFormer-S12: ~11.4M backbone parameters.
+
+    MetaFormer with simple pooling token mixer.
+    Proves that architecture matters more than complex attention.
+    2D only.
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_name="poolformer_s12",
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"PoolFormer_S12(in_shape={self.in_shape}, out_size={self.out_size}, "
+            f"pretrained={self.pretrained})"
+        )