wavedl 1.2.0__py3-none-any.whl

wavedl/utils/optimizers.py

@@ -0,0 +1,216 @@
+"""
+Optimizers for Deep Learning Training
+=====================================
+
+Provides a comprehensive set of optimizers with a factory function
+for easy selection via CLI arguments.
+
+Supported Optimizers:
+    - adamw: AdamW (default, best for most cases)
+    - adam: Adam (legacy)
+    - sgd: SGD with momentum
+    - nadam: NAdam (Adam + Nesterov momentum)
+    - radam: RAdam (variance-adaptive Adam)
+    - rmsprop: RMSprop (good for RNNs)
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+from collections.abc import Iterator
+
+import torch
+import torch.optim as optim
+from torch.nn import Parameter
+
+
+# ==============================================================================
+# OPTIMIZER REGISTRY
+# ==============================================================================
+def list_optimizers() -> list[str]:
+    """
+    Return list of available optimizer names.
+
+    Returns:
+        List of registered optimizer names
+    """
+    return ["adamw", "adam", "sgd", "nadam", "radam", "rmsprop"]
+
+
+def get_optimizer(
+    name: str,
+    params: Iterator[Parameter] | list[dict],
+    lr: float = 1e-3,
+    weight_decay: float = 1e-4,
+    # SGD-specific
+    momentum: float = 0.9,
+    nesterov: bool = False,
+    dampening: float = 0.0,
+    # Adam/AdamW-specific
+    betas: tuple = (0.9, 0.999),
+    eps: float = 1e-8,
+    amsgrad: bool = False,
+    # RMSprop-specific
+    alpha: float = 0.99,
+    centered: bool = False,
+    **kwargs,
+) -> optim.Optimizer:
+    """
+    Factory function to create optimizer by name.
+
+    Args:
+        name: Optimizer name (see list_optimizers())
+        params: Model parameters or parameter groups
+        lr: Learning rate
+        weight_decay: Weight decay (L2 penalty)
+        momentum: Momentum factor (SGD, RMSprop)
+        nesterov: Enable Nesterov momentum (SGD)
+        dampening: Dampening for momentum (SGD)
+        betas: Coefficients for computing running averages (Adam variants)
+        eps: Term for numerical stability (Adam variants, RMSprop)
+        amsgrad: Use AMSGrad variant (Adam variants)
+        alpha: Smoothing constant (RMSprop)
+        centered: Compute centered RMSprop
+        **kwargs: Additional optimizer-specific arguments
+
+    Returns:
+        Instantiated optimizer
+
+    Raises:
+        ValueError: If optimizer name is not recognized
+
+    Example:
+        >>> optimizer = get_optimizer("adamw", model.parameters(), lr=1e-3)
+        >>> optimizer = get_optimizer("sgd", model.parameters(), lr=1e-2, momentum=0.9)
+    """
+    name_lower = name.lower()
+
+    if name_lower == "adamw":
+        return optim.AdamW(
+            params,
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            amsgrad=amsgrad,
+            **kwargs,
+        )
+
+    elif name_lower == "adam":
+        return optim.Adam(
+            params,
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            amsgrad=amsgrad,
+            **kwargs,
+        )
+
+    elif name_lower == "sgd":
+        return optim.SGD(
+            params,
+            lr=lr,
+            momentum=momentum,
+            dampening=dampening,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            **kwargs,
+        )
+
+    elif name_lower == "nadam":
+        return optim.NAdam(
+            params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, **kwargs
+        )
+
+    elif name_lower == "radam":
+        return optim.RAdam(
+            params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, **kwargs
+        )
+
+    elif name_lower == "rmsprop":
+        return optim.RMSprop(
+            params,
+            lr=lr,
+            alpha=alpha,
+            eps=eps,
+            weight_decay=weight_decay,
+            momentum=momentum,
+            centered=centered,
+            **kwargs,
+        )
+
+    else:
+        available = ", ".join(list_optimizers())
+        raise ValueError(f"Unknown optimizer: '{name}'. Available options: {available}")
+
+
+def get_optimizer_with_param_groups(
+    name: str,
+    model: torch.nn.Module,
+    lr: float = 1e-3,
+    weight_decay: float = 1e-4,
+    no_decay_keywords: list[str] = None,
+    **kwargs,
+) -> optim.Optimizer:
+    """
+    Create optimizer with automatic parameter grouping.
+
+    Separates parameters into decay and no-decay groups based on
+    parameter names. By default, bias and normalization layer parameters
+    are excluded from weight decay.
+
+    Args:
+        name: Optimizer name
+        model: PyTorch model
+        lr: Learning rate
+        weight_decay: Weight decay for applicable parameters
+        no_decay_keywords: Keywords to identify no-decay parameters
+                          Default: ['bias', 'norm', 'bn', 'ln']
+        **kwargs: Additional optimizer arguments
+
+    Returns:
+        Optimizer with parameter groups configured
+
+    Example:
+        >>> optimizer = get_optimizer_with_param_groups(
+        ...     "adamw", model, lr=1e-3, weight_decay=1e-4
+        ... )
+    """
+    if no_decay_keywords is None:
+        no_decay_keywords = ["bias", "norm", "bn", "ln"]
+
+    decay_params = []
+    no_decay_params = []
+
+    for name_param, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        # Check if any no-decay keyword is in parameter name
+        if any(kw in name_param.lower() for kw in no_decay_keywords):
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+
+    param_groups = []
+    if decay_params:
+        param_groups.append(
+            {
+                "params": decay_params,
+                "weight_decay": weight_decay,
+            }
+        )
+    if no_decay_params:
+        param_groups.append(
+            {
+                "params": no_decay_params,
+                "weight_decay": 0.0,
+            }
+        )
+
+    # Fall back to all parameters if grouping fails
+    if not param_groups:
+        param_groups = [{"params": model.parameters()}]
+
+    return get_optimizer(name, param_groups, lr=lr, weight_decay=0.0, **kwargs)

wavedl/utils/schedulers.py

@@ -0,0 +1,251 @@
+"""
+Learning Rate Schedulers
+========================
+
+Provides a comprehensive set of learning rate schedulers with a factory
+function for easy selection via CLI arguments.
+
+Supported Schedulers:
+    - plateau: ReduceLROnPlateau (default, adaptive)
+    - cosine: CosineAnnealingLR
+    - cosine_restarts: CosineAnnealingWarmRestarts
+    - onecycle: OneCycleLR
+    - step: StepLR
+    - multistep: MultiStepLR
+    - exponential: ExponentialLR
+    - linear_warmup: LinearLR (warmup phase)
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+import torch.optim as optim
+from torch.optim.lr_scheduler import (
+    CosineAnnealingLR,
+    CosineAnnealingWarmRestarts,
+    ExponentialLR,
+    LinearLR,
+    LRScheduler,
+    MultiStepLR,
+    OneCycleLR,
+    ReduceLROnPlateau,
+    SequentialLR,
+    StepLR,
+)
+
+
+# ==============================================================================
+# SCHEDULER REGISTRY
+# ==============================================================================
+def list_schedulers() -> list[str]:
+    """
+    Return list of available scheduler names.
+
+    Returns:
+        List of registered scheduler names
+    """
+    return [
+        "plateau",
+        "cosine",
+        "cosine_restarts",
+        "onecycle",
+        "step",
+        "multistep",
+        "exponential",
+        "linear_warmup",
+    ]
+
+
+def get_scheduler(
+    name: str,
+    optimizer: optim.Optimizer,
+    # Common parameters
+    epochs: int = 100,
+    steps_per_epoch: int | None = None,
+    min_lr: float = 1e-6,
+    # ReduceLROnPlateau parameters
+    patience: int = 10,
+    factor: float = 0.5,
+    # Cosine parameters
+    T_max: int | None = None,
+    T_0: int = 10,
+    T_mult: int = 2,
+    # OneCycleLR parameters
+    max_lr: float | None = None,
+    pct_start: float = 0.3,
+    # Step/MultiStep parameters
+    step_size: int = 30,
+    milestones: list[int] | None = None,
+    gamma: float = 0.1,
+    # Linear warmup parameters
+    warmup_epochs: int = 5,
+    start_factor: float = 0.1,
+    **kwargs,
+) -> LRScheduler:
+    """
+    Factory function to create learning rate scheduler by name.
+
+    Args:
+        name: Scheduler name (see list_schedulers())
+        optimizer: Optimizer instance to schedule
+        epochs: Total training epochs (for cosine, onecycle)
+        steps_per_epoch: Steps per epoch (required for onecycle)
+        min_lr: Minimum learning rate (eta_min for cosine)
+        patience: Patience for ReduceLROnPlateau
+        factor: Reduction factor for plateau/step
+        T_max: Period for CosineAnnealingLR (default: epochs)
+        T_0: Initial period for CosineAnnealingWarmRestarts
+        T_mult: Period multiplier for warm restarts
+        max_lr: Maximum LR for OneCycleLR (default: optimizer's initial LR)
+        pct_start: Percentage of cycle spent increasing LR (OneCycleLR)
+        step_size: Period for StepLR
+        milestones: Epochs to decay LR for MultiStepLR
+        gamma: Decay factor for step/multistep/exponential
+        warmup_epochs: Number of warmup epochs for linear_warmup
+        start_factor: Starting LR factor for warmup (LR * start_factor)
+        **kwargs: Additional arguments passed to scheduler
+
+    Returns:
+        Instantiated learning rate scheduler
+
+    Raises:
+        ValueError: If scheduler name is not recognized
+
+    Example:
+        >>> scheduler = get_scheduler("plateau", optimizer, patience=15)
+        >>> scheduler = get_scheduler("cosine", optimizer, epochs=100)
+        >>> scheduler = get_scheduler(
+        ...     "onecycle", optimizer, epochs=100, steps_per_epoch=1000, max_lr=1e-3
+        ... )
+    """
+    name_lower = name.lower().replace("-", "_")
+
+    # Get initial LR from optimizer
+    base_lr = optimizer.param_groups[0]["lr"]
+
+    if name_lower == "plateau":
+        return ReduceLROnPlateau(
+            optimizer,
+            mode="min",
+            factor=factor,
+            patience=patience,
+            min_lr=min_lr,
+            **kwargs,
+        )
+
+    elif name_lower == "cosine":
+        return CosineAnnealingLR(
+            optimizer,
+            T_max=T_max if T_max is not None else epochs,
+            eta_min=min_lr,
+            **kwargs,
+        )
+
+    elif name_lower == "cosine_restarts":
+        return CosineAnnealingWarmRestarts(
+            optimizer, T_0=T_0, T_mult=T_mult, eta_min=min_lr, **kwargs
+        )
+
+    elif name_lower == "onecycle":
+        if steps_per_epoch is None:
+            raise ValueError(
+                "OneCycleLR requires 'steps_per_epoch'. "
+                "Pass len(train_dataloader) as steps_per_epoch."
+            )
+        return OneCycleLR(
+            optimizer,
+            max_lr=max_lr if max_lr is not None else base_lr,
+            epochs=epochs,
+            steps_per_epoch=steps_per_epoch,
+            pct_start=pct_start,
+            **kwargs,
+        )
+
+    elif name_lower == "step":
+        return StepLR(optimizer, step_size=step_size, gamma=gamma, **kwargs)
+
+    elif name_lower == "multistep":
+        if milestones is None:
+            # Default milestones at 30%, 60%, 90% of epochs
+            milestones = [int(epochs * 0.3), int(epochs * 0.6), int(epochs * 0.9)]
+        return MultiStepLR(optimizer, milestones=milestones, gamma=gamma, **kwargs)
+
+    elif name_lower == "exponential":
+        return ExponentialLR(optimizer, gamma=gamma, **kwargs)
+
+    elif name_lower == "linear_warmup":
+        return LinearLR(
+            optimizer,
+            start_factor=start_factor,
+            end_factor=1.0,
+            total_iters=warmup_epochs,
+            **kwargs,
+        )
+
+    else:
+        available = ", ".join(list_schedulers())
+        raise ValueError(f"Unknown scheduler: '{name}'. Available options: {available}")
+
+
+def get_scheduler_with_warmup(
+    name: str,
+    optimizer: optim.Optimizer,
+    warmup_epochs: int = 5,
+    start_factor: float = 0.1,
+    **kwargs,
+) -> LRScheduler:
+    """
+    Create a scheduler with linear warmup phase.
+
+    Combines LinearLR warmup with any other scheduler using SequentialLR.
+
+    Args:
+        name: Main scheduler name (after warmup)
+        optimizer: Optimizer instance
+        warmup_epochs: Number of warmup epochs
+        start_factor: Starting LR factor for warmup
+        **kwargs: Arguments for main scheduler (see get_scheduler)
+
+    Returns:
+        SequentialLR combining warmup and main scheduler
+
+    Example:
+        >>> scheduler = get_scheduler_with_warmup(
+        ...     "cosine", optimizer, warmup_epochs=5, epochs=100
+        ... )
+    """
+    # Create warmup scheduler
+    warmup_scheduler = LinearLR(
+        optimizer,
+        start_factor=start_factor,
+        end_factor=1.0,
+        total_iters=warmup_epochs,
+    )
+
+    # Create main scheduler
+    main_scheduler = get_scheduler(name, optimizer, **kwargs)
+
+    # Combine with SequentialLR
+    return SequentialLR(
+        optimizer,
+        schedulers=[warmup_scheduler, main_scheduler],
+        milestones=[warmup_epochs],
+    )
+
+
+def is_epoch_based(name: str) -> bool:
+    """
+    Check if scheduler should be stepped per epoch (True) or per batch (False).
+
+    Args:
+        name: Scheduler name
+
+    Returns:
+        True if scheduler should step per epoch, False for per batch
+    """
+    name_lower = name.lower().replace("-", "_")
+
+    # OneCycleLR steps per batch, all others step per epoch
+    per_batch_schedulers = {"onecycle"}
+
+    return name_lower not in per_batch_schedulers

wavedl-1.2.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Ductho Le
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.