wavedl 1.4.3__py3-none-any.whl → 1.4.5__py3-none-any.whl

wavedl/__init__.py

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

wavedl/hpc.py

@@ -26,7 +26,6 @@ import os
 import shutil
 import subprocess
 import sys
-import tempfile
 from pathlib import Path
 
 
@@ -60,26 +59,40 @@ def setup_hpc_environment() -> None:
     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", os.environ.get("TMPDIR", tempfile.gettempdir())
-    )
-
-    # Configure directories for systems with restricted home directories
-    os.environ.setdefault("MPLCONFIGDIR", f"{tmpdir}/matplotlib")
-    os.environ.setdefault(
-        "FONTCONFIG_PATH", os.environ.get("FONTCONFIG_PATH", "/etc/fonts")
-    )
-    os.environ.setdefault("XDG_CACHE_HOME", f"{tmpdir}/.cache")
-
-    # Ensure matplotlib config dir exists
-    Path(os.environ["MPLCONFIGDIR"]).mkdir(parents=True, exist_ok=True)
+    # Check if home is writable
+    home = os.path.expanduser("~")
+    home_writable = os.access(home, os.W_OK)
+
+    # Use SLURM_TMPDIR if available, otherwise CWD for HPC, or system temp
+    if home_writable:
+        # Local machine - let libraries use defaults
+        cache_base = None
+    else:
+        # HPC with restricted home - use CWD for persistent caches
+        cache_base = os.getcwd()
+
+    # Only set environment variables if home is not writable
+    if cache_base:
+        os.environ.setdefault("TORCH_HOME", f"{cache_base}/.torch_cache")
+        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 [
+            "TORCH_HOME",
+            "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"{tmpdir}/wandb")
-    os.environ.setdefault("WANDB_CACHE_DIR", f"{tmpdir}/wandb_cache")
-    os.environ.setdefault("WANDB_CONFIG_DIR", f"{tmpdir}/wandb_config")
+    if cache_base:
+        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
     os.environ.setdefault(

wavedl/test.py

@@ -29,29 +29,52 @@ Author: Ductho Le (ductho.le@outlook.com)
 # ==============================================================================
 # ENVIRONMENT CONFIGURATION (must be before matplotlib import)
 # ==============================================================================
+# Auto-configure writable cache directories when home is not writable.
+# Uses current working directory as fallback - works on HPC and local machines.
 import os
 
 
-os.environ.setdefault("MPLCONFIGDIR", os.getenv("TMPDIR", "/tmp") + "/matplotlib")
-os.environ.setdefault("FONTCONFIG_PATH", "/etc/fonts")
-
-import argparse
-import logging
-import pickle
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-import torch
-import torch.nn as nn
-from sklearn.metrics import mean_absolute_error, r2_score
-from torch.utils.data import DataLoader, TensorDataset
-from tqdm.auto import tqdm
+def _setup_cache_dir(env_var: str, subdir: str) -> None:
+    """Set cache directory to CWD if home is not writable."""
+    if env_var in os.environ:
+        return  # User already set, respect their choice
+
+    # Check if home is writable
+    home = os.path.expanduser("~")
+    if os.access(home, os.W_OK):
+        return  # Home is writable, let library use defaults
+
+    # Home not writable - use current working directory
+    cache_path = os.path.join(os.getcwd(), f".{subdir}")
+    os.makedirs(cache_path, exist_ok=True)
+    os.environ[env_var] = cache_path
+
+
+# Configure cache directories (before any library imports)
+_setup_cache_dir("TORCH_HOME", "torch_cache")
+_setup_cache_dir("MPLCONFIGDIR", "matplotlib")
+_setup_cache_dir("FONTCONFIG_CACHE", "fontconfig")
+_setup_cache_dir("XDG_DATA_HOME", "local/share")
+_setup_cache_dir("XDG_STATE_HOME", "local/state")
+_setup_cache_dir("XDG_CACHE_HOME", "cache")
+
+import argparse  # noqa: E402
+import logging  # noqa: E402
+import pickle  # noqa: E402
+from pathlib import Path  # noqa: E402
+
+import matplotlib.pyplot as plt  # noqa: E402
+import numpy as np  # noqa: E402
+import pandas as pd  # noqa: E402
+import torch  # noqa: E402
+import torch.nn as nn  # noqa: E402
+from sklearn.metrics import mean_absolute_error, r2_score  # noqa: E402
+from torch.utils.data import DataLoader, TensorDataset  # noqa: E402
+from tqdm.auto import tqdm  # noqa: E402
 
 # Local imports
-from wavedl.models import build_model, list_models
-from wavedl.utils import (
+from wavedl.models import build_model, list_models  # noqa: E402
+from wavedl.utils import (  # noqa: E402
     FIGURE_DPI,
     calc_pearson,
     load_test_data,
@@ -356,10 +379,18 @@ def load_checkpoint(
     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/train.py

@@ -40,45 +40,34 @@ from __future__ import annotations
 # =============================================================================
 # HPC Environment Setup (MUST be before any library imports)
 # =============================================================================
-# Set writable cache directories for matplotlib and fontconfig ONLY when
-# the default paths are not writable (common on HPC clusters).
+# Auto-configure writable cache directories when home is not writable.
+# Uses current working directory as fallback - works on HPC and local machines.
 import os
-import tempfile
 
 
-def _setup_cache_dir(env_var: str, default_subpath: str) -> None:
-    """Set cache directory only if default path is not writable."""
+def _setup_cache_dir(env_var: str, subdir: str) -> None:
+    """Set cache directory to CWD if home is not writable."""
     if env_var in os.environ:
         return  # User already set, respect their choice
 
-    # Check if default home config path is writable
+    # Check if home is writable
     home = os.path.expanduser("~")
-    default_path = os.path.join(home, ".config", default_subpath)
-    default_parent = os.path.dirname(default_path)
+    if os.access(home, os.W_OK):
+        return  # Home is writable, let library use defaults
 
-    # If default path or its parent is writable, let the library use defaults
-    if (
-        os.access(default_path, os.W_OK)
-        or (os.path.exists(default_parent) and os.access(default_parent, os.W_OK))
-        or os.access(os.path.join(home, ".config"), os.W_OK)
-    ):
-        return
-
-    # Default not writable - find alternative location
-    for cache_base in [
-        os.environ.get("SCRATCH"),
-        os.environ.get("SLURM_TMPDIR"),
-        tempfile.gettempdir(),
-    ]:
-        if cache_base and os.access(cache_base, os.W_OK):
-            cache_path = os.path.join(cache_base, f".{default_subpath}")
-            os.makedirs(cache_path, exist_ok=True)
-            os.environ[env_var] = cache_path
-            return
+    # Home not writable - use current working directory
+    cache_path = os.path.join(os.getcwd(), f".{subdir}")
+    os.makedirs(cache_path, exist_ok=True)
+    os.environ[env_var] = cache_path
 
 
+# Configure cache directories (before any library imports)
+_setup_cache_dir("TORCH_HOME", "torch_cache")
 _setup_cache_dir("MPLCONFIGDIR", "matplotlib")
 _setup_cache_dir("FONTCONFIG_CACHE", "fontconfig")
+_setup_cache_dir("XDG_DATA_HOME", "local/share")
+_setup_cache_dir("XDG_STATE_HOME", "local/state")
+_setup_cache_dir("XDG_CACHE_HOME", "cache")
 
 # =============================================================================
 # Standard imports (after environment setup)
@@ -1065,7 +1054,17 @@ def main():
                             f,
                         )
 
-                    unwrapped = accelerator.unwrap_model(model)
+                    # Unwrap model for saving (handle torch.compile compatibility)
+                    try:
+                        unwrapped = accelerator.unwrap_model(model)
+                    except KeyError:
+                        # torch.compile model may not have _orig_mod in expected location
+                        # Fall back to getting the module directly
+                        unwrapped = model.module if hasattr(model, "module") else model
+                        # If still compiled, try to get the underlying model
+                        if hasattr(unwrapped, "_orig_mod"):
+                            unwrapped = unwrapped._orig_mod
+
                     torch.save(
                         unwrapped.state_dict(),
                         os.path.join(args.output_dir, "best_model_weights.pth"),

{wavedl-1.4.3.dist-info → wavedl-1.4.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.4.3
+Version: 1.4.5
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -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>
@@ -753,7 +788,7 @@ accelerate launch -m wavedl.train --data_path train.npz --model cnn --lr 3.2e-4
 | `--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.3.dist-info → wavedl-1.4.5.dist-info}/RECORD

@@ -1,8 +1,8 @@
-wavedl/__init__.py,sha256=7uvgr9r21qLGu6RIZDUfQBhg1vXNAfmIWl_P6BMj_KQ,1177
-wavedl/hpc.py,sha256=de_GKERX8GS10sXRX9yXiGzMnk1jjq8JPzRw7QDs6d4,7967
+wavedl/__init__.py,sha256=2ro7SYQ3wCmq-ejiAm5sd6BeXf6sZgixC9U2vS7Ckbs,1177
+wavedl/hpc.py,sha256=0h8IZzOT0EzmEv3fU9cKyRVE9V1ivtBzbjuBCaxYadc,8445
 wavedl/hpo.py,sha256=YJXsnSGEBSVUqp_2ah7zu3_VClAUqZrdkuzDaSqQUjU,12952
-wavedl/test.py,sha256=jZmRJaivYYTMMTaccCi0yQjHOfp0a9YWR1wAPeKFH-k,36246
-wavedl/train.py,sha256=TjLABBPCqu9r7FEWlxJlKsT7uAMo6hiDxRfii2SKXe4,49052
+wavedl/test.py,sha256=81al6vQBDAJ3CpSEtxZn6xzR1c4-jo28R7tX_84KROc,37642
+wavedl/train.py,sha256=_pW7prvlNqfUGrGweHO2QelS87UiAYKvyJwqMAIj6yI,49292
 wavedl/models/__init__.py,sha256=lfSohEnAUztO14nuwayMJhPjpgySzRN3jGiyAUuBmAU,3206
 wavedl/models/_template.py,sha256=J_D8taSPmV8lBaucN_vU-WiG98iFr7CJrZVNNX_Tdts,4600
 wavedl/models/base.py,sha256=T9iDF9IQM2MYucG_ggQd31rieUkB2fob-nkHyNIl2ak,7337
@@ -29,9 +29,9 @@ wavedl/utils/losses.py,sha256=5762M-TBC_hz6uyj1NPbU1vZeFOJQq7fR3-j7OygJRo,7254
 wavedl/utils/metrics.py,sha256=mkCpqZwl_XUpNvA5Ekjf7y-HqApafR7eR6EuA8cBdM8,37287
 wavedl/utils/optimizers.py,sha256=PyIkJ_hRhFi_Fio81Gy5YQNhcME0JUUEl8OTSyu-0RA,6323
 wavedl/utils/schedulers.py,sha256=e6Sf0yj8VOqkdwkUHLMyUfGfHKTX4NMr-zfgxWqCTYI,7659
-wavedl-1.4.3.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
-wavedl-1.4.3.dist-info/METADATA,sha256=vs3nt8R2O5lD7q-si9M5ChyrriIKm0fFzDwi4HVIYxw,40386
-wavedl-1.4.3.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-wavedl-1.4.3.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
-wavedl-1.4.3.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
-wavedl-1.4.3.dist-info/RECORD,,
+wavedl-1.4.5.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
+wavedl-1.4.5.dist-info/METADATA,sha256=4ltxFDaqPqh4XUAW_K8nkFmvqBzPcL2cxmghH11GMWg,42191
+wavedl-1.4.5.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+wavedl-1.4.5.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
+wavedl-1.4.5.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
+wavedl-1.4.5.dist-info/RECORD,,