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

wavedl/models/regnet.py

@@ -0,0 +1,383 @@
+"""
+RegNet: Designing Network Design Spaces
+========================================
+
+RegNet provides a family of models with predictable scaling behavior,
+designed through systematic exploration of network design spaces.
+Models scale smoothly from mobile to server deployments.
+
+**Key Features**:
+    - Predictable scaling: accuracy increases linearly with compute
+    - Simple, uniform architecture (no complex compound scaling)
+    - Group convolutions for efficiency
+    - Optional Squeeze-and-Excitation (SE) attention
+
+**Variants** (RegNetY includes SE attention):
+    - regnet_y_400mf: Ultra-light (~4.0M params, 0.4 GFLOPs)
+    - regnet_y_800mf: Light (~5.8M params, 0.8 GFLOPs)
+    - regnet_y_1_6gf: Medium (~10.5M params, 1.6 GFLOPs) - Recommended
+    - regnet_y_3_2gf: Large (~18.3M params, 3.2 GFLOPs)
+    - regnet_y_8gf: Very large (~37.9M params, 8.0 GFLOPs)
+
+**When to Use RegNet**:
+    - When you need predictable performance at a given compute budget
+    - For systematic model selection experiments
+    - When interpretability of design choices matters
+    - As an efficient alternative to ResNet
+
+**Note**: RegNet is 2D-only. For 1D data, use TCN. For 3D data, use ResNet3D.
+
+References:
+    Radosavovic, I., et al. (2020). Designing Network Design Spaces.
+    CVPR 2020. https://arxiv.org/abs/2003.13678
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+try:
+    from torchvision.models import (
+        RegNet_Y_1_6GF_Weights,
+        RegNet_Y_3_2GF_Weights,
+        RegNet_Y_8GF_Weights,
+        RegNet_Y_400MF_Weights,
+        RegNet_Y_800MF_Weights,
+        regnet_y_1_6gf,
+        regnet_y_3_2gf,
+        regnet_y_8gf,
+        regnet_y_400mf,
+        regnet_y_800mf,
+    )
+
+    REGNET_AVAILABLE = True
+except ImportError:
+    REGNET_AVAILABLE = False
+
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+class RegNetBase(BaseModel):
+    """
+    Base RegNet class for regression tasks.
+
+    Wraps torchvision RegNetY (with SE attention) with:
+    - Optional pretrained weights (ImageNet-1K)
+    - Automatic input channel adaptation (grayscale → 3ch)
+    - Custom regression head
+
+    RegNet advantages:
+    - Simple, uniform design (easy to understand and modify)
+    - Predictable accuracy/compute trade-off
+    - Efficient group convolutions
+    - SE attention for channel weighting (RegNetY variants)
+
+    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 RegNet 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 REGNET_AVAILABLE:
+            raise ImportError(
+                "torchvision is required for RegNet. "
+                "Install with: pip install torchvision"
+            )
+
+        if len(in_shape) != 2:
+            raise ValueError(
+                f"RegNet 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)
+
+        # RegNet uses .fc as the classification head
+        in_features = self.backbone.fc.in_features
+
+        # Replace fc with regression head
+        self.backbone.fc = nn.Sequential(
+            nn.Dropout(dropout_rate),
+            nn.Linear(in_features, regression_hidden),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout_rate * 0.5),
+            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 fc layer."""
+        for name, param in self.backbone.named_parameters():
+            if "fc" 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 RegNet."""
+        return {
+            "pretrained": True,
+            "dropout_rate": 0.2,
+            "freeze_backbone": False,
+            "regression_hidden": 256,
+        }
+
+
+# =============================================================================
+# REGISTERED MODEL VARIANTS
+# =============================================================================
+
+
+@register_model("regnet_y_400mf")
+class RegNetY400MF(RegNetBase):
+    """
+    RegNetY-400MF: Ultra-lightweight for constrained environments.
+
+    ~4.0M parameters, 0.4 GFLOPs. Smallest RegNet variant with SE attention.
+
+    Recommended for:
+        - Edge deployment with moderate accuracy needs
+        - Quick training experiments
+        - Baseline comparisons
+
+    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 = RegNetY400MF(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_fn=regnet_y_400mf,
+            weights_class=RegNet_Y_400MF_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"RegNetY_400MF({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("regnet_y_800mf")
+class RegNetY800MF(RegNetBase):
+    """
+    RegNetY-800MF: Light variant with good accuracy.
+
+    ~6.4M parameters, 0.8 GFLOPs. Good balance for mobile deployment.
+
+    Recommended for:
+        - Mobile/portable devices
+        - When MobileNet isn't accurate enough
+        - Moderate compute budgets
+
+    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 = RegNetY800MF(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_fn=regnet_y_800mf,
+            weights_class=RegNet_Y_800MF_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"RegNetY_800MF({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("regnet_y_1_6gf")
+class RegNetY1_6GF(RegNetBase):
+    """
+    RegNetY-1.6GF: Recommended default for balanced performance.
+
+    ~11.2M parameters, 1.6 GFLOPs. Best trade-off of accuracy and efficiency.
+    Comparable to ResNet50 but more efficient.
+
+    Recommended for:
+        - Default choice for general wave-based tasks
+        - When you want predictable scaling
+        - Server deployment with efficiency needs
+
+    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 = RegNetY1_6GF(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_fn=regnet_y_1_6gf,
+            weights_class=RegNet_Y_1_6GF_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"RegNetY_1.6GF({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("regnet_y_3_2gf")
+class RegNetY3_2GF(RegNetBase):
+    """
+    RegNetY-3.2GF: Higher accuracy for demanding tasks.
+
+    ~19.4M parameters, 3.2 GFLOPs. Use when 1.6GF isn't sufficient.
+
+    Recommended for:
+        - Larger datasets requiring more capacity
+        - When accuracy is more important than efficiency
+        - Research experiments with multiple model sizes
+
+    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 = RegNetY3_2GF(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_fn=regnet_y_3_2gf,
+            weights_class=RegNet_Y_3_2GF_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"RegNetY_3.2GF({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("regnet_y_8gf")
+class RegNetY8GF(RegNetBase):
+    """
+    RegNetY-8GF: High capacity for large-scale tasks.
+
+    ~39.2M parameters, 8.0 GFLOPs. Use for maximum accuracy needs.
+
+    Recommended for:
+        - Very large datasets (>50k samples)
+        - Complex wave patterns
+        - 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.2)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 256)
+
+    Example:
+        >>> model = RegNetY8GF(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_fn=regnet_y_8gf,
+            weights_class=RegNet_Y_8GF_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"RegNetY_8GF({pt}, in={self.in_shape}, out={self.out_size})"

wavedl/models/resnet.py

@@ -11,12 +11,15 @@ Provides multiple depth variants (18, 34, 50) with optional pretrained weights f
     - 3D: Volumetric data, CT/MRI (N, 1, D, H, W) → Conv3d
 
 **Variants**:
-    - resnet18: Lightweight, fast training
-    - resnet34: Balanced capacity
-    - resnet50: Higher capacity with bottleneck blocks
+    - resnet18: Lightweight, fast training (~11M params)
+    - resnet34: Balanced capacity (~21M params)
+    - resnet50: Higher capacity with bottleneck blocks (~25M params)
+
+References:
+    He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep Residual Learning
+    for Image Recognition. CVPR 2016. https://arxiv.org/abs/1512.03385
 
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 from typing import Any

wavedl/models/resnet3d.py

@@ -0,0 +1,258 @@
+"""
+ResNet3D: 3D Residual Networks for Volumetric Data
+===================================================
+
+3D extension of ResNet for processing volumetric data such as C-scans,
+3D wavefield imaging, and spatiotemporal cubes. Wraps torchvision's
+video models adapted for regression tasks.
+
+**Key Features**:
+    - Native 3D convolutions for volumetric processing
+    - Pretrained weights from Kinetics-400 (video action recognition)
+    - Adapted for single-channel input (grayscale volumes)
+    - Custom regression head for parameter estimation
+
+**Variants**:
+    - resnet3d_18: Lightweight (33M params)
+    - resnet3d_34: Medium depth
+    - resnet3d_50: Higher capacity with bottleneck blocks
+
+**Use Cases**:
+    - C-scan volume analysis (ultrasonic NDT)
+    - 3D wavefield imaging and inversion
+    - Spatiotemporal data cubes (time × space × space)
+    - Medical imaging (CT/MRI volumes)
+
+**Note**: ResNet3D is 3D-only. For 1D/2D data, use TCN or standard ResNet.
+
+References:
+    Hara, K., Kataoka, H., & Satoh, Y. (2018). Can Spatiotemporal 3D CNNs
+    Retrace the History of 2D CNNs and ImageNet? CVPR 2018.
+    https://arxiv.org/abs/1711.09577
+
+    He, K., et al. (2016). Deep Residual Learning for Image Recognition.
+    CVPR 2016. https://arxiv.org/abs/1512.03385
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+try:
+    from torchvision.models.video import (
+        MC3_18_Weights,
+        R3D_18_Weights,
+        mc3_18,
+        r3d_18,
+    )
+
+    RESNET3D_AVAILABLE = True
+except ImportError:
+    RESNET3D_AVAILABLE = False
+
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+class ResNet3DBase(BaseModel):
+    """
+    Base ResNet3D class for volumetric regression tasks.
+
+    Wraps torchvision 3D ResNet with:
+    - Optional pretrained weights (Kinetics-400)
+    - Automatic input channel adaptation (grayscale → 3ch)
+    - Custom regression head
+
+    Note: This is 3D-only. Input shape must be (D, H, W).
+    """
+
+    def __init__(
+        self,
+        in_shape: tuple[int, 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 ResNet3D for regression.
+
+        Args:
+            in_shape: (D, H, W) input volume dimensions
+            out_size: Number of regression output targets
+            model_fn: torchvision model constructor
+            weights_class: Pretrained weights enum class
+            pretrained: Use Kinetics-400 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 RESNET3D_AVAILABLE:
+            raise ImportError(
+                "torchvision >= 0.12 is required for ResNet3D. "
+                "Install with: pip install torchvision>=0.12"
+            )
+
+        if len(in_shape) != 3:
+            raise ValueError(
+                f"ResNet3D requires 3D input (D, H, W), got {len(in_shape)}D. "
+                "For 1D data, use TCN. For 2D data, use standard ResNet."
+            )
+
+        self.pretrained = pretrained
+        self.dropout_rate = dropout_rate
+        self.freeze_backbone = freeze_backbone
+        self.regression_hidden = regression_hidden
+
+        # Load pretrained backbone
+        weights = weights_class.DEFAULT if pretrained else None
+        self.backbone = model_fn(weights=weights)
+
+        # Get the fc input features
+        in_features = self.backbone.fc.in_features
+
+        # Replace fc with regression head
+        self.backbone.fc = nn.Sequential(
+            nn.Dropout(dropout_rate),
+            nn.Linear(in_features, regression_hidden),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout_rate * 0.5),
+            nn.Linear(regression_hidden, regression_hidden // 2),
+            nn.ReLU(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 fc head."""
+        for name, param in self.backbone.named_parameters():
+            if "fc" 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, D, 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, -1)
+
+        return self.backbone(x)
+
+    @classmethod
+    def get_default_config(cls) -> dict[str, Any]:
+        """Return default configuration for ResNet3D."""
+        return {
+            "pretrained": True,
+            "dropout_rate": 0.3,
+            "freeze_backbone": False,
+            "regression_hidden": 512,
+        }
+
+
+# =============================================================================
+# REGISTERED MODEL VARIANTS
+# =============================================================================
+
+
+@register_model("resnet3d_18")
+class ResNet3D18(ResNet3DBase):
+    """
+    ResNet3D-18: Lightweight 3D ResNet for volumetric data.
+
+    ~33M parameters. Uses 3D convolutions throughout for true volumetric processing.
+    Pretrained on Kinetics-400 (video action recognition).
+
+    Recommended for:
+        - C-scan ultrasonic inspection volumes
+        - 3D wavefield data cubes
+        - Medical imaging (CT/MRI)
+        - Moderate compute budgets
+
+    Args:
+        in_shape: (D, H, W) volume dimensions
+        out_size: Number of regression targets
+        pretrained: Use Kinetics-400 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 = ResNet3D18(in_shape=(16, 112, 112), out_size=3)
+        >>> x = torch.randn(2, 1, 16, 112, 112)
+        >>> out = model(x)  # (2, 3)
+    """
+
+    def __init__(self, in_shape: tuple[int, int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_fn=r3d_18,
+            weights_class=R3D_18_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"ResNet3D_18({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("mc3_18")
+class MC3_18(ResNet3DBase):
+    """
+    MC3-18: Mixed Convolution 3D ResNet (3D stem + 2D residual blocks).
+
+    ~11M parameters. More efficient than pure 3D ResNet while maintaining
+    good spatiotemporal modeling. Uses 3D convolutions in early layers
+    and 2D convolutions in later layers.
+
+    Recommended for:
+        - When pure 3D is too expensive
+        - Volumes with limited temporal/depth extent
+        - Faster training with reasonable accuracy
+
+    Args:
+        in_shape: (D, H, W) volume dimensions
+        out_size: Number of regression targets
+        pretrained: Use Kinetics-400 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 = MC3_18(in_shape=(16, 112, 112), out_size=3)
+        >>> x = torch.randn(2, 1, 16, 112, 112)
+        >>> out = model(x)  # (2, 3)
+    """
+
+    def __init__(self, in_shape: tuple[int, int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_fn=mc3_18,
+            weights_class=MC3_18_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"MC3_18({pt}, in={self.in_shape}, out={self.out_size})"