wavedl 1.3.1__py3-none-any.whl → 1.4.1__py3-none-any.whl

wavedl/models/efficientnetv2.py

@@ -0,0 +1,292 @@
+"""
+EfficientNetV2: Faster Training and Better Accuracy
+====================================================
+
+Next-generation EfficientNet with improved training efficiency and performance.
+EfficientNetV2 replaces early depthwise convolutions with fused MBConv blocks,
+enabling 2-4× faster training while achieving better accuracy.
+
+**Key Improvements over EfficientNet**:
+    - Fused-MBConv in early stages (faster on accelerators)
+    - Progressive learning support (start small, grow)
+    - Better NAS-optimized architecture
+
+**Variants**:
+    - efficientnet_v2_s: Small (21.5M params) - Recommended default
+    - efficientnet_v2_m: Medium (54.1M params) - Higher accuracy
+    - efficientnet_v2_l: Large (118.5M params) - Maximum accuracy
+
+**Note**: EfficientNetV2 is 2D-only. For 1D data, use TCN. For 3D data, use ResNet3D.
+
+References:
+    Tan, M., & Le, Q. (2021). EfficientNetV2: Smaller Models and Faster Training.
+    ICML 2021. https://arxiv.org/abs/2104.00298
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+try:
+    from torchvision.models import (
+        EfficientNet_V2_L_Weights,
+        EfficientNet_V2_M_Weights,
+        EfficientNet_V2_S_Weights,
+        efficientnet_v2_l,
+        efficientnet_v2_m,
+        efficientnet_v2_s,
+    )
+
+    EFFICIENTNETV2_AVAILABLE = True
+except ImportError:
+    EFFICIENTNETV2_AVAILABLE = False
+
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+class EfficientNetV2Base(BaseModel):
+    """
+    Base EfficientNetV2 class for regression tasks.
+
+    Wraps torchvision EfficientNetV2 with:
+    - Optional pretrained weights (ImageNet-1K)
+    - Automatic input channel adaptation (grayscale → 3ch)
+    - Custom multi-layer regression head
+
+    Compared to EfficientNet (V1):
+    - 2-4× faster training on GPU/TPU
+    - Better accuracy at similar parameter counts
+    - More efficient at higher resolutions
+
+    Note: This is 2D-only. Input shape must be (H, W).
+    """
+
+    def __init__(
+        self,
+        in_shape: tuple[int, int],
+        out_size: int,
+        model_fn,
+        weights_class,
+        pretrained: bool = True,
+        dropout_rate: float = 0.3,
+        freeze_backbone: bool = False,
+        regression_hidden: int = 512,
+        **kwargs,
+    ):
+        """
+        Initialize EfficientNetV2 for regression.
+
+        Args:
+            in_shape: (H, W) input image dimensions
+            out_size: Number of regression output targets
+            model_fn: torchvision model constructor
+            weights_class: Pretrained weights enum class
+            pretrained: Use ImageNet pretrained weights (default: True)
+            dropout_rate: Dropout rate in regression head (default: 0.3)
+            freeze_backbone: Freeze backbone for fine-tuning (default: False)
+            regression_hidden: Hidden units in regression head (default: 512)
+        """
+        super().__init__(in_shape, out_size)
+
+        if not EFFICIENTNETV2_AVAILABLE:
+            raise ImportError(
+                "torchvision >= 0.13 is required for EfficientNetV2. "
+                "Install with: pip install torchvision>=0.13"
+            )
+
+        if len(in_shape) != 2:
+            raise ValueError(
+                f"EfficientNetV2 requires 2D input (H, W), got {len(in_shape)}D. "
+                "For 1D data, use TCN. For 3D data, use ResNet3D."
+            )
+
+        self.pretrained = pretrained
+        self.dropout_rate = dropout_rate
+        self.freeze_backbone = freeze_backbone
+        self.regression_hidden = regression_hidden
+
+        # Load pretrained backbone
+        weights = weights_class.IMAGENET1K_V1 if pretrained else None
+        self.backbone = model_fn(weights=weights)
+
+        # Get classifier input features (before the final classification layer)
+        in_features = self.backbone.classifier[1].in_features
+
+        # Replace classifier with regression head
+        # EfficientNetV2 benefits from a deeper regression head
+        self.backbone.classifier = nn.Sequential(
+            nn.Dropout(dropout_rate),
+            nn.Linear(in_features, regression_hidden),
+            nn.SiLU(inplace=True),  # SiLU (Swish) matches EfficientNet's activation
+            nn.Dropout(dropout_rate * 0.5),
+            nn.Linear(regression_hidden, regression_hidden // 2),
+            nn.SiLU(inplace=True),
+            nn.Linear(regression_hidden // 2, out_size),
+        )
+
+        # Optionally freeze backbone for fine-tuning
+        if freeze_backbone:
+            self._freeze_backbone()
+
+    def _freeze_backbone(self):
+        """Freeze all backbone parameters except the classifier."""
+        for name, param in self.backbone.named_parameters():
+            if "classifier" not in name:
+                param.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x: Input tensor of shape (B, C, H, W) where C is 1 or 3
+
+        Returns:
+            Output tensor of shape (B, out_size)
+        """
+        # Expand single channel to 3 channels for pretrained weights compatibility
+        if x.size(1) == 1:
+            x = x.expand(-1, 3, -1, -1)
+
+        return self.backbone(x)
+
+    @classmethod
+    def get_default_config(cls) -> dict[str, Any]:
+        """Return default configuration for EfficientNetV2."""
+        return {
+            "pretrained": True,
+            "dropout_rate": 0.3,
+            "freeze_backbone": False,
+            "regression_hidden": 512,
+        }
+
+
+# =============================================================================
+# REGISTERED MODEL VARIANTS
+# =============================================================================
+
+
+@register_model("efficientnet_v2_s")
+class EfficientNetV2S(EfficientNetV2Base):
+    """
+    EfficientNetV2-S: Small variant, recommended default.
+
+    ~21.5M parameters. Best balance of speed and accuracy for most tasks.
+    2× faster training than EfficientNet-B4 with better accuracy.
+
+    Recommended for:
+        - Default choice for 2D wave data
+        - Moderate compute budgets
+        - When training speed matters
+
+    Args:
+        in_shape: (H, W) image dimensions
+        out_size: Number of regression targets
+        pretrained: Use ImageNet pretrained weights (default: True)
+        dropout_rate: Dropout rate in head (default: 0.3)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 512)
+
+    Example:
+        >>> model = EfficientNetV2S(in_shape=(500, 500), out_size=3)
+        >>> x = torch.randn(4, 1, 500, 500)
+        >>> 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_fn=efficientnet_v2_s,
+            weights_class=EfficientNet_V2_S_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"EfficientNetV2_S({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("efficientnet_v2_m")
+class EfficientNetV2M(EfficientNetV2Base):
+    """
+    EfficientNetV2-M: Medium variant for higher accuracy.
+
+    ~54.1M parameters. Use when accuracy is more important than speed.
+
+    Recommended for:
+        - Large datasets (>50k samples)
+        - Complex wave patterns
+        - When compute is not a bottleneck
+
+    Args:
+        in_shape: (H, W) image dimensions
+        out_size: Number of regression targets
+        pretrained: Use ImageNet pretrained weights (default: True)
+        dropout_rate: Dropout rate in head (default: 0.3)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 512)
+
+    Example:
+        >>> model = EfficientNetV2M(in_shape=(500, 500), out_size=3)
+        >>> x = torch.randn(4, 1, 500, 500)
+        >>> 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_fn=efficientnet_v2_m,
+            weights_class=EfficientNet_V2_M_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"EfficientNetV2_M({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("efficientnet_v2_l")
+class EfficientNetV2L(EfficientNetV2Base):
+    """
+    EfficientNetV2-L: Large variant for maximum accuracy.
+
+    ~118.5M parameters. Use only with large datasets and sufficient compute.
+
+    Recommended for:
+        - Very large datasets (>100k samples)
+        - When maximum accuracy is critical
+        - HPC environments with ample GPU memory
+
+    Args:
+        in_shape: (H, W) image dimensions
+        out_size: Number of regression targets
+        pretrained: Use ImageNet pretrained weights (default: True)
+        dropout_rate: Dropout rate in head (default: 0.3)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 512)
+
+    Example:
+        >>> model = EfficientNetV2L(in_shape=(500, 500), out_size=3)
+        >>> x = torch.randn(4, 1, 500, 500)
+        >>> 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_fn=efficientnet_v2_l,
+            weights_class=EfficientNet_V2_L_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"EfficientNetV2_L({pt}, in={self.in_shape}, out={self.out_size})"

wavedl/models/mobilenetv3.py

@@ -0,0 +1,272 @@
+"""
+MobileNetV3: Efficient Networks for Edge Deployment
+====================================================
+
+Lightweight architecture optimized for mobile and embedded devices.
+MobileNetV3 combines neural architecture search (NAS) with hardware-aware
+optimization to achieve excellent accuracy with minimal computational cost.
+
+**Key Features**:
+    - Inverted residuals with depthwise separable convolutions
+    - Squeeze-and-Excitation (SE) attention for channel weighting
+    - h-swish activation: efficient approximation of swish
+    - Designed for real-time inference on CPUs and edge devices
+
+**Variants**:
+    - mobilenet_v3_small: Ultra-lightweight (~1.1M params) - Edge/embedded
+    - mobilenet_v3_large: Balanced (~3.2M params) - Mobile deployment
+
+**Use Cases**:
+    - Real-time structural health monitoring on embedded systems
+    - Field inspection with portable devices
+    - When model size and inference speed are critical
+
+**Note**: MobileNetV3 is 2D-only. For 1D data, use TCN. For 3D data, use ResNet3D.
+
+References:
+    Howard, A., et al. (2019). Searching for MobileNetV3.
+    ICCV 2019. https://arxiv.org/abs/1905.02244
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+try:
+    from torchvision.models import (
+        MobileNet_V3_Large_Weights,
+        MobileNet_V3_Small_Weights,
+        mobilenet_v3_large,
+        mobilenet_v3_small,
+    )
+
+    MOBILENETV3_AVAILABLE = True
+except ImportError:
+    MOBILENETV3_AVAILABLE = False
+
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+class MobileNetV3Base(BaseModel):
+    """
+    Base MobileNetV3 class for regression tasks.
+
+    Wraps torchvision MobileNetV3 with:
+    - Optional pretrained weights (ImageNet-1K)
+    - Automatic input channel adaptation (grayscale → 3ch)
+    - Lightweight regression head (maintains efficiency)
+
+    MobileNetV3 is ideal for:
+    - Edge deployment (Raspberry Pi, Jetson, mobile)
+    - Real-time inference requirements
+    - Memory-constrained environments
+    - Quick prototyping and experimentation
+
+    Note: This is 2D-only. Input shape must be (H, W).
+    """
+
+    def __init__(
+        self,
+        in_shape: tuple[int, int],
+        out_size: int,
+        model_fn,
+        weights_class,
+        pretrained: bool = True,
+        dropout_rate: float = 0.2,
+        freeze_backbone: bool = False,
+        regression_hidden: int = 256,
+        **kwargs,
+    ):
+        """
+        Initialize MobileNetV3 for regression.
+
+        Args:
+            in_shape: (H, W) input image dimensions
+            out_size: Number of regression output targets
+            model_fn: torchvision model constructor
+            weights_class: Pretrained weights enum class
+            pretrained: Use ImageNet pretrained weights (default: True)
+            dropout_rate: Dropout rate in regression head (default: 0.2)
+            freeze_backbone: Freeze backbone for fine-tuning (default: False)
+            regression_hidden: Hidden units in regression head (default: 256)
+        """
+        super().__init__(in_shape, out_size)
+
+        if not MOBILENETV3_AVAILABLE:
+            raise ImportError(
+                "torchvision is required for MobileNetV3. "
+                "Install with: pip install torchvision"
+            )
+
+        if len(in_shape) != 2:
+            raise ValueError(
+                f"MobileNetV3 requires 2D input (H, W), got {len(in_shape)}D. "
+                "For 1D data, use TCN. For 3D data, use ResNet3D."
+            )
+
+        self.pretrained = pretrained
+        self.dropout_rate = dropout_rate
+        self.freeze_backbone = freeze_backbone
+        self.regression_hidden = regression_hidden
+
+        # Load pretrained backbone
+        weights = weights_class.IMAGENET1K_V1 if pretrained else None
+        self.backbone = model_fn(weights=weights)
+
+        # MobileNetV3 classifier structure:
+        # classifier[0]: Linear (features → 1280 for Large, 1024 for Small)
+        # classifier[1]: Hardswish
+        # classifier[2]: Dropout
+        # classifier[3]: Linear (1280/1024 → num_classes)
+
+        # Get the input features to the final classifier
+        in_features = self.backbone.classifier[0].in_features
+
+        # Replace classifier with lightweight regression head
+        # Keep it efficient to maintain MobileNet's speed advantage
+        self.backbone.classifier = nn.Sequential(
+            nn.Linear(in_features, regression_hidden),
+            nn.Hardswish(inplace=True),  # Match MobileNetV3's activation
+            nn.Dropout(dropout_rate),
+            nn.Linear(regression_hidden, out_size),
+        )
+
+        # Optionally freeze backbone for fine-tuning
+        if freeze_backbone:
+            self._freeze_backbone()
+
+    def _freeze_backbone(self):
+        """Freeze all backbone parameters except the classifier."""
+        for name, param in self.backbone.named_parameters():
+            if "classifier" not in name:
+                param.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x: Input tensor of shape (B, C, H, W) where C is 1 or 3
+
+        Returns:
+            Output tensor of shape (B, out_size)
+        """
+        # Expand single channel to 3 channels for pretrained weights compatibility
+        if x.size(1) == 1:
+            x = x.expand(-1, 3, -1, -1)
+
+        return self.backbone(x)
+
+    @classmethod
+    def get_default_config(cls) -> dict[str, Any]:
+        """Return default configuration for MobileNetV3."""
+        return {
+            "pretrained": True,
+            "dropout_rate": 0.2,
+            "freeze_backbone": False,
+            "regression_hidden": 256,
+        }
+
+
+# =============================================================================
+# REGISTERED MODEL VARIANTS
+# =============================================================================
+
+
+@register_model("mobilenet_v3_small")
+class MobileNetV3Small(MobileNetV3Base):
+    """
+    MobileNetV3-Small: Ultra-lightweight for edge deployment.
+
+    ~1.1M parameters. Designed for the most constrained environments.
+    Achieves ~67% ImageNet accuracy with minimal compute.
+
+    Recommended for:
+        - Embedded systems (Raspberry Pi, Arduino with accelerators)
+        - Battery-powered devices
+        - Ultra-low latency requirements (<10ms)
+        - Quick training experiments
+
+    Performance (approximate):
+        - CPU inference: ~6ms (single core)
+        - Parameters: 2.5M
+        - MAdds: 56M
+
+    Args:
+        in_shape: (H, W) image dimensions
+        out_size: Number of regression targets
+        pretrained: Use ImageNet pretrained weights (default: True)
+        dropout_rate: Dropout rate in head (default: 0.2)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 256)
+
+    Example:
+        >>> model = MobileNetV3Small(in_shape=(224, 224), out_size=3)
+        >>> x = torch.randn(1, 1, 224, 224)
+        >>> out = model(x)  # (1, 3)
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_fn=mobilenet_v3_small,
+            weights_class=MobileNet_V3_Small_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"MobileNetV3_Small({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("mobilenet_v3_large")
+class MobileNetV3Large(MobileNetV3Base):
+    """
+    MobileNetV3-Large: Balanced efficiency and accuracy.
+
+    ~3.2M parameters. Best trade-off for mobile/portable deployment.
+    Achieves ~75% ImageNet accuracy with efficient inference.
+
+    Recommended for:
+        - Mobile deployment (smartphones, tablets)
+        - Portable inspection devices
+        - Real-time processing with moderate accuracy needs
+        - Default choice for edge deployment
+
+    Performance (approximate):
+        - CPU inference: ~20ms (single core)
+        - Parameters: 5.4M
+        - MAdds: 219M
+
+    Args:
+        in_shape: (H, W) image dimensions
+        out_size: Number of regression targets
+        pretrained: Use ImageNet pretrained weights (default: True)
+        dropout_rate: Dropout rate in head (default: 0.2)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 256)
+
+    Example:
+        >>> model = MobileNetV3Large(in_shape=(224, 224), out_size=3)
+        >>> x = torch.randn(1, 1, 224, 224)
+        >>> out = model(x)  # (1, 3)
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_fn=mobilenet_v3_large,
+            weights_class=MobileNet_V3_Large_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"MobileNetV3_Large({pt}, in={self.in_shape}, out={self.out_size})"

wavedl/models/registry.py

@@ -6,7 +6,6 @@ Provides the core model registration and factory functionality.
 This module has no dependencies on other model modules to prevent circular imports.
 
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 import torch.nn as nn