wavedl 1.2.0__tar.gz → 1.3.0__tar.gz

{wavedl-1.2.0/src/wavedl.egg-info → wavedl-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.2.0
+Version: 1.3.0
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -43,7 +43,7 @@ Provides-Extra: onnx
 Requires-Dist: onnx>=1.14.0; extra == "onnx"
 Requires-Dist: onnxruntime>=1.15.0; extra == "onnx"
 Provides-Extra: compile
-Requires-Dist: triton; extra == "compile"
+Requires-Dist: triton; sys_platform == "linux" and extra == "compile"
 Provides-Extra: hpo
 Requires-Dist: optuna>=3.0.0; extra == "hpo"
 Provides-Extra: all
@@ -53,7 +53,7 @@ Requires-Dist: ruff>=0.8.0; extra == "all"
 Requires-Dist: pre-commit>=3.5.0; extra == "all"
 Requires-Dist: onnx>=1.14.0; extra == "all"
 Requires-Dist: onnxruntime>=1.15.0; extra == "all"
-Requires-Dist: triton; extra == "all"
+Requires-Dist: triton; sys_platform == "linux" and extra == "all"
 Requires-Dist: optuna>=3.0.0; extra == "all"
 
 <div align="center">
@@ -211,40 +211,43 @@ Deploy models anywhere:
 ### Installation
 
 ```bash
-git clone https://github.com/ductho-le/WaveDL.git
-cd WaveDL
+# Install from PyPI (recommended)
+pip install wavedl
+
+# Or install with all extras (ONNX export, HPO, dev tools)
+pip install wavedl[all]
+```
 
-# Basic install (training + inference)
-pip install -e .
+#### From Source (for development)
 
-# Full install (adds ONNX export, torch.compile, HPO, dev tools)
-pip install -e ".[all]"
+```bash
+git clone https://github.com/ductho-le/WaveDL.git
+cd WaveDL
+pip install -e ".[dev]"
 ```
 
 > [!NOTE]
-> Dependencies are managed in `pyproject.toml`. Python 3.11+ required.
->
-> For development setup (running tests, contributing), see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
+> Python 3.11+ required. For development setup, see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
 
 ### Quick Start
 
 > [!TIP]
 > In all examples below, replace `<...>` placeholders with your values. See [Configuration](#️-configuration) for defaults and options.
 
-#### Option 1: Using the Helper Script (Recommended for HPC)
+#### Option 1: Using wavedl-hpc (Recommended for HPC)
 
-The `run_training.sh` wrapper automatically configures the environment for HPC systems:
+The `wavedl-hpc` command automatically configures the environment for HPC systems:
 
 ```bash
-# Make executable (first time only)
-chmod +x run_training.sh
-
 # Basic training (auto-detects available GPUs)
-./run_training.sh --model <model_name> --data_path <train_data> --batch_size <number> --output_dir <output_folder>
+wavedl-hpc --model <model_name> --data_path <train_data> --batch_size <number> --output_dir <output_folder>
 
 # Detailed configuration
-./run_training.sh --model <model_name> --data_path <train_data> --batch_size <number> \
+wavedl-hpc --model <model_name> --data_path <train_data> --batch_size <number> \
   --lr <number> --epochs <number> --patience <number> --compile --output_dir <output_folder>
+
+# Specify GPU count explicitly
+wavedl-hpc --num_gpus 4 --model cnn --data_path train.npz --output_dir results
 ```
 
 #### Option 2: Direct Accelerate Launch
@@ -261,13 +264,13 @@ accelerate launch -m wavedl.train --model <model_name> --data_path <train_data>
 accelerate launch -m wavedl.train --model <model_name> --data_path <train_data> --output_dir <output_folder> --fresh
 
 # List available models
-python -m wavedl.train --list_models
+wavedl-train --list_models
 ```
 
 > [!TIP]
 > **Auto-Resume**: If training crashes or is interrupted, simply re-run with the same `--output_dir`. The framework automatically detects incomplete training and resumes from the last checkpoint. Use `--fresh` to force a fresh start.
 >
-> **GPU Auto-Detection**: By default, `run_training.sh` automatically detects available GPUs using `nvidia-smi`. Set `NUM_GPUS` to override this behavior.
+> **GPU Auto-Detection**: `wavedl-hpc` automatically detects available GPUs using `nvidia-smi`. Use `--num_gpus` to override.
 
 ### Testing & Inference
 
@@ -299,6 +302,56 @@ python -m wavedl.test --checkpoint <checkpoint_folder> --data_path <test_data> \
 > [!NOTE]
 > `wavedl.test` auto-detects the model architecture from checkpoint metadata. If unavailable, it falls back to folder name parsing. Use `--model` to override if needed.
 
+### Adding Custom Models
+
+<details>
+<summary><b>Creating Your Own Architecture</b></summary>
+
+**Requirements** (your model must):
+1. Inherit from `BaseModel`
+2. Accept `in_channels`, `num_outputs`, `input_shape` in `__init__`
+3. Return a tensor of shape `(batch, num_outputs)` from `forward()`
+
+---
+
+**Step 1: Create `my_model.py`**
+
+```python
+import torch.nn as nn
+import torch.nn.functional as F
+from wavedl.models import BaseModel, register_model
+
+@register_model("my_model")  # This name is used with --model flag
+class MyModel(BaseModel):
+    def __init__(self, in_channels, num_outputs, input_shape):
+        # in_channels: number of input channels (auto-detected from data)
+        # num_outputs: number of parameters to predict (auto-detected from data)
+        # input_shape: spatial dimensions, e.g., (128,) or (64, 64) or (32, 32, 32)
+        super().__init__(in_channels, num_outputs, input_shape)
+
+        # Define your layers (this is just an example)
+        self.conv1 = nn.Conv2d(in_channels, 64, 3, padding=1)
+        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
+        self.fc = nn.Linear(128, num_outputs)
+
+    def forward(self, x):
+        # Input x has shape: (batch, in_channels, *input_shape)
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = x.mean(dim=[-2, -1])  # Global average pooling
+        return self.fc(x)  # Output shape: (batch, num_outputs)
+```
+
+**Step 2: Train**
+
+```bash
+wavedl-hpc --import my_model --model my_model --data_path train.npz
+```
+
+WaveDL handles everything else: training loop, logging, checkpoints, multi-GPU, early stopping, etc.
+
+</details>
+
 ---
 
 ## 📁 Project Structure
@@ -311,6 +364,7 @@ WaveDL/
 │       ├── train.py           # Training entry point
 │       ├── test.py            # Testing & inference script
 │       ├── hpo.py             # Hyperparameter optimization
+│       ├── hpc.py             # HPC distributed training launcher
 │       │
 │       ├── models/            # Model architectures
 │       │   ├── registry.py    # Model factory (@register_model)
@@ -332,7 +386,6 @@ WaveDL/
 │           ├── schedulers.py  # LR scheduler factory
 │           └── config.py      # YAML configuration support
 │
-├── run_training.sh            # HPC helper script
 ├── configs/                   # YAML config templates
 ├── examples/                  # Ready-to-run examples
 ├── notebooks/                 # Jupyter notebooks
@@ -347,12 +400,12 @@ WaveDL/
 ## ⚙️ Configuration
 
 > [!NOTE]
-> All configuration options below work with **both** `run_training.sh` and direct `accelerate launch`. The wrapper script passes all arguments directly to `train.py`.
+> All configuration options below work with **both** `wavedl-hpc` and direct `accelerate launch`. The wrapper script passes all arguments directly to `train.py`.
 >
 > **Examples:**
 > ```bash
-> # Using run_training.sh
-> ./run_training.sh --model cnn --batch_size 256 --lr 5e-4 --compile
+> # Using wavedl-hpc
+> wavedl-hpc --model cnn --batch_size 256 --lr 5e-4 --compile
 >
 > # Using accelerate launch directly
 > accelerate launch -m wavedl.train --model cnn --batch_size 256 --lr 5e-4 --compile
@@ -395,6 +448,7 @@ WaveDL/
 | Argument | Default | Description |
 |----------|---------|-------------|
 | `--model` | `cnn` | Model architecture |
+| `--import` | - | Python modules to import (for custom models) |
 | `--batch_size` | `128` | Per-GPU batch size |
 | `--lr` | `1e-3` | Learning rate |
 | `--epochs` | `1000` | Maximum epochs |
@@ -434,7 +488,7 @@ WaveDL/
 </details>
 
 <details>
-<summary><b>Environment Variables (run_training.sh)</b></summary>
+<summary><b>Environment Variables (wavedl-hpc)</b></summary>
 
 | Variable | Default | Description |
 |----------|---------|-------------|
@@ -527,15 +581,15 @@ For robust model evaluation, simply add the `--cv` flag:
 
 ```bash
 # 5-fold cross-validation (works with both methods!)
-./run_training.sh --model cnn --cv 5 --data_path train_data.npz
+wavedl-hpc --model cnn --cv 5 --data_path train_data.npz
 # OR
 accelerate launch -m wavedl.train --model cnn --cv 5 --data_path train_data.npz
 
 # Stratified CV (recommended for unbalanced data)
-./run_training.sh --model cnn --cv 5 --cv_stratify --loss huber --epochs 100
+wavedl-hpc --model cnn --cv 5 --cv_stratify --loss huber --epochs 100
 
 # Full configuration
-./run_training.sh --model cnn --cv 5 --cv_stratify \
+wavedl-hpc --model cnn --cv 5 --cv_stratify \
     --loss huber --optimizer adamw --scheduler cosine \
     --output_dir ./cv_results
 ```

{wavedl-1.2.0 → wavedl-1.3.0}/README.md

@@ -153,40 +153,43 @@ Deploy models anywhere:
 ### Installation
 
 ```bash
-git clone https://github.com/ductho-le/WaveDL.git
-cd WaveDL
+# Install from PyPI (recommended)
+pip install wavedl
+
+# Or install with all extras (ONNX export, HPO, dev tools)
+pip install wavedl[all]
+```
 
-# Basic install (training + inference)
-pip install -e .
+#### From Source (for development)
 
-# Full install (adds ONNX export, torch.compile, HPO, dev tools)
-pip install -e ".[all]"
+```bash
+git clone https://github.com/ductho-le/WaveDL.git
+cd WaveDL
+pip install -e ".[dev]"
 ```
 
 > [!NOTE]
-> Dependencies are managed in `pyproject.toml`. Python 3.11+ required.
->
-> For development setup (running tests, contributing), see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
+> Python 3.11+ required. For development setup, see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
 
 ### Quick Start
 
 > [!TIP]
 > In all examples below, replace `<...>` placeholders with your values. See [Configuration](#️-configuration) for defaults and options.
 
-#### Option 1: Using the Helper Script (Recommended for HPC)
+#### Option 1: Using wavedl-hpc (Recommended for HPC)
 
-The `run_training.sh` wrapper automatically configures the environment for HPC systems:
+The `wavedl-hpc` command automatically configures the environment for HPC systems:
 
 ```bash
-# Make executable (first time only)
-chmod +x run_training.sh
-
 # Basic training (auto-detects available GPUs)
-./run_training.sh --model <model_name> --data_path <train_data> --batch_size <number> --output_dir <output_folder>
+wavedl-hpc --model <model_name> --data_path <train_data> --batch_size <number> --output_dir <output_folder>
 
 # Detailed configuration
-./run_training.sh --model <model_name> --data_path <train_data> --batch_size <number> \
+wavedl-hpc --model <model_name> --data_path <train_data> --batch_size <number> \
   --lr <number> --epochs <number> --patience <number> --compile --output_dir <output_folder>
+
+# Specify GPU count explicitly
+wavedl-hpc --num_gpus 4 --model cnn --data_path train.npz --output_dir results
 ```
 
 #### Option 2: Direct Accelerate Launch
@@ -203,13 +206,13 @@ accelerate launch -m wavedl.train --model <model_name> --data_path <train_data>
 accelerate launch -m wavedl.train --model <model_name> --data_path <train_data> --output_dir <output_folder> --fresh
 
 # List available models
-python -m wavedl.train --list_models
+wavedl-train --list_models
 ```
 
 > [!TIP]
 > **Auto-Resume**: If training crashes or is interrupted, simply re-run with the same `--output_dir`. The framework automatically detects incomplete training and resumes from the last checkpoint. Use `--fresh` to force a fresh start.
 >
-> **GPU Auto-Detection**: By default, `run_training.sh` automatically detects available GPUs using `nvidia-smi`. Set `NUM_GPUS` to override this behavior.
+> **GPU Auto-Detection**: `wavedl-hpc` automatically detects available GPUs using `nvidia-smi`. Use `--num_gpus` to override.
 
 ### Testing & Inference
 
@@ -241,6 +244,56 @@ python -m wavedl.test --checkpoint <checkpoint_folder> --data_path <test_data> \
 > [!NOTE]
 > `wavedl.test` auto-detects the model architecture from checkpoint metadata. If unavailable, it falls back to folder name parsing. Use `--model` to override if needed.
 
+### Adding Custom Models
+
+<details>
+<summary><b>Creating Your Own Architecture</b></summary>
+
+**Requirements** (your model must):
+1. Inherit from `BaseModel`
+2. Accept `in_channels`, `num_outputs`, `input_shape` in `__init__`
+3. Return a tensor of shape `(batch, num_outputs)` from `forward()`
+
+---
+
+**Step 1: Create `my_model.py`**
+
+```python
+import torch.nn as nn
+import torch.nn.functional as F
+from wavedl.models import BaseModel, register_model
+
+@register_model("my_model")  # This name is used with --model flag
+class MyModel(BaseModel):
+    def __init__(self, in_channels, num_outputs, input_shape):
+        # in_channels: number of input channels (auto-detected from data)
+        # num_outputs: number of parameters to predict (auto-detected from data)
+        # input_shape: spatial dimensions, e.g., (128,) or (64, 64) or (32, 32, 32)
+        super().__init__(in_channels, num_outputs, input_shape)
+
+        # Define your layers (this is just an example)
+        self.conv1 = nn.Conv2d(in_channels, 64, 3, padding=1)
+        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
+        self.fc = nn.Linear(128, num_outputs)
+
+    def forward(self, x):
+        # Input x has shape: (batch, in_channels, *input_shape)
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = x.mean(dim=[-2, -1])  # Global average pooling
+        return self.fc(x)  # Output shape: (batch, num_outputs)
+```
+
+**Step 2: Train**
+
+```bash
+wavedl-hpc --import my_model --model my_model --data_path train.npz
+```
+
+WaveDL handles everything else: training loop, logging, checkpoints, multi-GPU, early stopping, etc.
+
+</details>
+
 ---
 
 ## 📁 Project Structure
@@ -253,6 +306,7 @@ WaveDL/
 │       ├── train.py           # Training entry point
 │       ├── test.py            # Testing & inference script
 │       ├── hpo.py             # Hyperparameter optimization
+│       ├── hpc.py             # HPC distributed training launcher
 │       │
 │       ├── models/            # Model architectures
 │       │   ├── registry.py    # Model factory (@register_model)
@@ -274,7 +328,6 @@ WaveDL/
 │           ├── schedulers.py  # LR scheduler factory
 │           └── config.py      # YAML configuration support
 │
-├── run_training.sh            # HPC helper script
 ├── configs/                   # YAML config templates
 ├── examples/                  # Ready-to-run examples
 ├── notebooks/                 # Jupyter notebooks
@@ -289,12 +342,12 @@ WaveDL/
 ## ⚙️ Configuration
 
 > [!NOTE]
-> All configuration options below work with **both** `run_training.sh` and direct `accelerate launch`. The wrapper script passes all arguments directly to `train.py`.
+> All configuration options below work with **both** `wavedl-hpc` and direct `accelerate launch`. The wrapper script passes all arguments directly to `train.py`.
 >
 > **Examples:**
 > ```bash
-> # Using run_training.sh
-> ./run_training.sh --model cnn --batch_size 256 --lr 5e-4 --compile
+> # Using wavedl-hpc
+> wavedl-hpc --model cnn --batch_size 256 --lr 5e-4 --compile
 >
 > # Using accelerate launch directly
 > accelerate launch -m wavedl.train --model cnn --batch_size 256 --lr 5e-4 --compile
@@ -337,6 +390,7 @@ WaveDL/
 | Argument | Default | Description |
 |----------|---------|-------------|
 | `--model` | `cnn` | Model architecture |
+| `--import` | - | Python modules to import (for custom models) |
 | `--batch_size` | `128` | Per-GPU batch size |
 | `--lr` | `1e-3` | Learning rate |
 | `--epochs` | `1000` | Maximum epochs |
@@ -376,7 +430,7 @@ WaveDL/
 </details>
 
 <details>
-<summary><b>Environment Variables (run_training.sh)</b></summary>
+<summary><b>Environment Variables (wavedl-hpc)</b></summary>
 
 | Variable | Default | Description |
 |----------|---------|-------------|
@@ -469,15 +523,15 @@ For robust model evaluation, simply add the `--cv` flag:
 
 ```bash
 # 5-fold cross-validation (works with both methods!)
-./run_training.sh --model cnn --cv 5 --data_path train_data.npz
+wavedl-hpc --model cnn --cv 5 --data_path train_data.npz
 # OR
 accelerate launch -m wavedl.train --model cnn --cv 5 --data_path train_data.npz
 
 # Stratified CV (recommended for unbalanced data)
-./run_training.sh --model cnn --cv 5 --cv_stratify --loss huber --epochs 100
+wavedl-hpc --model cnn --cv 5 --cv_stratify --loss huber --epochs 100
 
 # Full configuration
-./run_training.sh --model cnn --cv 5 --cv_stratify \
+wavedl-hpc --model cnn --cv 5 --cv_stratify \
     --loss huber --optimizer adamw --scheduler cosine \
     --output_dir ./cv_results
 ```

{wavedl-1.2.0 → wavedl-1.3.0}/pyproject.toml

@@ -67,12 +67,12 @@ dependencies = [
 [project.optional-dependencies]
 dev = ["pytest>=7.0.0", "pytest-xdist>=3.5.0", "ruff>=0.8.0", "pre-commit>=3.5.0"]
 onnx = ["onnx>=1.14.0", "onnxruntime>=1.15.0"]
-compile = ["triton"]  # Version resolved by PyTorch compatibility
+compile = ["triton; sys_platform == 'linux'"]  # Linux-only, enables torch.compile
 hpo = ["optuna>=3.0.0"]  # Hyperparameter optimization
 all = [
     "pytest>=7.0.0", "pytest-xdist>=3.5.0", "ruff>=0.8.0", "pre-commit>=3.5.0",
     "onnx>=1.14.0", "onnxruntime>=1.15.0",
-    "triton",
+    "triton; sys_platform == 'linux'",
     "optuna>=3.0.0",
 ]
 
@@ -80,6 +80,7 @@ all = [
 wavedl-train = "wavedl.train:main"
 wavedl-test = "wavedl.test:main"
 wavedl-hpo = "wavedl.hpo:main"
+wavedl-hpc = "wavedl.hpc:main"
 
 [project.urls]
 Homepage = "https://github.com/ductho-le/WaveDL"

{wavedl-1.2.0 → wavedl-1.3.0}/src/wavedl/__init__.py

@@ -18,7 +18,7 @@ For inference:
     # or: python -m wavedl.test --checkpoint best_checkpoint --data_path test.npz
 """
 
-__version__ = "1.2.0"
+__version__ = "1.3.0"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

wavedl-1.3.0/src/wavedl/hpc.py

@@ -0,0 +1,243 @@
+#!/usr/bin/env python
+"""
+WaveDL HPC Training Launcher.
+
+This module provides a Python-based HPC training launcher that wraps accelerate
+for distributed training on High-Performance Computing clusters.
+
+Usage:
+    wavedl-hpc --model cnn --data_path train.npz --num_gpus 4
+
+Example SLURM script:
+    #!/bin/bash
+    #SBATCH --nodes=1
+    #SBATCH --gpus-per-node=4
+    #SBATCH --time=12:00:00
+
+    wavedl-hpc --model cnn --data_path /scratch/data.npz --compile
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+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")
+        return 1
+
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--list-gpus"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        gpu_count = len(result.stdout.strip().split("\n"))
+        if gpu_count > 0:
+            print(f"Auto-detected {gpu_count} GPU(s)")
+            return gpu_count
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        pass
+
+    print("Warning: Could not detect GPUs, defaulting to 1")
+    return 1
+
+
+def setup_hpc_environment() -> None:
+    """Configure environment variables for HPC systems.
+
+    Handles restricted home directories (e.g., Compute Canada) and
+    offline logging configurations.
+    """
+    # Use SLURM_TMPDIR if available, otherwise system temp
+    tmpdir = os.environ.get("SLURM_TMPDIR", tempfile.gettempdir())
+
+    # Configure directories for systems with restricted home directories
+    os.environ.setdefault("MPLCONFIGDIR", f"{tmpdir}/matplotlib")
+    os.environ.setdefault("XDG_CACHE_HOME", f"{tmpdir}/.cache")
+
+    # Ensure matplotlib config dir exists
+    Path(os.environ["MPLCONFIGDIR"]).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"{tmpdir}/wandb")
+    os.environ.setdefault("WANDB_CACHE_DIR", f"{tmpdir}/wandb_cache")
+    os.environ.setdefault("WANDB_CONFIG_DIR", f"{tmpdir}/wandb_config")
+
+    # Suppress non-critical warnings
+    os.environ.setdefault(
+        "PYTHONWARNINGS",
+        "ignore::UserWarning,ignore::FutureWarning,ignore::DeprecationWarning",
+    )
+
+
+def parse_args() -> tuple[argparse.Namespace, list[str]]:
+    """Parse HPC-specific arguments, pass remaining to wavedl.train."""
+    parser = argparse.ArgumentParser(
+        description="WaveDL HPC Training Launcher",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Basic training with auto-detected GPUs
+  wavedl-hpc --model cnn --data_path train.npz --epochs 100
+
+  # Specify GPU count and mixed precision
+  wavedl-hpc --model cnn --data_path train.npz --num_gpus 4 --mixed_precision bf16
+
+  # Full configuration
+  wavedl-hpc --model resnet18 --data_path train.npz --num_gpus 8 \\
+             --batch_size 256 --lr 1e-3 --compile --output_dir ./results
+
+Environment Variables:
+  WANDB_MODE          WandB mode: offline|online (default: offline)
+  SLURM_TMPDIR        Temp directory for HPC systems
+""",
+    )
+
+    # HPC-specific arguments
+    parser.add_argument(
+        "--num_gpus",
+        type=int,
+        default=None,
+        help="Number of GPUs to use (default: auto-detect)",
+    )
+    parser.add_argument(
+        "--num_machines",
+        type=int,
+        default=1,
+        help="Number of machines for multi-node training (default: 1)",
+    )
+    parser.add_argument(
+        "--machine_rank",
+        type=int,
+        default=0,
+        help="Rank of this machine in multi-node setup (default: 0)",
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        choices=["bf16", "fp16", "no"],
+        default="bf16",
+        help="Mixed precision mode (default: bf16)",
+    )
+    parser.add_argument(
+        "--dynamo_backend",
+        type=str,
+        default="no",
+        help="PyTorch dynamo backend (default: no)",
+    )
+
+    # Parse known args, pass rest to wavedl.train
+    args, remaining = parser.parse_known_args()
+    return args, remaining
+
+
+def print_summary(exit_code: int, wandb_mode: str, wandb_dir: str) -> None:
+    """Print post-training summary and instructions."""
+    print()
+    print("=" * 50)
+
+    if exit_code == 0:
+        print("✅ Training completed successfully!")
+        print("=" * 50)
+
+        if wandb_mode == "offline":
+            print()
+            print("📊 WandB Sync Instructions:")
+            print("   From the login node, run:")
+            print(f"   wandb sync {wandb_dir}/wandb/offline-run-*")
+            print()
+            print("   This will upload your training logs to wandb.ai")
+    else:
+        print(f"❌ Training failed with exit code: {exit_code}")
+        print("=" * 50)
+        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()
+
+    print("=" * 50)
+    print()
+
+
+def main() -> int:
+    """Main entry point for wavedl-hpc command."""
+    # Parse arguments
+    args, train_args = parse_args()
+
+    # Setup HPC environment
+    setup_hpc_environment()
+
+    # Auto-detect GPUs if not specified
+    num_gpus = args.num_gpus if args.num_gpus is not None else detect_gpus()
+
+    # Build accelerate launch command
+    cmd = [
+        sys.executable,
+        "-m",
+        "accelerate.commands.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
+
+    # Create output directory if specified
+    for i, arg in enumerate(train_args):
+        if arg == "--output_dir" and i + 1 < len(train_args):
+            Path(train_args[i + 1]).mkdir(parents=True, exist_ok=True)
+            break
+        if arg.startswith("--output_dir="):
+            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)
+        exit_code = result.returncode
+    except KeyboardInterrupt:
+        print("\n\n⚠️  Training interrupted by user")
+        exit_code = 130
+
+    # Print summary
+    print_summary(
+        exit_code,
+        os.environ.get("WANDB_MODE", "offline"),
+        os.environ.get("WANDB_DIR", "/tmp/wandb"),
+    )
+
+    return exit_code
+
+
+if __name__ == "__main__":
+    sys.exit(main())

{wavedl-1.2.0 → wavedl-1.3.0}/src/wavedl/hpo.py

@@ -5,16 +5,16 @@ Automated hyperparameter search for finding optimal training configurations.
 
 Usage:
     # Basic HPO (50 trials)
-    python hpo.py --data_path train.npz --n_trials 50
+    wavedl-hpo --data_path train.npz --n_trials 50
 
     # Quick search (fewer parameters)
-    python hpo.py --data_path train.npz --n_trials 30 --quick
+    wavedl-hpo --data_path train.npz --n_trials 30 --quick
 
     # Full search with specific models
-    python hpo.py --data_path train.npz --n_trials 100 --models cnn resnet18 efficientnet_b0
+    wavedl-hpo --data_path train.npz --n_trials 100 --models cnn resnet18 efficientnet_b0
 
     # Parallel trials on multiple GPUs
-    python hpo.py --data_path train.npz --n_trials 100 --n_jobs 4
+    wavedl-hpo --data_path train.npz --n_trials 100 --n_jobs 4
 
 Author: Ductho Le (ductho.le@outlook.com)
 """
@@ -205,9 +205,9 @@ def main():
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog="""
 Examples:
-    python hpo.py --data_path train.npz --n_trials 50
-    python hpo.py --data_path train.npz --n_trials 30 --quick
-    python hpo.py --data_path train.npz --n_trials 100 --models cnn resnet18
+    wavedl-hpo --data_path train.npz --n_trials 50
+    wavedl-hpo --data_path train.npz --n_trials 30 --quick
+    wavedl-hpo --data_path train.npz --n_trials 100 --models cnn resnet18
         """,
     )
 
@@ -355,7 +355,7 @@ Examples:
     print("\n" + "=" * 60)
     print("TO TRAIN WITH BEST PARAMETERS:")
     print("=" * 60)
-    cmd_parts = ["accelerate launch train.py"]
+    cmd_parts = ["accelerate launch -m wavedl.train"]
     cmd_parts.append(f"--data_path {args.data_path}")
     for key, value in study.best_params.items():
         cmd_parts.append(f"--{key} {value}")

{wavedl-1.2.0 → wavedl-1.3.0}/src/wavedl/models/_template.py

@@ -11,7 +11,7 @@ Steps to Add a New 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 train.py --model your_model --wandb
+    5. Run: accelerate launch -m wavedl.train --model your_model --wandb
 
 Author: Ductho Le (ductho.le@outlook.com)
 Version: 1.0.0

{wavedl-1.2.0 → wavedl-1.3.0}/src/wavedl/test.py

@@ -13,14 +13,14 @@ Production-grade inference script for evaluating trained WaveDL models:
 
 Usage:
     # Basic inference
-    python test.py --checkpoint ./best_checkpoint --data_path test_data.npz
+    wavedl-test --checkpoint ./best_checkpoint --data_path test_data.npz
 
     # With visualization and detailed output
-    python test.py --checkpoint ./best_checkpoint --data_path test_data.npz \\
+    wavedl-test --checkpoint ./best_checkpoint --data_path test_data.npz \\
         --plot --plot_format png pdf --output_dir ./test_results --save_predictions
 
     # Export model to ONNX for deployment
-    python test.py --checkpoint ./best_checkpoint --data_path test_data.npz \\
+    wavedl-test --checkpoint ./best_checkpoint --data_path test_data.npz \\
         --export onnx --export_path model.onnx
 
 Author: Ductho Le (ductho.le@outlook.com)

{wavedl-1.2.0 → wavedl-1.3.0}/src/wavedl/train.py

@@ -12,26 +12,25 @@ A modular training framework for wave-based inverse problems and regression:
   6. Deep Observability: WandB integration with scatter analysis
 
 Usage:
-    # Recommended: Using the HPC helper script
-    ./run_training.sh --model cnn --batch_size 128 --wandb
+    # Recommended: Using the HPC launcher
+    wavedl-hpc --model cnn --batch_size 128 --wandb
 
     # Or with direct accelerate launch
-    accelerate launch train.py --model cnn --batch_size 128 --wandb
+    accelerate launch -m wavedl.train --model cnn --batch_size 128 --wandb
 
     # Multi-GPU with explicit config
-    accelerate launch --num_processes=4 --mixed_precision=bf16 \
-        train.py --model cnn --wandb --project_name "MyProject"
+    wavedl-hpc --num_gpus 4 --mixed_precision bf16 --model cnn --wandb
 
     # Resume from checkpoint
-    accelerate launch train.py --model cnn --resume best_checkpoint --wandb
+    accelerate launch -m wavedl.train --model cnn --resume best_checkpoint --wandb
 
     # List available models
-    python train.py --list_models
+    wavedl-train --list_models
 
 Note:
-    For HPC clusters (Compute Canada, etc.), use run_training.sh which handles
+    For HPC clusters (Compute Canada, etc.), use wavedl-hpc which handles
     environment configuration automatically. Mixed precision is controlled via
-    --precision flag (default: bf16).
+    --mixed_precision flag (default: bf16).
 
 Author: Ductho Le (ductho.le@outlook.com)
 """
@@ -122,6 +121,14 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--list_models", action="store_true", help="List all available models and exit"
     )
+    parser.add_argument(
+        "--import",
+        dest="import_modules",
+        type=str,
+        nargs="+",
+        default=[],
+        help="Python modules to import before training (for custom models)",
+    )
 
     # Configuration File
     parser.add_argument(
@@ -314,6 +321,37 @@ def parse_args() -> argparse.Namespace:
 def main():
     args, parser = parse_args()
 
+    # Import custom model modules if specified
+    if args.import_modules:
+        import importlib
+        import sys
+
+        for module_name in args.import_modules:
+            try:
+                # Handle both module names (my_model) and file paths (./my_model.py)
+                if module_name.endswith(".py"):
+                    # Import from file path
+                    import importlib.util
+
+                    spec = importlib.util.spec_from_file_location(
+                        "custom_module", module_name
+                    )
+                    if spec and spec.loader:
+                        module = importlib.util.module_from_spec(spec)
+                        sys.modules["custom_module"] = module
+                        spec.loader.exec_module(module)
+                        print(f"✓ Imported custom module from: {module_name}")
+                else:
+                    # Import as regular module
+                    importlib.import_module(module_name)
+                    print(f"✓ Imported module: {module_name}")
+            except ImportError as e:
+                print(f"✗ Failed to import '{module_name}': {e}", file=sys.stderr)
+                print(
+                    "  Make sure the module is in your Python path or current directory."
+                )
+                sys.exit(1)
+
     # Handle --list_models flag
     if args.list_models:
         print("Available models:")

{wavedl-1.2.0 → wavedl-1.3.0/src/wavedl.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.2.0
+Version: 1.3.0
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -43,7 +43,7 @@ Provides-Extra: onnx
 Requires-Dist: onnx>=1.14.0; extra == "onnx"
 Requires-Dist: onnxruntime>=1.15.0; extra == "onnx"
 Provides-Extra: compile
-Requires-Dist: triton; extra == "compile"
+Requires-Dist: triton; sys_platform == "linux" and extra == "compile"
 Provides-Extra: hpo
 Requires-Dist: optuna>=3.0.0; extra == "hpo"
 Provides-Extra: all
@@ -53,7 +53,7 @@ Requires-Dist: ruff>=0.8.0; extra == "all"
 Requires-Dist: pre-commit>=3.5.0; extra == "all"
 Requires-Dist: onnx>=1.14.0; extra == "all"
 Requires-Dist: onnxruntime>=1.15.0; extra == "all"
-Requires-Dist: triton; extra == "all"
+Requires-Dist: triton; sys_platform == "linux" and extra == "all"
 Requires-Dist: optuna>=3.0.0; extra == "all"
 
 <div align="center">
@@ -211,40 +211,43 @@ Deploy models anywhere:
 ### Installation
 
 ```bash
-git clone https://github.com/ductho-le/WaveDL.git
-cd WaveDL
+# Install from PyPI (recommended)
+pip install wavedl
+
+# Or install with all extras (ONNX export, HPO, dev tools)
+pip install wavedl[all]
+```
 
-# Basic install (training + inference)
-pip install -e .
+#### From Source (for development)
 
-# Full install (adds ONNX export, torch.compile, HPO, dev tools)
-pip install -e ".[all]"
+```bash
+git clone https://github.com/ductho-le/WaveDL.git
+cd WaveDL
+pip install -e ".[dev]"
 ```
 
 > [!NOTE]
-> Dependencies are managed in `pyproject.toml`. Python 3.11+ required.
->
-> For development setup (running tests, contributing), see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
+> Python 3.11+ required. For development setup, see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
 
 ### Quick Start
 
 > [!TIP]
 > In all examples below, replace `<...>` placeholders with your values. See [Configuration](#️-configuration) for defaults and options.
 
-#### Option 1: Using the Helper Script (Recommended for HPC)
+#### Option 1: Using wavedl-hpc (Recommended for HPC)
 
-The `run_training.sh` wrapper automatically configures the environment for HPC systems:
+The `wavedl-hpc` command automatically configures the environment for HPC systems:
 
 ```bash
-# Make executable (first time only)
-chmod +x run_training.sh
-
 # Basic training (auto-detects available GPUs)
-./run_training.sh --model <model_name> --data_path <train_data> --batch_size <number> --output_dir <output_folder>
+wavedl-hpc --model <model_name> --data_path <train_data> --batch_size <number> --output_dir <output_folder>
 
 # Detailed configuration
-./run_training.sh --model <model_name> --data_path <train_data> --batch_size <number> \
+wavedl-hpc --model <model_name> --data_path <train_data> --batch_size <number> \
   --lr <number> --epochs <number> --patience <number> --compile --output_dir <output_folder>
+
+# Specify GPU count explicitly
+wavedl-hpc --num_gpus 4 --model cnn --data_path train.npz --output_dir results
 ```
 
 #### Option 2: Direct Accelerate Launch
@@ -261,13 +264,13 @@ accelerate launch -m wavedl.train --model <model_name> --data_path <train_data>
 accelerate launch -m wavedl.train --model <model_name> --data_path <train_data> --output_dir <output_folder> --fresh
 
 # List available models
-python -m wavedl.train --list_models
+wavedl-train --list_models
 ```
 
 > [!TIP]
 > **Auto-Resume**: If training crashes or is interrupted, simply re-run with the same `--output_dir`. The framework automatically detects incomplete training and resumes from the last checkpoint. Use `--fresh` to force a fresh start.
 >
-> **GPU Auto-Detection**: By default, `run_training.sh` automatically detects available GPUs using `nvidia-smi`. Set `NUM_GPUS` to override this behavior.
+> **GPU Auto-Detection**: `wavedl-hpc` automatically detects available GPUs using `nvidia-smi`. Use `--num_gpus` to override.
 
 ### Testing & Inference
 
@@ -299,6 +302,56 @@ python -m wavedl.test --checkpoint <checkpoint_folder> --data_path <test_data> \
 > [!NOTE]
 > `wavedl.test` auto-detects the model architecture from checkpoint metadata. If unavailable, it falls back to folder name parsing. Use `--model` to override if needed.
 
+### Adding Custom Models
+
+<details>
+<summary><b>Creating Your Own Architecture</b></summary>
+
+**Requirements** (your model must):
+1. Inherit from `BaseModel`
+2. Accept `in_channels`, `num_outputs`, `input_shape` in `__init__`
+3. Return a tensor of shape `(batch, num_outputs)` from `forward()`
+
+---
+
+**Step 1: Create `my_model.py`**
+
+```python
+import torch.nn as nn
+import torch.nn.functional as F
+from wavedl.models import BaseModel, register_model
+
+@register_model("my_model")  # This name is used with --model flag
+class MyModel(BaseModel):
+    def __init__(self, in_channels, num_outputs, input_shape):
+        # in_channels: number of input channels (auto-detected from data)
+        # num_outputs: number of parameters to predict (auto-detected from data)
+        # input_shape: spatial dimensions, e.g., (128,) or (64, 64) or (32, 32, 32)
+        super().__init__(in_channels, num_outputs, input_shape)
+
+        # Define your layers (this is just an example)
+        self.conv1 = nn.Conv2d(in_channels, 64, 3, padding=1)
+        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
+        self.fc = nn.Linear(128, num_outputs)
+
+    def forward(self, x):
+        # Input x has shape: (batch, in_channels, *input_shape)
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = x.mean(dim=[-2, -1])  # Global average pooling
+        return self.fc(x)  # Output shape: (batch, num_outputs)
+```
+
+**Step 2: Train**
+
+```bash
+wavedl-hpc --import my_model --model my_model --data_path train.npz
+```
+
+WaveDL handles everything else: training loop, logging, checkpoints, multi-GPU, early stopping, etc.
+
+</details>
+
 ---
 
 ## 📁 Project Structure
@@ -311,6 +364,7 @@ WaveDL/
 │       ├── train.py           # Training entry point
 │       ├── test.py            # Testing & inference script
 │       ├── hpo.py             # Hyperparameter optimization
+│       ├── hpc.py             # HPC distributed training launcher
 │       │
 │       ├── models/            # Model architectures
 │       │   ├── registry.py    # Model factory (@register_model)
@@ -332,7 +386,6 @@ WaveDL/
 │           ├── schedulers.py  # LR scheduler factory
 │           └── config.py      # YAML configuration support
 │
-├── run_training.sh            # HPC helper script
 ├── configs/                   # YAML config templates
 ├── examples/                  # Ready-to-run examples
 ├── notebooks/                 # Jupyter notebooks
@@ -347,12 +400,12 @@ WaveDL/
 ## ⚙️ Configuration
 
 > [!NOTE]
-> All configuration options below work with **both** `run_training.sh` and direct `accelerate launch`. The wrapper script passes all arguments directly to `train.py`.
+> All configuration options below work with **both** `wavedl-hpc` and direct `accelerate launch`. The wrapper script passes all arguments directly to `train.py`.
 >
 > **Examples:**
 > ```bash
-> # Using run_training.sh
-> ./run_training.sh --model cnn --batch_size 256 --lr 5e-4 --compile
+> # Using wavedl-hpc
+> wavedl-hpc --model cnn --batch_size 256 --lr 5e-4 --compile
 >
 > # Using accelerate launch directly
 > accelerate launch -m wavedl.train --model cnn --batch_size 256 --lr 5e-4 --compile
@@ -395,6 +448,7 @@ WaveDL/
 | Argument | Default | Description |
 |----------|---------|-------------|
 | `--model` | `cnn` | Model architecture |
+| `--import` | - | Python modules to import (for custom models) |
 | `--batch_size` | `128` | Per-GPU batch size |
 | `--lr` | `1e-3` | Learning rate |
 | `--epochs` | `1000` | Maximum epochs |
@@ -434,7 +488,7 @@ WaveDL/
 </details>
 
 <details>
-<summary><b>Environment Variables (run_training.sh)</b></summary>
+<summary><b>Environment Variables (wavedl-hpc)</b></summary>
 
 | Variable | Default | Description |
 |----------|---------|-------------|
@@ -527,15 +581,15 @@ For robust model evaluation, simply add the `--cv` flag:
 
 ```bash
 # 5-fold cross-validation (works with both methods!)
-./run_training.sh --model cnn --cv 5 --data_path train_data.npz
+wavedl-hpc --model cnn --cv 5 --data_path train_data.npz
 # OR
 accelerate launch -m wavedl.train --model cnn --cv 5 --data_path train_data.npz
 
 # Stratified CV (recommended for unbalanced data)
-./run_training.sh --model cnn --cv 5 --cv_stratify --loss huber --epochs 100
+wavedl-hpc --model cnn --cv 5 --cv_stratify --loss huber --epochs 100
 
 # Full configuration
-./run_training.sh --model cnn --cv 5 --cv_stratify \
+wavedl-hpc --model cnn --cv 5 --cv_stratify \
     --loss huber --optimizer adamw --scheduler cosine \
     --output_dir ./cv_results
 ```

{wavedl-1.2.0 → wavedl-1.3.0}/src/wavedl.egg-info/SOURCES.txt

@@ -2,6 +2,7 @@ LICENSE
 README.md
 pyproject.toml
 src/wavedl/__init__.py
+src/wavedl/hpc.py
 src/wavedl/hpo.py
 src/wavedl/test.py
 src/wavedl/train.py

{wavedl-1.2.0 → wavedl-1.3.0}/src/wavedl.egg-info/entry_points.txt

@@ -1,4 +1,5 @@
 [console_scripts]
+wavedl-hpc = wavedl.hpc:main
 wavedl-hpo = wavedl.hpo:main
 wavedl-test = wavedl.test:main
 wavedl-train = wavedl.train:main

{wavedl-1.2.0 → wavedl-1.3.0}/src/wavedl.egg-info/requires.txt

@@ -19,10 +19,14 @@ ruff>=0.8.0
 pre-commit>=3.5.0
 onnx>=1.14.0
 onnxruntime>=1.15.0
-triton
 optuna>=3.0.0
 
+[all:sys_platform == "linux"]
+triton
+
 [compile]
+
+[compile:sys_platform == "linux"]
 triton
 
 [dev]