wavedl 1.5.7__py3-none-any.whl → 1.6.1__py3-none-any.whl

wavedl/utils/optimizers.py

@@ -1,216 +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)
+"""
+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)