wavedl 1.2.0__py3-none-any.whl

wavedl/__init__.py

@@ -0,0 +1,43 @@
+"""
+WaveDL - Deep Learning Framework for Wave-Based Inverse Problems
+=================================================================
+
+A scalable deep learning framework for wave-based inverse problems,
+from ultrasonic NDE and geophysics to biomedical tissue characterization.
+
+Quick Start:
+    from wavedl.models import build_model, list_models
+    from wavedl.utils import prepare_data, load_test_data
+
+For training:
+    wavedl-train --model cnn --data_path train.npz
+    # or: python -m wavedl.train --model cnn --data_path train.npz
+
+For inference:
+    wavedl-test --checkpoint best_checkpoint --data_path test.npz
+    # or: python -m wavedl.test --checkpoint best_checkpoint --data_path test.npz
+"""
+
+__version__ = "1.2.0"
+__author__ = "Ductho Le"
+__email__ = "ductho.le@outlook.com"
+
+# Re-export key APIs for convenience
+from wavedl.models import build_model, get_model, list_models, register_model
+from wavedl.utils import (
+    load_test_data,
+    load_training_data,
+    prepare_data,
+)
+
+
+__all__ = [
+    "__version__",
+    "build_model",
+    "get_model",
+    "list_models",
+    "load_test_data",
+    "load_training_data",
+    "prepare_data",
+    "register_model",
+]

wavedl/hpo.py

@@ -0,0 +1,366 @@
+"""
+WaveDL - Hyperparameter Optimization with Optuna
+=================================================
+Automated hyperparameter search for finding optimal training configurations.
+
+Usage:
+    # Basic HPO (50 trials)
+    python hpo.py --data_path train.npz --n_trials 50
+
+    # Quick search (fewer parameters)
+    python hpo.py --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
+
+    # Parallel trials on multiple GPUs
+    python hpo.py --data_path train.npz --n_trials 100 --n_jobs 4
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+import argparse
+import json
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+
+try:
+    import optuna
+    from optuna.trial import TrialState
+except ImportError:
+    print("Error: Optuna not installed. Run: pip install -e '.[hpo]'")
+    sys.exit(1)
+
+
+# =============================================================================
+# DEFAULT SEARCH SPACES
+# =============================================================================
+
+DEFAULT_MODELS = ["cnn", "resnet18", "resnet34"]
+QUICK_MODELS = ["cnn"]
+
+# All 6 optimizers
+DEFAULT_OPTIMIZERS = ["adamw", "adam", "sgd", "nadam", "radam", "rmsprop"]
+QUICK_OPTIMIZERS = ["adamw"]
+
+# All 8 schedulers
+DEFAULT_SCHEDULERS = [
+    "plateau",
+    "cosine",
+    "cosine_restarts",
+    "onecycle",
+    "step",
+    "multistep",
+    "exponential",
+    "linear_warmup",
+]
+QUICK_SCHEDULERS = ["plateau"]
+
+# All 6 losses
+DEFAULT_LOSSES = ["mse", "mae", "huber", "smooth_l1", "log_cosh", "weighted_mse"]
+QUICK_LOSSES = ["mse"]
+
+
+# =============================================================================
+# OBJECTIVE FUNCTION
+# =============================================================================
+
+
+def create_objective(args):
+    """Create Optuna objective function with configurable search space."""
+
+    def objective(trial):
+        # Select search space based on mode
+        # CLI arguments always take precedence over defaults
+        if args.quick:
+            models = args.models or QUICK_MODELS
+            optimizers = args.optimizers or QUICK_OPTIMIZERS
+            schedulers = args.schedulers or QUICK_SCHEDULERS
+            losses = args.losses or QUICK_LOSSES
+        else:
+            models = args.models or DEFAULT_MODELS
+            optimizers = args.optimizers or DEFAULT_OPTIMIZERS
+            schedulers = args.schedulers or DEFAULT_SCHEDULERS
+            losses = args.losses or DEFAULT_LOSSES
+
+        # Suggest hyperparameters
+        model = trial.suggest_categorical("model", models)
+        lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True)
+        batch_size = trial.suggest_categorical("batch_size", [64, 128, 256, 512])
+        optimizer = trial.suggest_categorical("optimizer", optimizers)
+        scheduler = trial.suggest_categorical("scheduler", schedulers)
+        loss = trial.suggest_categorical("loss", losses)
+        weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-3, log=True)
+        patience = trial.suggest_int("patience", 10, 30, step=5)
+
+        # Conditional hyperparameters
+        if loss == "huber":
+            huber_delta = trial.suggest_float("huber_delta", 0.1, 2.0)
+        else:
+            huber_delta = None
+
+        if optimizer == "sgd":
+            momentum = trial.suggest_float("momentum", 0.8, 0.99)
+        else:
+            momentum = None
+
+        # Build command
+        cmd = [
+            sys.executable,
+            "-m",
+            "wavedl.train",
+            "--data_path",
+            str(args.data_path),
+            "--model",
+            model,
+            "--lr",
+            str(lr),
+            "--batch_size",
+            str(batch_size),
+            "--optimizer",
+            optimizer,
+            "--scheduler",
+            scheduler,
+            "--loss",
+            loss,
+            "--weight_decay",
+            str(weight_decay),
+            "--patience",
+            str(patience),
+            "--epochs",
+            str(args.max_epochs),
+            "--seed",
+            str(args.seed),
+        ]
+
+        # Add conditional args
+        if huber_delta:
+            cmd.extend(["--huber_delta", str(huber_delta)])
+        if momentum:
+            cmd.extend(["--momentum", str(momentum)])
+
+        # Use temporary directory for trial output
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cmd.extend(["--output_dir", tmpdir])
+
+            # Run training
+            try:
+                result = subprocess.run(
+                    cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=args.timeout,
+                    cwd=Path(__file__).parent,
+                )
+
+                # Parse validation loss from output
+                # Look for "Best val_loss: X.XXXX" in stdout
+                val_loss = None
+                for line in result.stdout.split("\n"):
+                    if "Best val_loss:" in line:
+                        try:
+                            val_loss = float(line.split(":")[-1].strip())
+                        except ValueError:
+                            pass
+                    # Also check for final validation loss
+                    if "val_loss=" in line.lower():
+                        try:
+                            # Extract number after val_loss=
+                            parts = line.lower().split("val_loss=")
+                            if len(parts) > 1:
+                                val_str = parts[1].split()[0].strip(",")
+                                val_loss = float(val_str)
+                        except (ValueError, IndexError):
+                            pass
+
+                if val_loss is None:
+                    # Training failed or no loss found
+                    print(f"Trial {trial.number}: Training failed")
+                    return float("inf")
+
+                print(f"Trial {trial.number}: val_loss={val_loss:.6f}")
+                return val_loss
+
+            except subprocess.TimeoutExpired:
+                print(f"Trial {trial.number}: Timeout after {args.timeout}s")
+                return float("inf")
+            except Exception as e:
+                print(f"Trial {trial.number}: Error - {e}")
+                return float("inf")
+
+    return objective
+
+
+# =============================================================================
+# MAIN FUNCTION
+# =============================================================================
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="WaveDL Hyperparameter Optimization with Optuna",
+        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
+        """,
+    )
+
+    # Required
+    parser.add_argument(
+        "--data_path", type=str, required=True, help="Path to training data"
+    )
+
+    # HPO settings
+    parser.add_argument(
+        "--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)"
+    )
+    parser.add_argument(
+        "--quick",
+        action="store_true",
+        help="Quick mode: search fewer parameters",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=int,
+        default=3600,
+        help="Timeout per trial in seconds (default: 3600)",
+    )
+
+    # Search space customization
+    parser.add_argument(
+        "--models",
+        nargs="+",
+        default=None,
+        help=f"Models to search (default: {DEFAULT_MODELS})",
+    )
+    parser.add_argument(
+        "--optimizers",
+        nargs="+",
+        default=None,
+        help=f"Optimizers to search (default: {DEFAULT_OPTIMIZERS})",
+    )
+    parser.add_argument(
+        "--schedulers",
+        nargs="+",
+        default=None,
+        help=f"Schedulers to search (default: {DEFAULT_SCHEDULERS})",
+    )
+    parser.add_argument(
+        "--losses",
+        nargs="+",
+        default=None,
+        help=f"Losses to search (default: {DEFAULT_LOSSES})",
+    )
+
+    # Training settings for each trial
+    parser.add_argument(
+        "--max_epochs",
+        type=int,
+        default=50,
+        help="Max epochs per trial (default: 50, use early stopping)",
+    )
+    parser.add_argument(
+        "--seed", type=int, default=2025, help="Random seed (default: 2025)"
+    )
+
+    # Output
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="hpo_results.json",
+        help="Output file for best params (default: hpo_results.json)",
+    )
+    parser.add_argument(
+        "--study_name",
+        type=str,
+        default="wavedl_hpo",
+        help="Optuna study name (default: wavedl_hpo)",
+    )
+
+    args = parser.parse_args()
+
+    # Validate data path
+    if not Path(args.data_path).exists():
+        print(f"Error: Data file not found: {args.data_path}")
+        sys.exit(1)
+
+    # Create study
+    print("=" * 60)
+    print("WaveDL Hyperparameter Optimization")
+    print("=" * 60)
+    print(f"Data: {args.data_path}")
+    print(f"Trials: {args.n_trials}")
+    print(f"Mode: {'Quick' if args.quick else 'Full'}")
+    print(f"Parallel jobs: {args.n_jobs}")
+    print("=" * 60)
+
+    study = optuna.create_study(
+        study_name=args.study_name,
+        direction="minimize",
+        pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=10),
+    )
+
+    # Run optimization
+    objective = create_objective(args)
+    study.optimize(
+        objective,
+        n_trials=args.n_trials,
+        n_jobs=args.n_jobs,
+        show_progress_bar=True,
+    )
+
+    # Results
+    print("\n" + "=" * 60)
+    print("OPTIMIZATION COMPLETE")
+    print("=" * 60)
+
+    # Filter completed trials
+    completed_trials = [t for t in study.trials if t.state == TrialState.COMPLETE]
+
+    if not completed_trials:
+        print("No trials completed successfully.")
+        sys.exit(1)
+
+    print(f"\nCompleted trials: {len(completed_trials)}/{args.n_trials}")
+    print(f"Best trial: #{study.best_trial.number}")
+    print(f"Best val_loss: {study.best_value:.6f}")
+
+    print("\nBest hyperparameters:")
+    for key, value in study.best_params.items():
+        print(f"  {key}: {value}")
+
+    # Save results
+    results = {
+        "best_value": study.best_value,
+        "best_params": study.best_params,
+        "n_trials": len(completed_trials),
+        "study_name": args.study_name,
+    }
+
+    with open(args.output, "w") as f:
+        json.dump(results, f, indent=2)
+
+    print(f"\nResults saved to: {args.output}")
+
+    # Print command to train with best params
+    print("\n" + "=" * 60)
+    print("TO TRAIN WITH BEST PARAMETERS:")
+    print("=" * 60)
+    cmd_parts = ["accelerate launch train.py"]
+    cmd_parts.append(f"--data_path {args.data_path}")
+    for key, value in study.best_params.items():
+        cmd_parts.append(f"--{key} {value}")
+    print(" \\\n    ".join(cmd_parts))
+
+
+if __name__ == "__main__":
+    main()

wavedl/models/__init__.py

@@ -0,0 +1,86 @@
+"""
+Model Registry and Factory Pattern for Deep Learning Architectures
+===================================================================
+
+This module provides a centralized registry for neural network architectures,
+enabling dynamic model selection via command-line arguments.
+
+Usage:
+    from wavedl.models import get_model, list_models, MODEL_REGISTRY
+
+    # List available models
+    print(list_models())
+
+    # Get a model class by name
+    ModelClass = get_model("cnn")
+    model = ModelClass(in_shape=(500, 500), out_size=5)
+
+Adding New Models:
+    1. Create a new file in models/ (e.g., models/my_model.py)
+    2. Inherit from BaseModel
+    3. Use the @register_model decorator
+
+    Example:
+        from wavedl.models.base import BaseModel
+        from wavedl.models.registry import register_model
+
+        @register_model("my_model")
+        class MyModel(BaseModel):
+            def __init__(self, in_shape, out_size, **kwargs):
+                super().__init__(in_shape, out_size)
+                ...
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+# Import registry first (no dependencies)
+# Import base class (depends only on torch)
+from .base import BaseModel
+
+# Import model implementations (triggers registration via decorators)
+from .cnn import CNN
+from .convnext import ConvNeXtBase_, ConvNeXtSmall, ConvNeXtTiny
+from .densenet import DenseNet121, DenseNet169
+from .efficientnet import EfficientNetB0, EfficientNetB1, EfficientNetB2
+from .registry import (
+    MODEL_REGISTRY,
+    build_model,
+    get_model,
+    list_models,
+    register_model,
+)
+from .resnet import ResNet18, ResNet34, ResNet50
+from .unet import UNet, UNetRegression
+from .vit import ViTBase_, ViTSmall, ViTTiny
+
+
+# Export public API
+__all__ = [
+    # Models
+    "CNN",
+    # Registry
+    "MODEL_REGISTRY",
+    # Base class
+    "BaseModel",
+    "ConvNeXtBase_",
+    "ConvNeXtSmall",
+    "ConvNeXtTiny",
+    "DenseNet121",
+    "DenseNet169",
+    "EfficientNetB0",
+    "EfficientNetB1",
+    "EfficientNetB2",
+    "ResNet18",
+    "ResNet34",
+    "ResNet50",
+    "UNet",
+    "UNetRegression",
+    "ViTBase_",
+    "ViTSmall",
+    "ViTTiny",
+    "build_model",
+    "get_model",
+    "list_models",
+    "register_model",
+]

wavedl/models/_template.py

@@ -0,0 +1,157 @@
+"""
+Model Template for New Architectures
+=====================================
+
+Copy this file and modify to add new model architectures to the framework.
+The model will be automatically registered and available via --model flag.
+
+Steps to Add a New Model:
+    1. Copy this file to models/your_model.py
+    2. Rename the class and update @register_model("your_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
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from wavedl.models.base import BaseModel
+
+
+# Uncomment the decorator to register this model
+# @register_model("template")
+class TemplateModel(BaseModel):
+    """
+    Template Model Architecture.
+
+    Replace this docstring with your model description.
+    The first line will appear in --list_models output.
+
+    Args:
+        in_shape: Input spatial dimensions (H, W)
+        out_size: Number of regression output targets
+        hidden_dim: Size of hidden layers (default: 256)
+        num_layers: Number of convolutional layers (default: 4)
+        dropout: Dropout rate (default: 0.1)
+
+    Input Shape:
+        (B, 1, H, W) - Single-channel images
+
+    Output Shape:
+        (B, out_size) - Regression predictions
+    """
+
+    def __init__(
+        self,
+        in_shape: tuple[int, int],
+        out_size: int,
+        hidden_dim: int = 256,
+        num_layers: int = 4,
+        dropout: float = 0.1,
+        **kwargs,  # Accept extra kwargs for flexibility
+    ):
+        # REQUIRED: Call parent __init__ with in_shape and out_size
+        super().__init__(in_shape, out_size)
+
+        # Store hyperparameters as attributes (optional but recommended)
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.dropout_rate = dropout
+
+        # =================================================================
+        # BUILD YOUR ARCHITECTURE HERE
+        # =================================================================
+
+        # Example: Simple CNN encoder
+        self.encoder = nn.Sequential(
+            # Layer 1
+            nn.Conv2d(1, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            # Layer 2
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            # Layer 3
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            # Layer 4
+            nn.Conv2d(128, hidden_dim, kernel_size=3, padding=1),
+            nn.BatchNorm2d(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d(1),  # Global average pooling
+        )
+
+        # Example: Regression head
+        self.head = nn.Sequential(
+            nn.Flatten(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim // 2, out_size),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the model.
+
+        REQUIRED: Must accept (B, C, H, W) and return (B, out_size)
+
+        Args:
+            x: Input tensor of shape (B, 1, H, W)
+
+        Returns:
+            Output tensor of shape (B, out_size)
+        """
+        # Encode
+        features = self.encoder(x)
+
+        # Predict
+        output = self.head(features)
+
+        return output
+
+    @classmethod
+    def get_default_config(cls) -> dict[str, Any]:
+        """
+        Return default hyperparameters for this model.
+
+        OPTIONAL: Override to provide model-specific defaults.
+        These can be used for documentation or config files.
+        """
+        return {
+            "hidden_dim": 256,
+            "num_layers": 4,
+            "dropout": 0.1,
+        }
+
+
+# =============================================================================
+# USAGE EXAMPLE
+# =============================================================================
+if __name__ == "__main__":
+    # Quick test of the model
+    model = TemplateModel(in_shape=(500, 500), out_size=5)
+
+    # Print model summary
+    print(f"Model: {model.__class__.__name__}")
+    print(f"Parameters: {model.count_parameters():,}")
+    print(f"Default config: {model.get_default_config()}")
+
+    # Test forward pass
+    dummy_input = torch.randn(2, 1, 500, 500)
+    output = model(dummy_input)
+    print(f"Input shape: {dummy_input.shape}")
+    print(f"Output shape: {output.shape}")