wavedl 1.9.0

For fine-grained control over distributed training, you can use accelerate launch directly:

# Custom accelerate configuration
accelerate launch -m wavedl.train --model <model_name> --data_path <train_data> --output_dir <output_folder>

# Multi-node training
accelerate launch --num_machines 2 --main_process_ip <ip> -m wavedl.train --model cnn --data_path train.npz

Requirements (for a vector-output model):

1. Inherit from BaseModel
2. Accept in_shape, out_size in __init__
3. Return a tensor of shape (batch, out_size) from forward()

Step 1: Create my_model.py

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_shape, out_size, **kwargs):
        # in_shape: spatial dimensions, e.g., (128,) or (64, 64) or (32, 32, 32)
        # out_size: number of parameters to predict (auto-detected from data)
        super().__init__(in_shape, out_size)

        # Define your layers (this is just an example for 2D)
        self.conv1 = nn.Conv2d(1, 64, 3, padding=1)  # Input always has 1 channel
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.fc = nn.Linear(128, out_size)

    def forward(self, x):
        # Input x has shape: (batch, 1, *in_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, out_size)

Step 2: Train

wavedl-train --import my_model.py --model my_model --data_path train.npz

WaveDL handles everything else: training loop, logging, checkpoints, multi-GPU, early stopping, etc.

For a custom dense 2D model, also accept the inferred out_shape=(C, H_out, W_out), set supports_spatial_output = True (and requires_spatial_output = True if it cannot emit vectors), and return exactly (batch, C, H_out, W_out).

Output support

The dimensionality column below describes the input. A model accepting 2D input does not automatically support dense output; look at the output-support table above when selecting a model.

⭐ = Pretrained on ImageNet (recommended for smaller datasets). Weights are downloaded automatically on first use.

🗺️ = Dense 2D field output. Unmarked models produce regression vectors.

• Cache location: ~/.cache/torch/hub/checkpoints/ (or ./.torch_cache/ on HPC if home is not writable)
• Use pretrained weights: Pass --pretrained (default: train from scratch). Only effective for *_pretrained model variants marked with ⭐ above.

💡 HPC Users: If compute nodes block internet, pre-download weights on the login node:

# Run once on login node (with internet) — downloads ALL pretrained weights
# Uses download-only approach (no model instantiation) to avoid CPU time limits
python -c "
import os, torch, warnings
warnings.filterwarnings('ignore', category=UserWarning, module='pydantic')
os.environ['TORCH_HOME'] = '.torch_cache'  # Match WaveDL's HPC cache location
os.environ['HF_HOME'] = '.hf_cache'        # Match WaveDL's HPC cache for timm models

from torchvision import models as m
from torchvision.models import video as v

# === TorchVision + Video Models — download only, no model instantiation ===
urls = [
    ('ResNet18',         m.ResNet18_Weights.IMAGENET1K_V1.url),
    ('ResNet50',         m.ResNet50_Weights.IMAGENET1K_V1.url),
    ('EfficientNet_B0',  m.EfficientNet_B0_Weights.IMAGENET1K_V1.url),
    ('EfficientNet_B2',  m.EfficientNet_B2_Weights.IMAGENET1K_V1.url),
    ('EfficientNet_B4',  m.EfficientNet_B4_Weights.IMAGENET1K_V1.url),
    ('EfficientNet_B7',  m.EfficientNet_B7_Weights.IMAGENET1K_V1.url),
    ('EfficientNetV2_S', m.EfficientNet_V2_S_Weights.IMAGENET1K_V1.url),
    ('EfficientNetV2_M', m.EfficientNet_V2_M_Weights.IMAGENET1K_V1.url),
    ('EfficientNetV2_L', m.EfficientNet_V2_L_Weights.IMAGENET1K_V1.url),
    ('MobileNetV3_S',    m.MobileNet_V3_Small_Weights.IMAGENET1K_V1.url),
    ('MobileNetV3_L',    m.MobileNet_V3_Large_Weights.IMAGENET1K_V1.url),
    ('RegNet_Y_400MF',   m.RegNet_Y_400MF_Weights.IMAGENET1K_V1.url),
    ('RegNet_Y_800MF',   m.RegNet_Y_800MF_Weights.IMAGENET1K_V1.url),
    ('RegNet_Y_1_6GF',   m.RegNet_Y_1_6GF_Weights.IMAGENET1K_V1.url),
    ('RegNet_Y_3_2GF',   m.RegNet_Y_3_2GF_Weights.IMAGENET1K_V1.url),
    ('RegNet_Y_8GF',     m.RegNet_Y_8GF_Weights.IMAGENET1K_V1.url),
    ('Swin_T',           m.Swin_T_Weights.IMAGENET1K_V1.url),
    ('Swin_S',           m.Swin_S_Weights.IMAGENET1K_V1.url),
    ('Swin_B',           m.Swin_B_Weights.IMAGENET1K_V1.url),
    ('ConvNeXt_Tiny',    m.ConvNeXt_Tiny_Weights.IMAGENET1K_V1.url),
    ('DenseNet121',      m.DenseNet121_Weights.IMAGENET1K_V1.url),
    ('R3D_18',           v.R3D_18_Weights.KINETICS400_V1.url),
    ('MC3_18',           v.MC3_18_Weights.KINETICS400_V1.url),
]
cache = os.path.join(os.environ['TORCH_HOME'], 'hub', 'checkpoints')
os.makedirs(cache, exist_ok=True)
for name, url in urls:
    torch.hub.download_url_to_file(url, os.path.join(cache, os.path.basename(url)))
    print(f'  ✓ {name}')

# === Timm Models — download only via HuggingFace Hub ===
import timm
from huggingface_hub import hf_hub_download
timm_models = [
    'maxvit_tiny_tf_224', 'maxvit_small_tf_224', 'maxvit_base_tf_224',
    'fastvit_t8', 'fastvit_t12', 'fastvit_s12', 'fastvit_sa12',
    'caformer_s18', 'caformer_s36', 'caformer_m36', 'poolformer_s12',
    'convnextv2_tiny',
    'efficientvit_m1', 'efficientvit_b1', 'efficientvit_b2', 'efficientvit_l2',
]
for name in timm_models:
    cfg = timm.get_pretrained_cfg(name)
    if cfg.hf_hub_id:
        hf_hub_download(cfg.hf_hub_id, 'model.safetensors')
    elif cfg.url:
        torch.hub.download_url_to_file(cfg.url, os.path.join(cache, os.path.basename(cfg.url)))
    print(f'  ✓ {name}')

print(f'\n✓ All {len(urls) + len(timm_models)} pretrained weight files cached!')
"

Automatic GPU Optimizations:

WaveDL automatically enables performance optimizations for modern GPUs:

HPC Best Practices:

• Stage data to $SLURM_TMPDIR (local NVMe) for maximum I/O throughput
• Use --compile for training runs > 50 epochs
• Increase --workers manually if auto-detection is suboptimal

Environment Variables (for logging):

Example:

# Use Huber loss for noisy NDE data
wavedl-train --model cnn --loss huber --huber_delta 0.5

# Weighted MSE: prioritize thickness (first target)
wavedl-train --model cnn --loss weighted_mse --loss_weights "2.0,1.0,1.0"

# Dense maps: weight the first property channel twice as strongly
wavedl-train --model unet_field --loss weighted_mse --loss_weights "2.0,1.0" \
    --data_path field_data.npz --batch_size 4

Example:

# SGD with Nesterov momentum (often better generalization)
wavedl-train --model cnn --optimizer sgd --lr 0.01 --momentum 0.9 --nesterov

# RAdam for more stable training
wavedl-train --model cnn --optimizer radam --lr 1e-3

Example:

# Cosine annealing for 1000 epochs
wavedl-train --model cnn --scheduler cosine --epochs 1000 --min_lr 1e-7

# OneCycleLR for super-convergence
wavedl-train --model cnn --scheduler onecycle --lr 1e-2 --epochs 50

# MultiStep with custom milestones
wavedl-train --model cnn --scheduler multistep --milestones "100,200,300"

For robust model evaluation, simply add the --cv flag:

# 5-fold cross-validation
wavedl-train --model cnn --cv 5 --data_path train_data.npz

# Dense 2D field cross-validation (non-stratified)
wavedl-train --model unet_field --cv 5 --data_path field_data.npz \
    --batch_size 4

# Stratified CV for unbalanced vector targets
wavedl-train --model cnn --cv 5 --cv_stratify --loss huber --epochs 100

# Full configuration
wavedl-train --model cnn --cv 5 --cv_stratify \
    --loss huber --optimizer adamw --scheduler cosine \
    --output_dir ./cv_results

Output:

• cv_summary.json: Aggregated metrics (mean ± std)
• cv_results.csv: Per-fold detailed results
• fold_*/: Individual fold models and scalers

Dense field CV can require substantial host memory because folds are currently materialized in memory. Do not combine it with --cv_stratify; use smaller datasets/batches or run separate held-out experiments for very large field maps. CV preserves custom loss weights, Huber delta, and pretrained initialization. Physical constraints and gradient accumulation require standard training; CV rejects these options instead of silently ignoring them. CV runs in one process and currently selects folds using the configured loss (--selection_metric loss).

Use YAML files for reproducible experiments. CLI arguments can override any config value.

# Use a config file
wavedl-train --config configs/config.yaml --data_path train.npz

# Override specific values from config
wavedl-train --config configs/config.yaml --lr 5e-4 --epochs 500

Example config (configs/config.yaml):

# Model & Training
model: cnn
batch_size: 128
lr: 0.001
epochs: 1000
patience: 20

# Loss, Optimizer, Scheduler
loss: mse
optimizer: adamw
scheduler: plateau

# Cross-Validation (0 = disabled)
cv: 0

# Performance
precision: bf16
compile: false
seed: 2025

See configs/config.yaml for the complete template with all available options documented.

Add penalty terms to the loss function to enforce physical laws:

Total Loss = Data Loss + weight × penalty(violation)

Expression Constraints

# Positivity
--constraint "y0 > 0"

# Bounds
--constraint "y0 >= 0" "y0 <= 1"

# Equations (penalize deviations from zero)
--constraint "y2 - y0 * y1"

# Input-dependent constraints
--constraint "y0 - 2*x[0]"

# Multiple constraints with different weights
--constraint "y0 > 0" "y1 - y2" --constraint_weight 0.1 1.0

Custom Python Constraints

For complex physics (matrix operations, implicit equations):

# my_constraint.py
import torch

def constraint(pred, inputs=None):
    """
    Args:
        pred:   (batch, num_outputs) vectors or (batch, C, H, W) fields
        inputs: (batch, features) or (batch, C, H, W) or (batch, C, D, H, W)
    Returns:
        Violation tensor; WaveDL reduces all returned values to one penalty.
    """
    # For vectors, y0/y1/y2 are (batch,). For dense outputs they are
    # (batch, H, W), so this same expression applies at every pixel.
    y0, y1, y2 = pred[:, 0], pred[:, 1], pred[:, 2]

    # Inputs — Tabular: (batch, features)
    # x0 = inputs[:, 0]                    # Feature 0
    # x_sum = inputs.sum(dim=1)            # Sum all features

    # Inputs — Images: (batch, C, H, W)
    # pixel = inputs[:, 0, 3, 5]           # Pixel at (3,5), channel 0
    # img_mean = inputs.mean(dim=(1,2,3))  # Mean over C,H,W

    # Inputs — 3D Volumes: (batch, C, D, H, W)
    # voxel = inputs[:, 0, 2, 3, 5]        # Voxel at (2,3,5), channel 0

    # Example constraints:
    # return y2 - y0 * y1                                    # Wave equation
    # return y0 - 2 * inputs[:, 0]                           # Output = 2×input
    # return inputs[:, 0, 3, 5] * y0 + inputs[:, 0, 6, 7] * y1  # Mixed

    return y0 - y1 * y2

--constraint_file my_constraint.py --constraint_weight 1.0

Reference

Expression Syntax

Operators: +, -, *, /, **, >, <, >=, <=, ==

Functions: sin, cos, exp, log, sqrt, sigmoid, softplus, tanh, relu, abs

Automatically find the best training configuration using Optuna.

Run HPO:

# Basic HPO (50 trials, auto-detects GPUs)
wavedl-hpo --data_path train.npz --n_trials 50

# Quick search (minimal search space, fastest)
wavedl-hpo --data_path train.npz --n_trials 30 --quick

# Medium search (balanced between quick and full)
wavedl-hpo --data_path train.npz --n_trials 50 --medium

# Full search with specific models
wavedl-hpo --data_path train.npz --n_trials 100 --models cnn resnet18 efficientnet_b0

# Dense field search (auto-selects unet_field and small batch candidates)
wavedl-hpo --data_path field_data.npz --n_trials 30

# In-process mode (enables pruning, faster, single-GPU)
wavedl-hpo --data_path train.npz --n_trials 50 --inprocess

HPO inspects the target shape before creating a study. Vector data uses the standard CNN/ResNet model presets and batch candidates 16/32/64/128. Dense data uses unet_field and batch candidates 1/2/4/8. Explicit --models values are validated against the target contract before any trials launch; use --batch_sizes to override the memory-conscious dense defaults.

HPO compares trials using the same normalized validation MSE, even when their training losses differ. It seeds the sampler and records a study signature covering the dataset, search settings, and objective. Use a new study name when those settings change or when migrating a legacy study without this signature; mixing incompatible trials is rejected. Parallel trials lease distinct visible GPUs.

Train with best parameters

After HPO completes, it prints the optimal command:

wavedl-train --data_path train.npz --model cnn --lr 3.2e-4 --batch_size 128 ...

What Gets Searched:

Search Presets:

Execution Modes:

All Arguments:

See Available Models for all 20+ architectures (70+ variants) you can search.

import numpy as np

X = np.array(images, dtype=np.float32)  # (N, H, W)
y = np.array(labels, dtype=np.float32)  # (N, T)

np.savez('train_data.npz', input_train=X, output_train=y)

import numpy as np
from PIL import Image
from pathlib import Path
import pandas as pd

# Load images
images = [np.array(Image.open(f).convert('L'), dtype=np.float32)
          for f in sorted(Path("images/").glob("*.png"))]
X = np.stack(images)

# Load labels
y = pd.read_csv("labels.csv").values.astype(np.float32)

np.savez('train_data.npz', input_train=X, output_train=y)

import numpy as np
from scipy.io import loadmat

data = loadmat('simulation_data.mat')
X = data['spectrograms'].astype(np.float32)  # Adjust key
y = data['parameters'].astype(np.float32)

# Transpose if needed: (H, W, N) → (N, H, W)
if X.ndim == 3 and X.shape[2] < X.shape[0]:
    X = np.transpose(X, (2, 0, 1))

np.savez('train_data.npz', input_train=X, output_train=y)

import numpy as np

X = np.random.randn(1000, 256, 256).astype(np.float32)
y = np.random.randn(1000, 5).astype(np.float32)

np.savez('test_data.npz', input_test=X, output_test=y)

import numpy as np

data = np.load('train_data.npz')
assert data['input_train'].ndim >= 2, "Input must be at least 2D: (N, ...) "
assert data['output_train'].ndim in (1, 2, 3, 4), \
    "Output must be (N,), (N,T), (N,H,W), or (N,C,H,W)"
assert len(data['input_train']) == len(data['output_train']), "Sample mismatch"

print(f"✓ Input:  {data['input_train'].shape} {data['input_train'].dtype}")
print(f"✓ Output: {data['output_train'].shape} {data['output_train'].dtype}")