wavedl 1.4.4__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.4"
+__version__ = "1.4.5"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

wavedl/test.py

@@ -379,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-1.4.4.dist-info → wavedl-1.4.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.4.4
+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.4.dist-info → wavedl-1.4.5.dist-info}/RECORD

@@ -1,7 +1,7 @@
-wavedl/__init__.py,sha256=n0XNSrp0aGEE6HQpzbF1wiKK-jORKgcc0Q2Op4MtQGk,1177
+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=Wajcze8gFEyJ9VyN_Bq-YadZ_VZtVaX_HicvUmW6MXM,37365
+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
@@ -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.4.dist-info/LICENSE,sha256=cEUCvcvH-9BT9Y-CNGY__PwWONCKu9zsoIqWA-NeHJ4,1066
-wavedl-1.4.4.dist-info/METADATA,sha256=4tgSvkwzJmZP3PLCrqx-FYV_w6VT6Mi4XIsB0Dvb6_0,40386
-wavedl-1.4.4.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-wavedl-1.4.4.dist-info/entry_points.txt,sha256=f1RNDkXFZwBzrBzTMFocJ6xhfTvTmaEDTi5YyDEUaF8,140
-wavedl-1.4.4.dist-info/top_level.txt,sha256=ccneUt3D5Qzbh3bsBSSrq9bqrhGiogcWKY24ZC4Q6Xw,7
-wavedl-1.4.4.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,,