wavedl 1.3.1__py3-none-any.whl → 1.4.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.3.1"
+__version__ = "1.4.0"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

wavedl/hpc.py

@@ -33,7 +33,7 @@ from pathlib import Path
 def detect_gpus() -> int:
     """Auto-detect available GPUs using nvidia-smi."""
     if shutil.which("nvidia-smi") is None:
-        print("Warning: nvidia-smi not found, defaulting to 1 GPU")
+        print("Warning: nvidia-smi not found, defaulting to NUM_GPUS=1")
         return 1
 
     try:
@@ -50,7 +50,7 @@ def detect_gpus() -> int:
     except (subprocess.CalledProcessError, FileNotFoundError):
         pass
 
-    print("Warning: Could not detect GPUs, defaulting to 1")
+    print("Warning: No GPUs detected, defaulting to NUM_GPUS=1")
     return 1
 
 
@@ -61,10 +61,15 @@ def setup_hpc_environment() -> None:
     offline logging configurations.
     """
     # Use SLURM_TMPDIR if available, otherwise system temp
-    tmpdir = os.environ.get("SLURM_TMPDIR", tempfile.gettempdir())
+    tmpdir = os.environ.get(
+        "SLURM_TMPDIR", os.environ.get("TMPDIR", tempfile.gettempdir())
+    )
 
     # Configure directories for systems with restricted home directories
     os.environ.setdefault("MPLCONFIGDIR", f"{tmpdir}/matplotlib")
+    os.environ.setdefault(
+        "FONTCONFIG_PATH", os.environ.get("FONTCONFIG_PATH", "/etc/fonts")
+    )
     os.environ.setdefault("XDG_CACHE_HOME", f"{tmpdir}/.cache")
 
     # Ensure matplotlib config dir exists
@@ -147,11 +152,11 @@ Environment Variables:
 def print_summary(exit_code: int, wandb_mode: str, wandb_dir: str) -> None:
     """Print post-training summary and instructions."""
     print()
-    print("=" * 50)
+    print("=" * 40)
 
     if exit_code == 0:
         print("✅ Training completed successfully!")
-        print("=" * 50)
+        print("=" * 40)
 
         if wandb_mode == "offline":
             print()
@@ -162,15 +167,15 @@ def print_summary(exit_code: int, wandb_mode: str, wandb_dir: str) -> None:
             print("   This will upload your training logs to wandb.ai")
     else:
         print(f"❌ Training failed with exit code: {exit_code}")
-        print("=" * 50)
+        print("=" * 40)
         print()
         print("Common issues:")
         print("  - Missing data file (check --data_path)")
         print("  - Insufficient GPU memory (reduce --batch_size)")
-        print("  - Invalid model name (run: wavedl-train --list_models)")
+        print("  - Invalid model name (run: python train.py --list_models)")
         print()
 
-    print("=" * 50)
+    print("=" * 40)
     print()
 
 
@@ -182,17 +187,27 @@ def main() -> int:
     # Setup HPC environment
     setup_hpc_environment()
 
+    # Check if wavedl package is importable
+    try:
+        import wavedl  # noqa: F401
+    except ImportError:
+        print("Error: wavedl package not found. Run: pip install -e .", file=sys.stderr)
+        return 1
+
     # Auto-detect GPUs if not specified
-    num_gpus = args.num_gpus if args.num_gpus is not None else detect_gpus()
+    if args.num_gpus is not None:
+        num_gpus = args.num_gpus
+        print(f"Using NUM_GPUS={num_gpus} (set via command line)")
+    else:
+        num_gpus = detect_gpus()
 
     # Build accelerate launch command
     cmd = [
-        sys.executable,
-        "-m",
-        "accelerate.commands.launch",
+        "accelerate",
+        "launch",
         f"--num_processes={num_gpus}",
         f"--num_machines={args.num_machines}",
-        f"--machine_rank={args.machine_rank}",
+        "--machine_rank=0",
         f"--mixed_precision={args.mixed_precision}",
         f"--dynamo_backend={args.dynamo_backend}",
         "-m",
@@ -208,19 +223,6 @@ def main() -> int:
             Path(arg.split("=", 1)[1]).mkdir(parents=True, exist_ok=True)
             break
 
-    # Print launch configuration
-    print()
-    print("=" * 50)
-    print("🚀 WaveDL HPC Training Launcher")
-    print("=" * 50)
-    print(f"   GPUs: {num_gpus}")
-    print(f"   Machines: {args.num_machines}")
-    print(f"   Mixed Precision: {args.mixed_precision}")
-    print(f"   Dynamo Backend: {args.dynamo_backend}")
-    print(f"   WandB Mode: {os.environ.get('WANDB_MODE', 'offline')}")
-    print("=" * 50)
-    print()
-
     # Launch training
     try:
         result = subprocess.run(cmd, check=False)

wavedl/models/__init__.py

@@ -5,6 +5,12 @@ Model Registry and Factory Pattern for Deep Learning Architectures
 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
+
 Usage:
     from wavedl.models import get_model, list_models, MODEL_REGISTRY
 
@@ -31,7 +37,6 @@ Adding New Models:
                 ...
 
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 # Import registry first (no dependencies)
@@ -43,6 +48,8 @@ from .cnn import CNN
 from .convnext import ConvNeXtBase_, ConvNeXtSmall, ConvNeXtTiny
 from .densenet import DenseNet121, DenseNet169
 from .efficientnet import EfficientNetB0, EfficientNetB1, EfficientNetB2
+from .efficientnetv2 import EfficientNetV2L, EfficientNetV2M, EfficientNetV2S
+from .mobilenetv3 import MobileNetV3Large, MobileNetV3Small
 from .registry import (
     MODEL_REGISTRY,
     build_model,
@@ -50,18 +57,22 @@ from .registry import (
     list_models,
     register_model,
 )
+from .regnet import RegNetY1_6GF, RegNetY3_2GF, RegNetY8GF, RegNetY400MF, RegNetY800MF
 from .resnet import ResNet18, ResNet34, ResNet50
-from .unet import UNet, UNetRegression
+from .resnet3d import MC3_18, ResNet3D18
+from .swin import SwinBase, SwinSmall, SwinTiny
+from .tcn import TCN, TCNLarge, TCNSmall
+from .unet import UNetRegression
 from .vit import ViTBase_, ViTSmall, ViTTiny
 
 
-# Export public API
+# Export public API (sorted alphabetically per RUF022)
+# See module docstring for dimensionality support details
 __all__ = [
-    # Models
     "CNN",
-    # Registry
+    "MC3_18",
     "MODEL_REGISTRY",
-    # Base class
+    "TCN",
     "BaseModel",
     "ConvNeXtBase_",
     "ConvNeXtSmall",
@@ -71,10 +82,25 @@ __all__ = [
     "EfficientNetB0",
     "EfficientNetB1",
     "EfficientNetB2",
+    "EfficientNetV2L",
+    "EfficientNetV2M",
+    "EfficientNetV2S",
+    "MobileNetV3Large",
+    "MobileNetV3Small",
+    "RegNetY1_6GF",
+    "RegNetY3_2GF",
+    "RegNetY8GF",
+    "RegNetY400MF",
+    "RegNetY800MF",
+    "ResNet3D18",
     "ResNet18",
     "ResNet34",
     "ResNet50",
-    "UNet",
+    "SwinBase",
+    "SwinSmall",
+    "SwinTiny",
+    "TCNLarge",
+    "TCNSmall",
     "UNetRegression",
     "ViTBase_",
     "ViTSmall",

wavedl/models/_template.py

@@ -14,7 +14,6 @@ Steps to Add a New Model:
     5. Run: accelerate launch -m wavedl.train --model your_model --wandb
 
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 from typing import Any

wavedl/models/base.py

@@ -6,7 +6,6 @@ Defines the interface contract that all models must implement for compatibility
 with the training pipeline. Provides common utilities and enforces consistency.
 
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 from abc import ABC, abstractmethod

wavedl/models/cnn.py

@@ -17,7 +17,6 @@ Use this as:
     - A starting point for custom modifications
 
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 from typing import Any

wavedl/models/convnext.py

@@ -15,8 +15,11 @@ Features: inverted bottleneck, LayerNorm, GELU activation, depthwise convolution
     - convnext_small: Medium (~50M params for 2D)
     - convnext_base: Standard (~89M params for 2D)
 
+References:
+    Liu, Z., et al. (2022). A ConvNet for the 2020s.
+    CVPR 2022. https://arxiv.org/abs/2201.03545
+
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 from typing import Any

wavedl/models/densenet.py

@@ -14,8 +14,11 @@ Features: feature reuse, efficient gradient flow, compact model.
     - densenet121: Standard (121 layers, ~8M params for 2D)
     - densenet169: Deeper (169 layers, ~14M params for 2D)
 
+References:
+    Huang, G., et al. (2017). Densely Connected Convolutional Networks.
+    CVPR 2017 (Best Paper). https://arxiv.org/abs/1608.06993
+
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 from typing import Any

wavedl/models/efficientnet.py

@@ -6,14 +6,18 @@ Wrapper around torchvision's EfficientNet with a regression head.
 Provides optional ImageNet pretrained weights for transfer learning.
 
 **Variants**:
-    - efficientnet_b0: Smallest, fastest (5.3M params)
-    - efficientnet_b1: Light (7.8M params)
-    - efficientnet_b2: Balanced (9.1M params)
+    - efficientnet_b0: Smallest, fastest (~4.7M params)
+    - efficientnet_b1: Light (~7.2M params)
+    - efficientnet_b2: Balanced (~8.4M params)
 
-**Note**: EfficientNet is 2D-only. For 1D/3D data, use ResNet or CNN.
+**Note**: EfficientNet is 2D-only. For 1D data, use TCN. For 3D data, use ResNet3D.
+
+References:
+    Tan, M., & Le, Q. (2019). EfficientNet: Rethinking Model Scaling
+    for Convolutional Neural Networks. ICML 2019.
+    https://arxiv.org/abs/1905.11946
 
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
 from typing import Any

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})"