wavedl 1.2.0__py3-none-any.whl

wavedl/utils/distributed.py

@@ -0,0 +1,138 @@
+"""
+Distributed Training Utilities
+==============================
+
+Provides DDP-safe utilities for multi-GPU training including:
+- Early stopping synchronization across ranks
+- Value broadcasting from rank 0
+- Tensor synchronization
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+import torch
+import torch.distributed as dist
+from accelerate import Accelerator
+
+
+def broadcast_early_stop(should_stop: bool, accelerator: Accelerator) -> bool:
+    """
+    Broadcast early stopping decision from rank 0 to all processes.
+
+    In DDP training, early stopping state (patience counter, best loss) is typically
+    tracked only on rank 0. This function ensures all ranks receive the same stop signal
+    to prevent deadlocks from inconsistent termination.
+
+    Args:
+        should_stop: Whether to stop training (only meaningful on rank 0)
+        accelerator: Accelerator instance for device and process info
+
+    Returns:
+        True if training should stop (synchronized across all ranks)
+
+    Example:
+        # On main process: patience_ctr >= patience
+        # On other processes: unknown/stale value
+        should_stop = patience_ctr >= args.patience if accelerator.is_main_process else False
+        if broadcast_early_stop(should_stop, accelerator):
+            break  # All ranks exit loop together
+    """
+    stop_tensor = torch.tensor(
+        1 if should_stop else 0, device=accelerator.device, dtype=torch.int32
+    )
+
+    if accelerator.num_processes > 1:
+        dist.broadcast(stop_tensor, src=0)
+
+    return stop_tensor.item() == 1
+
+
+def broadcast_value(value: int | float, accelerator: Accelerator) -> int | float:
+    """
+    Broadcast a scalar value from rank 0 to all processes.
+
+    Useful for synchronizing hyperparameters or computed values across ranks.
+
+    Args:
+        value: Scalar value to broadcast (only rank 0's value is used)
+        accelerator: Accelerator instance for device and process info
+
+    Returns:
+        Value from rank 0 (synchronized across all ranks)
+    """
+    is_int = isinstance(value, int)
+    dtype = torch.int64 if is_int else torch.float32
+
+    tensor = torch.tensor(value, device=accelerator.device, dtype=dtype)
+
+    if accelerator.num_processes > 1:
+        dist.broadcast(tensor, src=0)
+
+    result = tensor.item()
+    return int(result) if is_int else result
+
+
+def sync_tensor(
+    tensor: torch.Tensor, accelerator: Accelerator, reduction: str = "sum"
+) -> torch.Tensor:
+    """
+    Synchronize a tensor across all processes with specified reduction.
+
+    Wrapper around accelerator.reduce with additional validation.
+
+    Args:
+        tensor: Tensor to synchronize
+        accelerator: Accelerator instance
+        reduction: Reduction operation ("sum", "mean", "max", "min")
+
+    Returns:
+        Reduced tensor (synchronized across all ranks)
+
+    Raises:
+        ValueError: If reduction type is not recognized
+    """
+    valid_reductions = {"sum", "mean", "max", "min"}
+    if reduction not in valid_reductions:
+        raise ValueError(
+            f"Invalid reduction '{reduction}'. Must be one of {valid_reductions}"
+        )
+
+    return accelerator.reduce(tensor, reduction=reduction)
+
+
+def get_world_info(accelerator: Accelerator) -> dict:
+    """
+    Get distributed training world information.
+
+    Args:
+        accelerator: Accelerator instance
+
+    Returns:
+        Dictionary with world_size, rank, local_rank, is_main, device
+    """
+    return {
+        "world_size": accelerator.num_processes,
+        "rank": accelerator.process_index,
+        "local_rank": accelerator.local_process_index,
+        "is_main": accelerator.is_main_process,
+        "device": str(accelerator.device),
+    }
+
+
+def print_rank0(message: str, accelerator: Accelerator, logger=None):
+    """
+    Print message only on rank 0.
+
+    Convenience function for logging in distributed setting.
+
+    Args:
+        message: Message to print
+        accelerator: Accelerator instance
+        logger: Optional logger (uses print if None)
+    """
+    if accelerator.is_main_process:
+        if logger:
+            logger.info(message)
+        else:
+            print(message)

wavedl/utils/losses.py

@@ -0,0 +1,216 @@
+"""
+Loss Functions for Regression Tasks
+====================================
+
+Provides a comprehensive set of loss functions for regression problems,
+with a factory function for easy selection via CLI arguments.
+
+Supported Losses:
+    - mse: Mean Squared Error (default)
+    - mae: Mean Absolute Error (L1)
+    - huber: Huber Loss (smooth blend of MSE and MAE)
+    - smooth_l1: Smooth L1 Loss (PyTorch native Huber variant)
+    - log_cosh: Log-Cosh Loss (smooth approximation to MAE)
+    - weighted_mse: Per-target weighted MSE
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+import torch
+import torch.nn as nn
+
+
+# ==============================================================================
+# CUSTOM LOSS FUNCTIONS
+# ==============================================================================
+class LogCoshLoss(nn.Module):
+    """
+    Log-Cosh Loss: A smooth approximation to Mean Absolute Error.
+
+    The loss is defined as: loss = log(cosh(pred - target))
+
+    Properties:
+        - Smooth everywhere (differentiable)
+        - Behaves like L2 for small errors, L1 for large errors
+        - More robust to outliers than MSE
+
+    Example:
+        >>> criterion = LogCoshLoss()
+        >>> loss = criterion(predictions, targets)
+    """
+
+    def __init__(self, reduction: str = "mean"):
+        """
+        Args:
+            reduction: Specifies the reduction: 'none' | 'mean' | 'sum'
+        """
+        super().__init__()
+        if reduction not in ("none", "mean", "sum"):
+            raise ValueError(f"Invalid reduction mode: {reduction}")
+        self.reduction = reduction
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """
+        Compute Log-Cosh loss.
+
+        Args:
+            pred: Predicted values of shape (N, *)
+            target: Target values of shape (N, *)
+
+        Returns:
+            Loss value (scalar if reduction is 'mean' or 'sum')
+        """
+        diff = pred - target
+        # log(cosh(x)) = x + softplus(-2x) - log(2)
+        # This formulation is numerically stable
+        loss = diff + torch.nn.functional.softplus(-2.0 * diff) - 0.693147  # log(2)
+
+        if self.reduction == "none":
+            return loss
+        elif self.reduction == "sum":
+            return loss.sum()
+        else:  # mean
+            return loss.mean()
+
+
+class WeightedMSELoss(nn.Module):
+    """
+    Weighted Mean Squared Error Loss.
+
+    Applies different weights to each target dimension, allowing
+    prioritization of specific outputs (e.g., prioritize thickness
+    over velocity in NDE applications).
+
+    Example:
+        >>> # 3 targets, prioritize first target
+        >>> criterion = WeightedMSELoss(weights=[2.0, 1.0, 1.0])
+        >>> loss = criterion(predictions, targets)
+    """
+
+    def __init__(
+        self, weights: list[float] | torch.Tensor | None = None, reduction: str = "mean"
+    ):
+        """
+        Args:
+            weights: Per-target weights. If None, equal weights (standard MSE).
+                     Length must match number of output targets.
+            reduction: Specifies the reduction: 'none' | 'mean' | 'sum'
+        """
+        super().__init__()
+        if reduction not in ("none", "mean", "sum"):
+            raise ValueError(f"Invalid reduction mode: {reduction}")
+        self.reduction = reduction
+
+        if weights is not None:
+            if isinstance(weights, list):
+                weights = torch.tensor(weights, dtype=torch.float32)
+            self.register_buffer("weights", weights)
+        else:
+            self.weights = None
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """
+        Compute weighted MSE loss.
+
+        Args:
+            pred: Predicted values of shape (N, T) where T is number of targets
+            target: Target values of shape (N, T)
+
+        Returns:
+            Loss value (scalar if reduction is 'mean' or 'sum')
+
+        Raises:
+            ValueError: If weight dimension doesn't match target dimension
+        """
+        mse = (pred - target) ** 2
+
+        if self.weights is not None:
+            # Validate weight dimension matches target dimension
+            if self.weights.shape[0] != pred.shape[-1]:
+                raise ValueError(
+                    f"Weight dimension ({self.weights.shape[0]}) must match "
+                    f"output dimension ({pred.shape[-1]}). "
+                    f"Check your --loss_weights argument."
+                )
+            # Use local variable to avoid mutating registered buffer
+            # (mutating self.weights breaks state_dict semantics)
+            weights = self.weights.to(mse.device)
+            # Apply per-target weights with correct broadcasting: (N, T) * (T,) -> (N, T)
+            mse = mse * weights
+
+        if self.reduction == "none":
+            return mse
+        elif self.reduction == "sum":
+            return mse.sum()
+        else:  # mean
+            return mse.mean()
+
+
+# ==============================================================================
+# LOSS REGISTRY
+# ==============================================================================
+_LOSS_REGISTRY = {
+    "mse": nn.MSELoss,
+    "mae": nn.L1Loss,
+    "l1": nn.L1Loss,  # Alias for mae
+    "huber": nn.HuberLoss,
+    "smooth_l1": nn.SmoothL1Loss,
+    "log_cosh": LogCoshLoss,
+    "logcosh": LogCoshLoss,  # Alias
+    "weighted_mse": WeightedMSELoss,
+}
+
+
+def list_losses() -> list[str]:
+    """
+    Return list of available loss function names.
+
+    Returns:
+        List of registered loss function names (excluding aliases)
+    """
+    # Return unique loss names (exclude aliases)
+    primary_names = ["mse", "mae", "huber", "smooth_l1", "log_cosh", "weighted_mse"]
+    return primary_names
+
+
+def get_loss(
+    name: str, weights: list[float] | None = None, delta: float = 1.0, **kwargs
+) -> nn.Module:
+    """
+    Factory function to create loss function by name.
+
+    Args:
+        name: Loss function name (see list_losses())
+        weights: Per-target weights for weighted_mse
+        delta: Delta parameter for Huber loss (default: 1.0)
+        **kwargs: Additional arguments passed to loss constructor
+
+    Returns:
+        Instantiated loss function (nn.Module)
+
+    Raises:
+        ValueError: If loss name is not recognized
+
+    Example:
+        >>> criterion = get_loss("mse")
+        >>> criterion = get_loss("huber", delta=0.5)
+        >>> criterion = get_loss("weighted_mse", weights=[2.0, 1.0, 1.0])
+    """
+    name_lower = name.lower().replace("-", "_")
+
+    if name_lower not in _LOSS_REGISTRY:
+        available = ", ".join(list_losses())
+        raise ValueError(
+            f"Unknown loss function: '{name}'. Available options: {available}"
+        )
+
+    loss_cls = _LOSS_REGISTRY[name_lower]
+
+    # Special handling for specific loss types
+    if name_lower == "huber":
+        return loss_cls(delta=delta, **kwargs)
+    elif name_lower == "weighted_mse":
+        return loss_cls(weights=weights, **kwargs)
+    else:
+        return loss_cls(**kwargs)