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

wavedl/{hpc.py → launcher.py}

@@ -1,12 +1,21 @@
 #!/usr/bin/env python
 """
-WaveDL HPC Training Launcher.
+WaveDL Training Launcher.
 
-This module provides a Python-based HPC training launcher that wraps accelerate
-for distributed training on High-Performance Computing clusters.
+This module provides a universal training launcher that wraps accelerate
+for distributed training. It works seamlessly on both:
+- Local machines (uses standard cache locations)
+- HPC clusters (uses local caching, offline WandB)
+
+The environment is auto-detected based on scheduler variables (SLURM, PBS, etc.)
+and home directory writability.
 
 Usage:
-    wavedl-hpc --model cnn --data_path train.npz --num_gpus 4
+    # Local machine or HPC - same command!
+    wavedl-train --model cnn --data_path train.npz --output_dir results
+
+    # Multi-GPU is automatic (uses all available GPUs)
+    wavedl-train --model resnet18 --data_path train.npz --num_gpus 4
 
 Example SLURM script:
     #!/bin/bash
@@ -14,7 +23,7 @@ Example SLURM script:
     #SBATCH --gpus-per-node=4
     #SBATCH --time=12:00:00
 
-    wavedl-hpc --model cnn --data_path /scratch/data.npz --compile
+    wavedl-train --model cnn --data_path /scratch/data.npz --compile
 
 Author: Ductho Le (ductho.le@outlook.com)
 """
@@ -53,78 +62,138 @@ def detect_gpus() -> int:
     return 1
 
 
-def setup_hpc_environment() -> None:
-    """Configure environment variables for HPC systems.
+def is_hpc_environment() -> bool:
+    """Detect if running on an HPC cluster.
+
+    Checks for:
+    1. Common HPC scheduler environment variables (SLURM, PBS, LSF, SGE, Cobalt)
+    2. Non-writable home directory (common on HPC systems)
 
-    Handles restricted home directories (e.g., Compute Canada) and
-    offline logging configurations. Always uses CWD-based TORCH_HOME
-    since compute nodes typically lack internet access.
+    Returns:
+        True if HPC environment detected, False otherwise.
     """
-    # Use CWD for cache base since HPC compute nodes typically lack internet
-    cache_base = os.getcwd()
+    # Check for common HPC scheduler environment variables
+    hpc_indicators = [
+        "SLURM_JOB_ID",  # SLURM
+        "PBS_JOBID",  # PBS/Torque
+        "LSB_JOBID",  # LSF
+        "SGE_TASK_ID",  # Sun Grid Engine
+        "COBALT_JOBID",  # Cobalt
+    ]
+    if any(var in os.environ for var in hpc_indicators):
+        return True
 
-    # TORCH_HOME always set to CWD - compute nodes need pre-cached weights
-    os.environ.setdefault("TORCH_HOME", f"{cache_base}/.torch_cache")
-    Path(os.environ["TORCH_HOME"]).mkdir(parents=True, exist_ok=True)
+    # Check if home directory is not writable (common on HPC)
+    home = os.path.expanduser("~")
+    return not os.access(home, os.W_OK)
 
-    # Triton/Inductor caches - prevents permission errors with --compile
-    # These MUST be set before any torch.compile calls
-    os.environ.setdefault("TRITON_CACHE_DIR", f"{cache_base}/.triton_cache")
-    os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", f"{cache_base}/.inductor_cache")
-    Path(os.environ["TRITON_CACHE_DIR"]).mkdir(parents=True, exist_ok=True)
-    Path(os.environ["TORCHINDUCTOR_CACHE_DIR"]).mkdir(parents=True, exist_ok=True)
 
-    # Check if home is writable for other caches
-    home = os.path.expanduser("~")
-    home_writable = os.access(home, os.W_OK)
-
-    # Other caches only if home is not writable
-    if not home_writable:
-        os.environ.setdefault("MPLCONFIGDIR", f"{cache_base}/.matplotlib")
-        os.environ.setdefault("FONTCONFIG_CACHE", f"{cache_base}/.fontconfig")
-        os.environ.setdefault("XDG_CACHE_HOME", f"{cache_base}/.cache")
-
-        # Ensure directories exist
-        for env_var in [
-            "MPLCONFIGDIR",
-            "FONTCONFIG_CACHE",
-            "XDG_CACHE_HOME",
-        ]:
-            Path(os.environ[env_var]).mkdir(parents=True, exist_ok=True)
-
-    # WandB configuration (offline by default for HPC)
-    os.environ.setdefault("WANDB_MODE", "offline")
-    os.environ.setdefault("WANDB_DIR", f"{cache_base}/.wandb")
-    os.environ.setdefault("WANDB_CACHE_DIR", f"{cache_base}/.wandb_cache")
-    os.environ.setdefault("WANDB_CONFIG_DIR", f"{cache_base}/.wandb_config")
-
-    # Suppress non-critical warnings
+def setup_environment() -> None:
+    """Configure environment for HPC or local machine.
+
+    Automatically detects the environment and configures accordingly:
+    - HPC: Uses CWD-based caching, offline WandB (compute nodes lack internet)
+    - Local: Uses standard cache locations (~/.cache), doesn't override WandB
+    """
+    is_hpc = is_hpc_environment()
+
+    if is_hpc:
+        # HPC: use CWD-based caching (compute nodes lack internet)
+        cache_base = os.getcwd()
+
+        # TORCH_HOME set to CWD - compute nodes need pre-cached weights
+        os.environ.setdefault("TORCH_HOME", f"{cache_base}/.torch_cache")
+        Path(os.environ["TORCH_HOME"]).mkdir(parents=True, exist_ok=True)
+
+        # Triton/Inductor caches - prevents permission errors with --compile
+        os.environ.setdefault("TRITON_CACHE_DIR", f"{cache_base}/.triton_cache")
+        os.environ.setdefault(
+            "TORCHINDUCTOR_CACHE_DIR", f"{cache_base}/.inductor_cache"
+        )
+        Path(os.environ["TRITON_CACHE_DIR"]).mkdir(parents=True, exist_ok=True)
+        Path(os.environ["TORCHINDUCTOR_CACHE_DIR"]).mkdir(parents=True, exist_ok=True)
+
+        # Check if home is writable for other caches
+        home = os.path.expanduser("~")
+        home_writable = os.access(home, os.W_OK)
+
+        # Other caches only if home is not writable
+        if not home_writable:
+            os.environ.setdefault("MPLCONFIGDIR", f"{cache_base}/.matplotlib")
+            os.environ.setdefault("FONTCONFIG_CACHE", f"{cache_base}/.fontconfig")
+            os.environ.setdefault("XDG_CACHE_HOME", f"{cache_base}/.cache")
+
+            for env_var in [
+                "MPLCONFIGDIR",
+                "FONTCONFIG_CACHE",
+                "XDG_CACHE_HOME",
+            ]:
+                Path(os.environ[env_var]).mkdir(parents=True, exist_ok=True)
+
+        # WandB configuration (offline by default for HPC)
+        os.environ.setdefault("WANDB_MODE", "offline")
+        os.environ.setdefault("WANDB_DIR", f"{cache_base}/.wandb")
+        os.environ.setdefault("WANDB_CACHE_DIR", f"{cache_base}/.wandb_cache")
+        os.environ.setdefault("WANDB_CONFIG_DIR", f"{cache_base}/.wandb_config")
+
+        print("🖥️  HPC environment detected - using local caching")
+    else:
+        # Local machine: use standard locations, don't override user settings
+        # TORCH_HOME defaults to ~/.cache/torch (PyTorch default)
+        # WANDB_MODE defaults to online (WandB default)
+        print("💻 Local environment detected - using standard cache locations")
+
+    # Suppress non-critical warnings (both environments)
     os.environ.setdefault(
         "PYTHONWARNINGS",
         "ignore::UserWarning,ignore::FutureWarning,ignore::DeprecationWarning",
     )
 
 
+def handle_fast_path_args() -> int | None:
+    """Handle utility flags that don't need accelerate launch.
+
+    Returns:
+        Exit code if handled (0 for success), None if should continue to full launch.
+    """
+    # --list_models: print models and exit immediately
+    if "--list_models" in sys.argv:
+        from wavedl.models import list_models
+
+        print("Available models:")
+        for name in list_models():
+            print(f"  {name}")
+        return 0
+
+    return None  # Continue to full launch
+
+
 def parse_args() -> tuple[argparse.Namespace, list[str]]:
-    """Parse HPC-specific arguments, pass remaining to wavedl.train."""
+    """Parse launcher-specific arguments, pass remaining to wavedl.train."""
     parser = argparse.ArgumentParser(
-        description="WaveDL HPC Training Launcher",
+        description="WaveDL Training Launcher (works on local machines and HPC clusters)",
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog="""
 Examples:
-  # Basic training with auto-detected GPUs
-  wavedl-hpc --model cnn --data_path train.npz --epochs 100
+  # Basic training (auto-detects GPUs and environment)
+  wavedl-train --model cnn --data_path train.npz --output_dir results
 
-  # Specify GPU count and mixed precision
-  wavedl-hpc --model cnn --data_path train.npz --num_gpus 4 --mixed_precision bf16
+  # Specify GPU count explicitly
+  wavedl-train --model cnn --data_path train.npz --num_gpus 4
 
   # Full configuration
-  wavedl-hpc --model resnet18 --data_path train.npz --num_gpus 8 \\
-             --batch_size 256 --lr 1e-3 --compile --output_dir ./results
+  wavedl-train --model resnet18 --data_path train.npz --batch_size 256 \\
+               --lr 1e-3 --epochs 100 --compile --output_dir ./results
+
+  # List available models
+  wavedl-train --list_models
 
-Environment Variables:
-  WANDB_MODE          WandB mode: offline|online (default: offline)
-  SLURM_TMPDIR        Temp directory for HPC systems
+Environment Detection:
+  The launcher automatically detects your environment:
+  - HPC (SLURM, PBS, etc.): Uses local caching, offline WandB
+  - Local machine: Uses standard cache locations (~/.cache)
+
+For full training options, see: python -m wavedl.train --help
 """,
     )
 
@@ -204,7 +273,7 @@ def print_summary(
         print("Common issues:")
         print("  - Missing data file (check --data_path)")
         print("  - Insufficient GPU memory (reduce --batch_size)")
-        print("  - Invalid model name (run: python train.py --list_models)")
+        print("  - Invalid model name (run: wavedl-train --list_models)")
         print()
 
     print("=" * 40)
@@ -212,12 +281,17 @@ def print_summary(
 
 
 def main() -> int:
-    """Main entry point for wavedl-hpc command."""
+    """Main entry point for wavedl-train command."""
+    # Fast path for utility flags (avoid accelerate launch overhead)
+    exit_code = handle_fast_path_args()
+    if exit_code is not None:
+        return exit_code
+
     # Parse arguments
     args, train_args = parse_args()
 
-    # Setup HPC environment
-    setup_hpc_environment()
+    # Setup environment (smart detection)
+    setup_environment()
 
     # Check if wavedl package is importable
     try:

wavedl/models/__init__.py

@@ -80,8 +80,24 @@ from .vit import ViTBase_, ViTSmall, ViTTiny
 # Optional timm-based models (imported conditionally)
 try:
     from .caformer import CaFormerS18, CaFormerS36, PoolFormerS12
+    from .efficientvit import (
+        EfficientViTB0,
+        EfficientViTB1,
+        EfficientViTB2,
+        EfficientViTB3,
+        EfficientViTL1,
+        EfficientViTL2,
+        EfficientViTM0,
+        EfficientViTM1,
+        EfficientViTM2,
+    )
     from .fastvit import FastViTS12, FastViTSA12, FastViTT8, FastViTT12
     from .maxvit import MaxViTBaseLarge, MaxViTSmall, MaxViTTiny
+    from .unireplknet import (
+        UniRepLKNetBaseLarge,
+        UniRepLKNetSmall,
+        UniRepLKNetTiny,
+    )
 
     _HAS_TIMM_MODELS = True
 except ImportError:
@@ -148,6 +164,15 @@ if _HAS_TIMM_MODELS:
         [
             "CaFormerS18",
             "CaFormerS36",
+            "EfficientViTB0",
+            "EfficientViTB1",
+            "EfficientViTB2",
+            "EfficientViTB3",
+            "EfficientViTL1",
+            "EfficientViTL2",
+            "EfficientViTM0",
+            "EfficientViTM1",
+            "EfficientViTM2",
             "FastViTS12",
             "FastViTSA12",
             "FastViTT8",
@@ -156,5 +181,8 @@ if _HAS_TIMM_MODELS:
             "MaxViTSmall",
             "MaxViTTiny",
             "PoolFormerS12",
+            "UniRepLKNetBaseLarge",
+            "UniRepLKNetSmall",
+            "UniRepLKNetTiny",
         ]
     )

wavedl/models/{_timm_utils.py → _pretrained_utils.py}

@@ -236,3 +236,131 @@ def adapt_input_channels(
         return new_conv
     else:
         raise NotImplementedError(f"Unsupported layer type: {type(conv_layer)}")
+
+
+def adapt_first_conv_for_single_channel(
+    module: nn.Module,
+    conv_path: str,
+    pretrained: bool = True,
+) -> None:
+    """
+    Adapt the first convolutional layer of a pretrained model for single-channel input.
+
+    This is a convenience function for torchvision-style models where the path
+    to the first conv layer is known. It modifies the model in-place.
+
+    For pretrained models, the RGB weights are averaged to create grayscale weights,
+    which provides a reasonable initialization for single-channel inputs.
+
+    Args:
+        module: The model or submodule containing the conv layer
+        conv_path: Dot-separated path to the conv layer (e.g., "conv1", "features.0.0")
+        pretrained: Whether to adapt pretrained weights by averaging RGB channels
+
+    Example:
+        >>> # For torchvision ResNet
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "conv1", pretrained=True
+        ... )
+        >>> # For torchvision ConvNeXt
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "features.0.0", pretrained=True
+        ... )
+        >>> # For torchvision DenseNet
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "features.conv0", pretrained=True
+        ... )
+    """
+    # Navigate to parent and get the conv layer
+    parts = conv_path.split(".")
+    parent = module
+    for part in parts[:-1]:
+        if part.isdigit():
+            parent = parent[int(part)]
+        else:
+            parent = getattr(parent, part)
+
+    # Get the final attribute name and the old conv
+    final_attr = parts[-1]
+    if final_attr.isdigit():
+        old_conv = parent[int(final_attr)]
+    else:
+        old_conv = getattr(parent, final_attr)
+
+    # Create and set the new conv
+    new_conv = adapt_input_channels(old_conv, new_in_channels=1, pretrained=pretrained)
+
+    if final_attr.isdigit():
+        parent[int(final_attr)] = new_conv
+    else:
+        setattr(parent, final_attr, new_conv)
+
+
+def find_and_adapt_input_convs(
+    backbone: nn.Module,
+    pretrained: bool = True,
+    adapt_all: bool = False,
+) -> int:
+    """
+    Find and adapt Conv2d layers with 3 input channels for single-channel input.
+
+    This is useful for timm-style models where the exact path to the first
+    conv layer may vary or where multiple layers need adaptation.
+
+    Args:
+        backbone: The backbone model to adapt
+        pretrained: Whether to adapt pretrained weights by averaging RGB channels
+        adapt_all: If True, adapt all Conv2d layers with 3 input channels.
+                   If False (default), only adapt the first one found.
+
+    Returns:
+        Number of layers adapted
+
+    Example:
+        >>> # For timm models (adapt first conv only)
+        >>> count = find_and_adapt_input_convs(model.backbone, pretrained=True)
+        >>> # For models with multiple input convs (e.g., FastViT)
+        >>> count = find_and_adapt_input_convs(
+        ...     model.backbone, pretrained=True, adapt_all=True
+        ... )
+    """
+    adapted_count = 0
+
+    for name, module in backbone.named_modules():
+        if not hasattr(module, "in_channels") or module.in_channels != 3:
+            continue
+
+        # Check if this is a wrapper with inner .conv attribute
+        if hasattr(module, "conv") and isinstance(module.conv, nn.Conv2d):
+            old_conv = module.conv
+            module.conv = adapt_input_channels(
+                old_conv, new_in_channels=1, pretrained=pretrained
+            )
+            adapted_count += 1
+
+        elif isinstance(module, nn.Conv2d):
+            # Direct Conv2d - need to replace it in parent
+            parts = name.split(".")
+            parent = backbone
+            for part in parts[:-1]:
+                if part.isdigit():
+                    parent = parent[int(part)]
+                else:
+                    parent = getattr(parent, part)
+
+            child_name = parts[-1]
+            new_conv = adapt_input_channels(
+                module, new_in_channels=1, pretrained=pretrained
+            )
+
+            if child_name.isdigit():
+                parent[int(child_name)] = new_conv
+            else:
+                setattr(parent, child_name, new_conv)
+
+            adapted_count += 1
+
+        if not adapt_all and adapted_count > 0:
+            break
+
+    return adapted_count

wavedl/models/base.py

@@ -15,6 +15,54 @@ import torch
 import torch.nn as nn
 
 
+# =============================================================================
+# TYPE ALIASES
+# =============================================================================
+
+# Spatial shape type aliases for model input dimensions
+SpatialShape1D = tuple[int]
+SpatialShape2D = tuple[int, int]
+SpatialShape3D = tuple[int, int, int]
+SpatialShape = SpatialShape1D | SpatialShape2D | SpatialShape3D
+
+
+# =============================================================================
+# UTILITY FUNCTIONS
+# =============================================================================
+
+
+def compute_num_groups(num_channels: int, preferred_groups: int = 32) -> int:
+    """
+    Compute valid num_groups for GroupNorm that divides num_channels evenly.
+
+    GroupNorm requires num_channels to be divisible by num_groups. This function
+    finds the largest valid divisor up to preferred_groups.
+
+    Args:
+        num_channels: Number of channels to normalize (must be positive)
+        preferred_groups: Preferred number of groups (default: 32)
+
+    Returns:
+        Valid num_groups that satisfies num_channels % num_groups == 0
+
+    Example:
+        >>> compute_num_groups(64)  # Returns 32
+        >>> compute_num_groups(48)  # Returns 16 (48 % 32 != 0)
+        >>> compute_num_groups(7)  # Returns 1 (prime number)
+    """
+    # Try preferred groups first, then common divisors
+    for groups in [preferred_groups, 16, 8, 4, 2, 1]:
+        if groups <= num_channels and num_channels % groups == 0:
+            return groups
+
+    # Fallback: find any valid divisor (always returns at least 1)
+    for groups in range(min(32, num_channels), 0, -1):
+        if num_channels % groups == 0:
+            return groups
+
+    return 1  # Always valid
+
+
 class BaseModel(nn.Module, ABC):
     """
     Abstract base class for all regression models.

wavedl/models/caformer.py

@@ -33,7 +33,7 @@ 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._pretrained_utils import build_regression_head
 from wavedl.models.base import BaseModel
 from wavedl.models.registry import register_model
 

wavedl/models/cnn.py

@@ -24,14 +24,10 @@ from typing import Any
 import torch
 import torch.nn as nn
 
-from wavedl.models.base import BaseModel
+from wavedl.models.base import BaseModel, SpatialShape, compute_num_groups
 from wavedl.models.registry import register_model
 
 
-# Type alias for spatial shapes
-SpatialShape = tuple[int] | tuple[int, int] | tuple[int, int, int]
-
-
 def _get_conv_layers(
     dim: int,
 ) -> tuple[type[nn.Module], type[nn.Module], type[nn.Module]]:
@@ -163,27 +159,6 @@ class CNN(BaseModel):
             nn.Linear(64, out_size),
         )
 
-    @staticmethod
-    def _compute_num_groups(num_channels: int, target_groups: int = 4) -> int:
-        """
-        Compute valid num_groups for GroupNorm that divides num_channels.
-
-        Finds the largest divisor of num_channels that is <= target_groups,
-        or falls back to 1 if no suitable divisor exists.
-
-        Args:
-            num_channels: Number of channels (must be positive)
-            target_groups: Desired number of groups (default: 4)
-
-        Returns:
-            Valid num_groups that satisfies num_channels % num_groups == 0
-        """
-        # Try target_groups down to 1, return first valid divisor
-        for g in range(min(target_groups, num_channels), 0, -1):
-            if num_channels % g == 0:
-                return g
-        return 1  # Fallback (always valid)
-
     def _make_conv_block(
         self, in_channels: int, out_channels: int, dropout: float = 0.0
     ) -> nn.Sequential:
@@ -198,7 +173,7 @@ class CNN(BaseModel):
         Returns:
             Sequential block: Conv → GroupNorm → LeakyReLU → MaxPool [→ Dropout]
         """
-        num_groups = self._compute_num_groups(out_channels, target_groups=4)
+        num_groups = compute_num_groups(out_channels, preferred_groups=4)
 
         layers = [
             self._Conv(in_channels, out_channels, kernel_size=3, padding=1),

wavedl/models/convnext.py

@@ -28,14 +28,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from wavedl.models.base import BaseModel
+from wavedl.models.base import BaseModel, SpatialShape
 from wavedl.models.registry import register_model
 
 
-# Type alias for spatial shapes
-SpatialShape = tuple[int] | tuple[int, int] | tuple[int, int, int]
-
-
 def _get_conv_layer(dim: int) -> type[nn.Module]:
     """Get dimension-appropriate Conv class."""
     if dim == 1:
@@ -468,20 +464,11 @@ class ConvNeXtTinyPretrained(BaseModel):
         )
 
         # Modify first conv for single-channel input
-        old_conv = self.backbone.features[0][0]
-        self.backbone.features[0][0] = 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,
+        from wavedl.models._pretrained_utils import adapt_first_conv_for_single_channel
+
+        adapt_first_conv_for_single_channel(
+            self.backbone, "features.0.0", pretrained=pretrained
         )
-        if pretrained:
-            with torch.no_grad():
-                self.backbone.features[0][0].weight = nn.Parameter(
-                    old_conv.weight.mean(dim=1, keepdim=True)
-                )
 
         if freeze_backbone:
             self._freeze_backbone()

wavedl/models/convnext_v2.py

@@ -31,20 +31,17 @@ from typing import Any
 import torch
 import torch.nn as nn
 
-from wavedl.models._timm_utils import (
+from wavedl.models._pretrained_utils import (
     LayerNormNd,
     build_regression_head,
     get_conv_layer,
     get_grn_layer,
     get_pool_layer,
 )
-from wavedl.models.base import BaseModel
+from wavedl.models.base import BaseModel, SpatialShape
 from wavedl.models.registry import register_model
 
 
-# Type alias for spatial shapes
-SpatialShape = tuple[int] | tuple[int, int] | tuple[int, int, int]
-
 __all__ = [
     "ConvNeXtV2Base",
     "ConvNeXtV2BaseLarge",
@@ -469,24 +466,11 @@ class ConvNeXtV2TinyPretrained(BaseModel):
 
     def _adapt_input_channels(self):
         """Adapt first conv layer for single-channel input."""
-        old_conv = self.backbone.features[0][0]
-        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,
-        )
+        from wavedl.models._pretrained_utils import adapt_first_conv_for_single_channel
 
-        if self.pretrained:
-            with torch.no_grad():
-                # Average RGB weights for grayscale
-                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)
-
-        self.backbone.features[0][0] = new_conv
+        adapt_first_conv_for_single_channel(
+            self.backbone, "features.0.0", pretrained=self.pretrained
+        )
 
     def _freeze_backbone(self):
         """Freeze all backbone parameters except classifier."""