wavedl 1.3.1__py3-none-any.whl → 1.4.1__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.1"
 __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
@@ -125,6 +130,18 @@ Environment Variables:
         default=0,
         help="Rank of this machine in multi-node setup (default: 0)",
     )
+    parser.add_argument(
+        "--main_process_ip",
+        type=str,
+        default=None,
+        help="IP address of the main process for multi-node training",
+    )
+    parser.add_argument(
+        "--main_process_port",
+        type=int,
+        default=None,
+        help="Port for multi-node communication (default: accelerate auto-selects)",
+    )
     parser.add_argument(
         "--mixed_precision",
         type=str,
@@ -147,11 +164,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 +179,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,22 +199,38 @@ 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}",
         f"--mixed_precision={args.mixed_precision}",
         f"--dynamo_backend={args.dynamo_backend}",
-        "-m",
-        "wavedl.train",
-    ] + train_args
+    ]
+
+    # Add multi-node networking args if specified (required for some clusters)
+    if args.main_process_ip:
+        cmd.append(f"--main_process_ip={args.main_process_ip}")
+    if args.main_process_port:
+        cmd.append(f"--main_process_port={args.main_process_port}")
+
+    cmd += ["-m", "wavedl.train"] + train_args
 
     # Create output directory if specified
     for i, arg in enumerate(train_args):
@@ -208,19 +241,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

@@ -1,24 +1,26 @@
 """
-Model Template for New Architectures
-=====================================
+Model Template for Custom Architectures
+========================================
 
-Copy this file and modify to add new model architectures to the framework.
+Copy this file and modify to add custom model architectures to WaveDL.
 The model will be automatically registered and available via --model flag.
 
-Steps to Add a New Model:
-    1. Copy this file to models/your_model.py
-    2. Rename the class and update @register_model("your_model")
-    3. Implement the __init__ and forward methods
-    4. Import your model in models/__init__.py:
-       from wavedl.models.your_model import YourModel
-    5. Run: accelerate launch -m wavedl.train --model your_model --wandb
+Quick Start:
+    1. Copy this file to your project: cp _template.py my_model.py
+    2. Rename the class and update @register_model("my_model")
+    3. Implement your architecture in __init__ and forward
+    4. Train: wavedl-train --import my_model --model my_model --data_path data.npz
+
+Requirements (your model MUST):
+    1. Inherit from BaseModel
+    2. Accept (in_shape, out_size, **kwargs) in __init__
+    3. Return tensor of shape (batch, out_size) from forward()
+
+See README.md "Adding Custom Models" section for more details.
 
 Author: Ductho Le (ductho.le@outlook.com)
-Version: 1.0.0
 """
 
-from typing import Any
-
 import torch
 import torch.nn as nn
 
@@ -26,7 +28,7 @@ from wavedl.models.base import BaseModel
 
 
 # Uncomment the decorator to register this model
-# @register_model("template")
+# @register_model("my_model")
 class TemplateModel(BaseModel):
     """
     Template Model Architecture.
@@ -35,14 +37,16 @@ class TemplateModel(BaseModel):
     The first line will appear in --list_models output.
 
     Args:
-        in_shape: Input spatial dimensions (H, W)
-        out_size: Number of regression output targets
+        in_shape: Input spatial dimensions (auto-detected from data)
+                  - 1D: (L,) for signals
+                  - 2D: (H, W) for images
+                  - 3D: (D, H, W) for volumes
+        out_size: Number of regression targets (auto-detected from data)
         hidden_dim: Size of hidden layers (default: 256)
-        num_layers: Number of convolutional layers (default: 4)
         dropout: Dropout rate (default: 0.1)
 
     Input Shape:
-        (B, 1, H, W) - Single-channel images
+        (B, 1, *in_shape) - e.g., (B, 1, 64, 64) for 2D
 
     Output Shape:
         (B, out_size) - Regression predictions
@@ -50,10 +54,9 @@ class TemplateModel(BaseModel):
 
     def __init__(
         self,
-        in_shape: tuple[int, int],
+        in_shape: tuple,
         out_size: int,
         hidden_dim: int = 256,
-        num_layers: int = 4,
         dropout: float = 0.1,
         **kwargs,  # Accept extra kwargs for flexibility
     ):
@@ -62,14 +65,13 @@ class TemplateModel(BaseModel):
 
         # Store hyperparameters as attributes (optional but recommended)
         self.hidden_dim = hidden_dim
-        self.num_layers = num_layers
         self.dropout_rate = dropout
 
         # =================================================================
         # BUILD YOUR ARCHITECTURE HERE
         # =================================================================
 
-        # Example: Simple CNN encoder
+        # Example: Simple CNN encoder (assumes 2D input with 1 channel)
         self.encoder = nn.Sequential(
             # Layer 1
             nn.Conv2d(1, 32, kernel_size=3, padding=1),
@@ -107,10 +109,10 @@ class TemplateModel(BaseModel):
         """
         Forward pass of the model.
 
-        REQUIRED: Must accept (B, C, H, W) and return (B, out_size)
+        REQUIRED: Must accept (B, C, *spatial) and return (B, out_size)
 
         Args:
-            x: Input tensor of shape (B, 1, H, W)
+            x: Input tensor of shape (B, 1, *in_shape)
 
         Returns:
             Output tensor of shape (B, out_size)
@@ -123,35 +125,20 @@ class TemplateModel(BaseModel):
 
         return output
 
-    @classmethod
-    def get_default_config(cls) -> dict[str, Any]:
-        """
-        Return default hyperparameters for this model.
-
-        OPTIONAL: Override to provide model-specific defaults.
-        These can be used for documentation or config files.
-        """
-        return {
-            "hidden_dim": 256,
-            "num_layers": 4,
-            "dropout": 0.1,
-        }
-
 
 # =============================================================================
 # USAGE EXAMPLE
 # =============================================================================
 if __name__ == "__main__":
     # Quick test of the model
-    model = TemplateModel(in_shape=(500, 500), out_size=5)
+    model = TemplateModel(in_shape=(64, 64), out_size=5)
 
     # Print model summary
     print(f"Model: {model.__class__.__name__}")
     print(f"Parameters: {model.count_parameters():,}")
-    print(f"Default config: {model.get_default_config()}")
 
     # Test forward pass
-    dummy_input = torch.randn(2, 1, 500, 500)
+    dummy_input = torch.randn(2, 1, 64, 64)
     output = model(dummy_input)
     print(f"Input shape: {dummy_input.shape}")
     print(f"Output shape: {output.shape}")

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
@@ -76,13 +75,61 @@ class BaseModel(nn.Module, ABC):
         Forward pass of the model.
 
         Args:
-            x: Input tensor of shape (B, C, H, W)
+            x: Input tensor of shape (B, C, *spatial_dims)
+               - 1D: (B, C, L)
+               - 2D: (B, C, H, W)
+               - 3D: (B, C, D, H, W)
 
         Returns:
             Output tensor of shape (B, out_size)
         """
         pass
 
+    def validate_input_shape(self, x: torch.Tensor) -> None:
+        """
+        Validate input tensor shape against model's expected shape.
+
+        Call this at the start of forward() for explicit shape contract enforcement.
+        Provides clear, actionable error messages instead of cryptic Conv layer errors.
+
+        Args:
+            x: Input tensor to validate
+
+        Raises:
+            ValueError: If shape doesn't match expected dimensions
+
+        Example:
+            def forward(self, x):
+                self.validate_input_shape(x)  # Optional but recommended
+                return self.model(x)
+        """
+        expected_ndim = len(self.in_shape) + 2  # +2 for (batch, channel)
+
+        if x.ndim != expected_ndim:
+            dim_names = {
+                3: "1D (B, C, L)",
+                4: "2D (B, C, H, W)",
+                5: "3D (B, C, D, H, W)",
+            }
+            expected_name = dim_names.get(expected_ndim, f"{expected_ndim}D")
+            actual_name = dim_names.get(x.ndim, f"{x.ndim}D")
+            raise ValueError(
+                f"Input shape mismatch: model expects {expected_name} input, "
+                f"got {actual_name} with shape {tuple(x.shape)}.\n"
+                f"Expected in_shape: {self.in_shape} -> input should be (B, C, {', '.join(map(str, self.in_shape))})\n"
+                f"Hint: Check your data preprocessing - you may need to add/remove dimensions."
+            )
+
+        # Validate spatial dimensions match
+        spatial_dims = tuple(x.shape[2:])  # Skip batch and channel
+        if spatial_dims != tuple(self.in_shape):
+            raise ValueError(
+                f"Spatial dimension mismatch: model expects {self.in_shape}, "
+                f"got {spatial_dims}.\n"
+                f"Full input shape: {tuple(x.shape)} (B={x.shape[0]}, C={x.shape[1]})\n"
+                f"Hint: Ensure your data dimensions match the model's in_shape."
+            )
+
     def count_parameters(self, trainable_only: bool = True) -> int:
         """
         Count the number of parameters in the model.

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