wavedl 1.4.4__tar.gz → 1.4.6__tar.gz

{wavedl-1.4.4/src/wavedl.egg-info → wavedl-1.4.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.4.4
+Version: 1.4.6
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -49,7 +49,7 @@ Requires-Dist: triton>=2.0.0; sys_platform == "linux"
 
 ### A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 
-[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg?style=plastic&logo=python&logoColor=white)](https://www.python.org/downloads/)
+[![Python 3.11+](https://img.shields.io/badge/Python-3.11+-blue.svg?style=plastic&logo=python&logoColor=white)](https://www.python.org/downloads/)
 [![PyTorch 2.x](https://img.shields.io/badge/PyTorch-2.x-ee4c2c.svg?style=plastic&logo=pytorch&logoColor=white)](https://pytorch.org/)
 [![Accelerate](https://img.shields.io/badge/Accelerate-Enabled-yellow.svg?style=plastic&logo=huggingface&logoColor=white)](https://huggingface.co/docs/accelerate/)
 <br>
@@ -57,7 +57,7 @@ Requires-Dist: triton>=2.0.0; sys_platform == "linux"
 [![Lint](https://img.shields.io/github/actions/workflow/status/ductho-le/WaveDL/lint.yml?branch=main&style=plastic&logo=ruff&logoColor=white&label=Lint)](https://github.com/ductho-le/WaveDL/actions/workflows/lint.yml)
 [![Try it on Colab](https://img.shields.io/badge/Try_it_on_Colab-8E44AD?style=plastic&logo=googlecolab&logoColor=white)](https://colab.research.google.com/github/ductho-le/WaveDL/blob/main/notebooks/demo.ipynb)
 <br>
-[![Downloads](https://img.shields.io/pepy/dt/wavedl?style=plastic&logo=pypi&logoColor=white&color=9ACD32)](https://pepy.tech/project/wavedl)
+[![Downloads](https://img.shields.io/badge/dynamic/json?url=https://pypistats.org/api/packages/wavedl/recent?period=month%26mirrors=false&query=data.last_month&style=plastic&logo=pypi&logoColor=white&color=9ACD32&label=Downloads&suffix=/month)](https://pypistats.org/packages/wavedl)
 [![License: MIT](https://img.shields.io/badge/License-MIT-orange.svg?style=plastic)](LICENSE)
 [![DOI](https://img.shields.io/badge/DOI-10.5281/zenodo.18012338-008080.svg?style=plastic)](https://doi.org/10.5281/zenodo.18012338)
 
@@ -462,7 +462,43 @@ WaveDL/
 | **U-Net** — U-shaped Network |||
 | `unet_regression` | 31.1M | 1D/2D/3D |
 
-> ⭐ = Pretrained on ImageNet. Recommended for smaller datasets.
+⭐ = **Pretrained on ImageNet** (recommended for smaller datasets). Weights are downloaded automatically on first use.
+- **Cache location**: `~/.cache/torch/hub/checkpoints/` (or `./.torch_cache/` on HPC if home is not writable)
+- **Size**: ~20–350 MB per model depending on architecture
+
+**💡 HPC Users**: If compute nodes block internet, pre-download weights on the login node:
+
+```bash
+# Run once on login node (with internet) — downloads ALL pretrained weights (~1.5 GB total)
+python -c "
+import os
+os.environ['TORCH_HOME'] = '.torch_cache'  # Match WaveDL's HPC cache location
+
+from torchvision import models as m
+from torchvision.models import video as v
+
+# Model name -> Weights class mapping
+weights = {
+    'resnet18': m.ResNet18_Weights, 'resnet50': m.ResNet50_Weights,
+    'efficientnet_b0': m.EfficientNet_B0_Weights, 'efficientnet_b1': m.EfficientNet_B1_Weights,
+    'efficientnet_b2': m.EfficientNet_B2_Weights, 'efficientnet_v2_s': m.EfficientNet_V2_S_Weights,
+    'efficientnet_v2_m': m.EfficientNet_V2_M_Weights, 'efficientnet_v2_l': m.EfficientNet_V2_L_Weights,
+    'mobilenet_v3_small': m.MobileNet_V3_Small_Weights, 'mobilenet_v3_large': m.MobileNet_V3_Large_Weights,
+    'regnet_y_400mf': m.RegNet_Y_400MF_Weights, 'regnet_y_800mf': m.RegNet_Y_800MF_Weights,
+    'regnet_y_1_6gf': m.RegNet_Y_1_6GF_Weights, 'regnet_y_3_2gf': m.RegNet_Y_3_2GF_Weights,
+    'regnet_y_8gf': m.RegNet_Y_8GF_Weights, 'swin_t': m.Swin_T_Weights, 'swin_s': m.Swin_S_Weights,
+    'swin_b': m.Swin_B_Weights, 'convnext_tiny': m.ConvNeXt_Tiny_Weights, 'densenet121': m.DenseNet121_Weights,
+}
+for name, w in weights.items():
+    getattr(m, name)(weights=w.DEFAULT); print(f'✓ {name}')
+
+# 3D video models
+v.r3d_18(weights=v.R3D_18_Weights.DEFAULT); print('✓ r3d_18')
+v.mc3_18(weights=v.MC3_18_Weights.DEFAULT); print('✓ mc3_18')
+print('\\n✓ All pretrained weights cached!')
+"
+```
+
 
 </details>
 
@@ -687,7 +723,6 @@ compile: false
 seed: 2025
 ```
 
-> [!TIP]
 > See [`configs/config.yaml`](configs/config.yaml) for the complete template with all available options documented.
 
 </details>
@@ -699,18 +734,20 @@ Automatically find the best training configuration using [Optuna](https://optuna
 
 **Run HPO:**
 
-You specify which models to search and how many trials to run:
 ```bash
-# Search 3 models with 100 trials
-python -m wavedl.hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 100
+# Basic HPO (auto-detects GPUs for parallel trials)
+wavedl-hpo --data_path train.npz --models cnn --n_trials 100
 
-# Search 1 model (faster)
-python -m wavedl.hpo --data_path train.npz --models cnn --n_trials 50
+# Search multiple models
+wavedl-hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 200
 
-# Search all your candidate models
-python -m wavedl.hpo --data_path train.npz --models cnn resnet18 resnet50 vit_small densenet121 --n_trials 200
+# Quick mode (fewer parameters, faster)
+wavedl-hpo --data_path train.npz --models cnn --n_trials 50 --quick
 ```
 
+> [!TIP]
+> **Auto GPU Detection**: HPO automatically detects available GPUs and runs one trial per GPU in parallel. On a 4-GPU system, 4 trials run simultaneously. Use `--n_jobs 1` to force serial execution.
+
 **Train with best parameters**
 
 After HPO completes, it prints the optimal command:
@@ -749,11 +786,11 @@ accelerate launch -m wavedl.train --data_path train.npz --model cnn --lr 3.2e-4
 | `--optimizers` | all 6 | Optimizers to search |
 | `--schedulers` | all 8 | Schedulers to search |
 | `--losses` | all 6 | Losses to search |
-| `--n_jobs` | `1` | Parallel trials (multi-GPU) |
+| `--n_jobs` | `-1` | Parallel trials (-1 = auto-detect GPUs) |
 | `--max_epochs` | `50` | Max epochs per trial |
 | `--output` | `hpo_results.json` | Output file |
 
-> [!TIP]
+
 > See [Available Models](#available-models) for all 38 architectures you can search.
 
 </details>

{wavedl-1.4.4 → wavedl-1.4.6}/README.md

@@ -4,7 +4,7 @@
 
 ### A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 
-[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg?style=plastic&logo=python&logoColor=white)](https://www.python.org/downloads/)
+[![Python 3.11+](https://img.shields.io/badge/Python-3.11+-blue.svg?style=plastic&logo=python&logoColor=white)](https://www.python.org/downloads/)
 [![PyTorch 2.x](https://img.shields.io/badge/PyTorch-2.x-ee4c2c.svg?style=plastic&logo=pytorch&logoColor=white)](https://pytorch.org/)
 [![Accelerate](https://img.shields.io/badge/Accelerate-Enabled-yellow.svg?style=plastic&logo=huggingface&logoColor=white)](https://huggingface.co/docs/accelerate/)
 <br>
@@ -12,7 +12,7 @@
 [![Lint](https://img.shields.io/github/actions/workflow/status/ductho-le/WaveDL/lint.yml?branch=main&style=plastic&logo=ruff&logoColor=white&label=Lint)](https://github.com/ductho-le/WaveDL/actions/workflows/lint.yml)
 [![Try it on Colab](https://img.shields.io/badge/Try_it_on_Colab-8E44AD?style=plastic&logo=googlecolab&logoColor=white)](https://colab.research.google.com/github/ductho-le/WaveDL/blob/main/notebooks/demo.ipynb)
 <br>
-[![Downloads](https://img.shields.io/pepy/dt/wavedl?style=plastic&logo=pypi&logoColor=white&color=9ACD32)](https://pepy.tech/project/wavedl)
+[![Downloads](https://img.shields.io/badge/dynamic/json?url=https://pypistats.org/api/packages/wavedl/recent?period=month%26mirrors=false&query=data.last_month&style=plastic&logo=pypi&logoColor=white&color=9ACD32&label=Downloads&suffix=/month)](https://pypistats.org/packages/wavedl)
 [![License: MIT](https://img.shields.io/badge/License-MIT-orange.svg?style=plastic)](LICENSE)
 [![DOI](https://img.shields.io/badge/DOI-10.5281/zenodo.18012338-008080.svg?style=plastic)](https://doi.org/10.5281/zenodo.18012338)
 
@@ -417,7 +417,43 @@ WaveDL/
 | **U-Net** — U-shaped Network |||
 | `unet_regression` | 31.1M | 1D/2D/3D |
 
-> ⭐ = Pretrained on ImageNet. Recommended for smaller datasets.
+⭐ = **Pretrained on ImageNet** (recommended for smaller datasets). Weights are downloaded automatically on first use.
+- **Cache location**: `~/.cache/torch/hub/checkpoints/` (or `./.torch_cache/` on HPC if home is not writable)
+- **Size**: ~20–350 MB per model depending on architecture
+
+**💡 HPC Users**: If compute nodes block internet, pre-download weights on the login node:
+
+```bash
+# Run once on login node (with internet) — downloads ALL pretrained weights (~1.5 GB total)
+python -c "
+import os
+os.environ['TORCH_HOME'] = '.torch_cache'  # Match WaveDL's HPC cache location
+
+from torchvision import models as m
+from torchvision.models import video as v
+
+# Model name -> Weights class mapping
+weights = {
+    'resnet18': m.ResNet18_Weights, 'resnet50': m.ResNet50_Weights,
+    'efficientnet_b0': m.EfficientNet_B0_Weights, 'efficientnet_b1': m.EfficientNet_B1_Weights,
+    'efficientnet_b2': m.EfficientNet_B2_Weights, 'efficientnet_v2_s': m.EfficientNet_V2_S_Weights,
+    'efficientnet_v2_m': m.EfficientNet_V2_M_Weights, 'efficientnet_v2_l': m.EfficientNet_V2_L_Weights,
+    'mobilenet_v3_small': m.MobileNet_V3_Small_Weights, 'mobilenet_v3_large': m.MobileNet_V3_Large_Weights,
+    'regnet_y_400mf': m.RegNet_Y_400MF_Weights, 'regnet_y_800mf': m.RegNet_Y_800MF_Weights,
+    'regnet_y_1_6gf': m.RegNet_Y_1_6GF_Weights, 'regnet_y_3_2gf': m.RegNet_Y_3_2GF_Weights,
+    'regnet_y_8gf': m.RegNet_Y_8GF_Weights, 'swin_t': m.Swin_T_Weights, 'swin_s': m.Swin_S_Weights,
+    'swin_b': m.Swin_B_Weights, 'convnext_tiny': m.ConvNeXt_Tiny_Weights, 'densenet121': m.DenseNet121_Weights,
+}
+for name, w in weights.items():
+    getattr(m, name)(weights=w.DEFAULT); print(f'✓ {name}')
+
+# 3D video models
+v.r3d_18(weights=v.R3D_18_Weights.DEFAULT); print('✓ r3d_18')
+v.mc3_18(weights=v.MC3_18_Weights.DEFAULT); print('✓ mc3_18')
+print('\\n✓ All pretrained weights cached!')
+"
+```
+
 
 </details>
 
@@ -642,7 +678,6 @@ compile: false
 seed: 2025
 ```
 
-> [!TIP]
 > See [`configs/config.yaml`](configs/config.yaml) for the complete template with all available options documented.
 
 </details>
@@ -654,18 +689,20 @@ Automatically find the best training configuration using [Optuna](https://optuna
 
 **Run HPO:**
 
-You specify which models to search and how many trials to run:
 ```bash
-# Search 3 models with 100 trials
-python -m wavedl.hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 100
+# Basic HPO (auto-detects GPUs for parallel trials)
+wavedl-hpo --data_path train.npz --models cnn --n_trials 100
 
-# Search 1 model (faster)
-python -m wavedl.hpo --data_path train.npz --models cnn --n_trials 50
+# Search multiple models
+wavedl-hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 200
 
-# Search all your candidate models
-python -m wavedl.hpo --data_path train.npz --models cnn resnet18 resnet50 vit_small densenet121 --n_trials 200
+# Quick mode (fewer parameters, faster)
+wavedl-hpo --data_path train.npz --models cnn --n_trials 50 --quick
 ```
 
+> [!TIP]
+> **Auto GPU Detection**: HPO automatically detects available GPUs and runs one trial per GPU in parallel. On a 4-GPU system, 4 trials run simultaneously. Use `--n_jobs 1` to force serial execution.
+
 **Train with best parameters**
 
 After HPO completes, it prints the optimal command:
@@ -704,11 +741,11 @@ accelerate launch -m wavedl.train --data_path train.npz --model cnn --lr 3.2e-4
 | `--optimizers` | all 6 | Optimizers to search |
 | `--schedulers` | all 8 | Schedulers to search |
 | `--losses` | all 6 | Losses to search |
-| `--n_jobs` | `1` | Parallel trials (multi-GPU) |
+| `--n_jobs` | `-1` | Parallel trials (-1 = auto-detect GPUs) |
 | `--max_epochs` | `50` | Max epochs per trial |
 | `--output` | `hpo_results.json` | Output file |
 
-> [!TIP]
+
 > See [Available Models](#available-models) for all 38 architectures you can search.
 
 </details>

{wavedl-1.4.4 → wavedl-1.4.6}/src/wavedl/__init__.py

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

{wavedl-1.4.4 → wavedl-1.4.6}/src/wavedl/hpc.py

@@ -174,7 +174,9 @@ Environment Variables:
     return args, remaining
 
 
-def print_summary(exit_code: int, wandb_mode: str, wandb_dir: str) -> None:
+def print_summary(
+    exit_code: int, wandb_enabled: bool, wandb_mode: str, wandb_dir: str
+) -> None:
     """Print post-training summary and instructions."""
     print()
     print("=" * 40)
@@ -183,7 +185,8 @@ def print_summary(exit_code: int, wandb_mode: str, wandb_dir: str) -> None:
         print("✅ Training completed successfully!")
         print("=" * 40)
 
-        if wandb_mode == "offline":
+        # Only show WandB sync instructions if user enabled wandb
+        if wandb_enabled and wandb_mode == "offline":
             print()
             print("📊 WandB Sync Instructions:")
             print("   From the login node, run:")
@@ -237,6 +240,10 @@ def main() -> int:
         f"--dynamo_backend={args.dynamo_backend}",
     ]
 
+    # Explicitly set multi_gpu to suppress accelerate auto-detection warning
+    if num_gpus > 1:
+        cmd.append("--multi_gpu")
+
     # 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}")
@@ -263,8 +270,10 @@ def main() -> int:
         exit_code = 130
 
     # Print summary
+    wandb_enabled = "--wandb" in train_args
     print_summary(
         exit_code,
+        wandb_enabled,
         os.environ.get("WANDB_MODE", "offline"),
         os.environ.get("WANDB_DIR", "/tmp/wandb"),
     )

{wavedl-1.4.4 → wavedl-1.4.6}/src/wavedl/hpo.py

@@ -31,7 +31,7 @@ try:
     import optuna
     from optuna.trial import TrialState
 except ImportError:
-    print("Error: Optuna not installed. Run: pip install -e '.[hpo]'")
+    print("Error: Optuna not installed. Run: pip install wavedl")
     sys.exit(1)
 
 
@@ -147,6 +147,32 @@ def create_objective(args):
             cmd.extend(["--output_dir", tmpdir])
             history_file = Path(tmpdir) / "training_history.csv"
 
+            # GPU isolation for parallel trials: assign each trial to a specific GPU
+            # This prevents multiple trials from competing for all GPUs
+            env = None
+            if args.n_jobs > 1:
+                import os
+
+                # Detect available GPUs
+                n_gpus = 1
+                try:
+                    import subprocess as sp
+
+                    result_gpu = sp.run(
+                        ["nvidia-smi", "--list-gpus"],
+                        capture_output=True,
+                        text=True,
+                    )
+                    if result_gpu.returncode == 0:
+                        n_gpus = len(result_gpu.stdout.strip().split("\n"))
+                except Exception:
+                    pass
+
+                # Assign trial to a specific GPU (round-robin)
+                gpu_id = trial.number % n_gpus
+                env = os.environ.copy()
+                env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+
             # Run training
             try:
                 result = subprocess.run(
@@ -155,6 +181,7 @@ def create_objective(args):
                     text=True,
                     timeout=args.timeout,
                     cwd=Path(__file__).parent,
+                    env=env,
                 )
 
                 # Read best val_loss from training_history.csv (reliable machine-readable)
@@ -248,7 +275,10 @@ Examples:
         "--n_trials", type=int, default=50, help="Number of HPO trials (default: 50)"
     )
     parser.add_argument(
-        "--n_jobs", type=int, default=1, help="Parallel trials (default: 1)"
+        "--n_jobs",
+        type=int,
+        default=-1,
+        help="Parallel trials (-1 = auto-detect GPUs, default: -1)",
     )
     parser.add_argument(
         "--quick",
@@ -315,11 +345,30 @@ Examples:
 
     args = parser.parse_args()
 
+    # Convert to absolute path (child processes may run in different cwd)
+    args.data_path = str(Path(args.data_path).resolve())
+
     # Validate data path
     if not Path(args.data_path).exists():
         print(f"Error: Data file not found: {args.data_path}")
         sys.exit(1)
 
+    # Auto-detect GPUs for n_jobs if not specified
+    if args.n_jobs == -1:
+        try:
+            result_gpu = subprocess.run(
+                ["nvidia-smi", "--list-gpus"],
+                capture_output=True,
+                text=True,
+            )
+            if result_gpu.returncode == 0:
+                args.n_jobs = max(1, len(result_gpu.stdout.strip().split("\n")))
+            else:
+                args.n_jobs = 1
+        except Exception:
+            args.n_jobs = 1
+        print(f"Auto-detected {args.n_jobs} GPU(s) for parallel trials")
+
     # Create study
     print("=" * 60)
     print("WaveDL Hyperparameter Optimization")

{wavedl-1.4.4 → wavedl-1.4.6}/src/wavedl/test.py

@@ -366,23 +366,37 @@ def load_checkpoint(
     logging.info(f"   Building model: {model_name}")
     model = build_model(model_name, in_shape=in_shape, out_size=out_size)
 
-    # Load weights (prefer safetensors)
-    weight_path = checkpoint_dir / "model.safetensors"
-    if not weight_path.exists():
-        weight_path = checkpoint_dir / "pytorch_model.bin"
-
-    if not weight_path.exists():
-        raise FileNotFoundError(f"No model weights found in {checkpoint_dir}")
+    # Load weights (check multiple formats in order of preference)
+    weight_path = None
+    for fname in ["model.safetensors", "model.bin", "pytorch_model.bin"]:
+        candidate = checkpoint_dir / fname
+        if candidate.exists():
+            weight_path = candidate
+            break
+
+    if weight_path is None:
+        raise FileNotFoundError(
+            f"No model weights found in {checkpoint_dir}. "
+            f"Expected one of: model.safetensors, model.bin, pytorch_model.bin"
+        )
 
     if HAS_SAFETENSORS and weight_path.suffix == ".safetensors":
         state_dict = load_safetensors(str(weight_path))
     else:
         state_dict = torch.load(weight_path, map_location="cpu", weights_only=True)
 
-    # Remove 'module.' prefix from DDP checkpoints (leading only, not all occurrences)
-    state_dict = {
-        (k[7:] if k.startswith("module.") else k): v for k, v in state_dict.items()
-    }
+    # Remove wrapper prefixes from checkpoints:
+    # - 'module.' from DDP (DistributedDataParallel)
+    # - '_orig_mod.' from torch.compile()
+    cleaned_dict = {}
+    for k, v in state_dict.items():
+        key = k
+        if key.startswith("module."):
+            key = key[7:]  # Remove 'module.' (7 chars)
+        if key.startswith("_orig_mod."):
+            key = key[10:]  # Remove '_orig_mod.' (10 chars)
+        cleaned_dict[key] = v
+    state_dict = cleaned_dict
 
     model.load_state_dict(state_dict)
     model.eval()

{wavedl-1.4.4 → wavedl-1.4.6}/src/wavedl/train.py

@@ -148,6 +148,24 @@ torch.set_float32_matmul_precision("high")  # Use TF32 for float32 ops
 torch.backends.cudnn.benchmark = True
 
 
+# ==============================================================================
+# LOGGING UTILITIES
+# ==============================================================================
+from contextlib import contextmanager
+
+
+@contextmanager
+def suppress_accelerate_logging():
+    """Temporarily suppress accelerate's verbose checkpoint save messages."""
+    accelerate_logger = logging.getLogger("accelerate.checkpointing")
+    original_level = accelerate_logger.level
+    accelerate_logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        accelerate_logger.setLevel(original_level)
+
+
 # ==============================================================================
 # ARGUMENT PARSING
 # ==============================================================================
@@ -1033,7 +1051,8 @@ def main():
             # Step 3: Save checkpoint with all ranks participating
             if is_best_epoch:
                 ckpt_dir = os.path.join(args.output_dir, "best_checkpoint")
-                accelerator.save_state(ckpt_dir)  # All ranks must call this
+                with suppress_accelerate_logging():
+                    accelerator.save_state(ckpt_dir, safe_serialization=False)
 
                 # Step 4: Rank 0 handles metadata and updates tracking variables
                 if accelerator.is_main_process:
@@ -1096,7 +1115,8 @@ def main():
             if periodic_checkpoint_needed:
                 ckpt_name = f"epoch_{epoch + 1}_checkpoint"
                 ckpt_dir = os.path.join(args.output_dir, ckpt_name)
-                accelerator.save_state(ckpt_dir)  # All ranks participate
+                with suppress_accelerate_logging():
+                    accelerator.save_state(ckpt_dir, safe_serialization=False)
 
                 if accelerator.is_main_process:
                     with open(os.path.join(ckpt_dir, "training_meta.pkl"), "wb") as f:
@@ -1147,7 +1167,11 @@ def main():
 
     except KeyboardInterrupt:
         logger.warning("Training interrupted. Saving emergency checkpoint...")
-        accelerator.save_state(os.path.join(args.output_dir, "interrupted_checkpoint"))
+        with suppress_accelerate_logging():
+            accelerator.save_state(
+                os.path.join(args.output_dir, "interrupted_checkpoint"),
+                safe_serialization=False,
+            )
 
     except Exception as e:
         logger.error(f"Critical error: {e}", exc_info=True)

{wavedl-1.4.4 → wavedl-1.4.6/src/wavedl.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.4.4
+Version: 1.4.6
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -49,7 +49,7 @@ Requires-Dist: triton>=2.0.0; sys_platform == "linux"
 
 ### A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 
-[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg?style=plastic&logo=python&logoColor=white)](https://www.python.org/downloads/)
+[![Python 3.11+](https://img.shields.io/badge/Python-3.11+-blue.svg?style=plastic&logo=python&logoColor=white)](https://www.python.org/downloads/)
 [![PyTorch 2.x](https://img.shields.io/badge/PyTorch-2.x-ee4c2c.svg?style=plastic&logo=pytorch&logoColor=white)](https://pytorch.org/)
 [![Accelerate](https://img.shields.io/badge/Accelerate-Enabled-yellow.svg?style=plastic&logo=huggingface&logoColor=white)](https://huggingface.co/docs/accelerate/)
 <br>
@@ -57,7 +57,7 @@ Requires-Dist: triton>=2.0.0; sys_platform == "linux"
 [![Lint](https://img.shields.io/github/actions/workflow/status/ductho-le/WaveDL/lint.yml?branch=main&style=plastic&logo=ruff&logoColor=white&label=Lint)](https://github.com/ductho-le/WaveDL/actions/workflows/lint.yml)
 [![Try it on Colab](https://img.shields.io/badge/Try_it_on_Colab-8E44AD?style=plastic&logo=googlecolab&logoColor=white)](https://colab.research.google.com/github/ductho-le/WaveDL/blob/main/notebooks/demo.ipynb)
 <br>
-[![Downloads](https://img.shields.io/pepy/dt/wavedl?style=plastic&logo=pypi&logoColor=white&color=9ACD32)](https://pepy.tech/project/wavedl)
+[![Downloads](https://img.shields.io/badge/dynamic/json?url=https://pypistats.org/api/packages/wavedl/recent?period=month%26mirrors=false&query=data.last_month&style=plastic&logo=pypi&logoColor=white&color=9ACD32&label=Downloads&suffix=/month)](https://pypistats.org/packages/wavedl)
 [![License: MIT](https://img.shields.io/badge/License-MIT-orange.svg?style=plastic)](LICENSE)
 [![DOI](https://img.shields.io/badge/DOI-10.5281/zenodo.18012338-008080.svg?style=plastic)](https://doi.org/10.5281/zenodo.18012338)
 
@@ -462,7 +462,43 @@ WaveDL/
 | **U-Net** — U-shaped Network |||
 | `unet_regression` | 31.1M | 1D/2D/3D |
 
-> ⭐ = Pretrained on ImageNet. Recommended for smaller datasets.
+⭐ = **Pretrained on ImageNet** (recommended for smaller datasets). Weights are downloaded automatically on first use.
+- **Cache location**: `~/.cache/torch/hub/checkpoints/` (or `./.torch_cache/` on HPC if home is not writable)
+- **Size**: ~20–350 MB per model depending on architecture
+
+**💡 HPC Users**: If compute nodes block internet, pre-download weights on the login node:
+
+```bash
+# Run once on login node (with internet) — downloads ALL pretrained weights (~1.5 GB total)
+python -c "
+import os
+os.environ['TORCH_HOME'] = '.torch_cache'  # Match WaveDL's HPC cache location
+
+from torchvision import models as m
+from torchvision.models import video as v
+
+# Model name -> Weights class mapping
+weights = {
+    'resnet18': m.ResNet18_Weights, 'resnet50': m.ResNet50_Weights,
+    'efficientnet_b0': m.EfficientNet_B0_Weights, 'efficientnet_b1': m.EfficientNet_B1_Weights,
+    'efficientnet_b2': m.EfficientNet_B2_Weights, 'efficientnet_v2_s': m.EfficientNet_V2_S_Weights,
+    'efficientnet_v2_m': m.EfficientNet_V2_M_Weights, 'efficientnet_v2_l': m.EfficientNet_V2_L_Weights,
+    'mobilenet_v3_small': m.MobileNet_V3_Small_Weights, 'mobilenet_v3_large': m.MobileNet_V3_Large_Weights,
+    'regnet_y_400mf': m.RegNet_Y_400MF_Weights, 'regnet_y_800mf': m.RegNet_Y_800MF_Weights,
+    'regnet_y_1_6gf': m.RegNet_Y_1_6GF_Weights, 'regnet_y_3_2gf': m.RegNet_Y_3_2GF_Weights,
+    'regnet_y_8gf': m.RegNet_Y_8GF_Weights, 'swin_t': m.Swin_T_Weights, 'swin_s': m.Swin_S_Weights,
+    'swin_b': m.Swin_B_Weights, 'convnext_tiny': m.ConvNeXt_Tiny_Weights, 'densenet121': m.DenseNet121_Weights,
+}
+for name, w in weights.items():
+    getattr(m, name)(weights=w.DEFAULT); print(f'✓ {name}')
+
+# 3D video models
+v.r3d_18(weights=v.R3D_18_Weights.DEFAULT); print('✓ r3d_18')
+v.mc3_18(weights=v.MC3_18_Weights.DEFAULT); print('✓ mc3_18')
+print('\\n✓ All pretrained weights cached!')
+"
+```
+
 
 </details>
 
@@ -687,7 +723,6 @@ compile: false
 seed: 2025
 ```
 
-> [!TIP]
 > See [`configs/config.yaml`](configs/config.yaml) for the complete template with all available options documented.
 
 </details>
@@ -699,18 +734,20 @@ Automatically find the best training configuration using [Optuna](https://optuna
 
 **Run HPO:**
 
-You specify which models to search and how many trials to run:
 ```bash
-# Search 3 models with 100 trials
-python -m wavedl.hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 100
+# Basic HPO (auto-detects GPUs for parallel trials)
+wavedl-hpo --data_path train.npz --models cnn --n_trials 100
 
-# Search 1 model (faster)
-python -m wavedl.hpo --data_path train.npz --models cnn --n_trials 50
+# Search multiple models
+wavedl-hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 200
 
-# Search all your candidate models
-python -m wavedl.hpo --data_path train.npz --models cnn resnet18 resnet50 vit_small densenet121 --n_trials 200
+# Quick mode (fewer parameters, faster)
+wavedl-hpo --data_path train.npz --models cnn --n_trials 50 --quick
 ```
 
+> [!TIP]
+> **Auto GPU Detection**: HPO automatically detects available GPUs and runs one trial per GPU in parallel. On a 4-GPU system, 4 trials run simultaneously. Use `--n_jobs 1` to force serial execution.
+
 **Train with best parameters**
 
 After HPO completes, it prints the optimal command:
@@ -749,11 +786,11 @@ accelerate launch -m wavedl.train --data_path train.npz --model cnn --lr 3.2e-4
 | `--optimizers` | all 6 | Optimizers to search |
 | `--schedulers` | all 8 | Schedulers to search |
 | `--losses` | all 6 | Losses to search |
-| `--n_jobs` | `1` | Parallel trials (multi-GPU) |
+| `--n_jobs` | `-1` | Parallel trials (-1 = auto-detect GPUs) |
 | `--max_epochs` | `50` | Max epochs per trial |
 | `--output` | `hpo_results.json` | Output file |
 
-> [!TIP]
+
 > See [Available Models](#available-models) for all 38 architectures you can search.
 
 </details>