Migrate optimizers to torch

docs/full/lib/model.rst

@@ -33,6 +33,22 @@ networks package
    :include-all-objects:
    :noindex:
 
+optimizers package
+==================
+
+.. automodapi:: lib.model.optimizers.adabelief
+   :include-all-objects:
+   :noindex:
+
+|
+.. automodapi:: lib.model.optimizers.lion
+   :include-all-objects:
+   :noindex:
+
+|
+.. automodapi:: lib.model.optimizers.keras_legacy
+   :include-all-objects:
+   :noindex:
 
 model package
 =============
@@ -62,7 +78,3 @@ model package
 |
 .. automodapi:: lib.model.normalization
    :include-all-objects:
-
-|
-.. automodapi:: lib.model.optimizers
-   :include-all-objects:

docs/full/lib/training.rst

@@ -20,6 +20,11 @@ The training Package handles libraries to assist with training a model
    :include-all-objects:
    :no-inheritance-diagram:
 
+|
+.. automodapi:: lib.training.optimizer
+   :include-all-objects:
+   :no-inheritance-diagram:
+
 |
 .. automodapi:: lib.training.preview
    :include-all-objects:

lib/align/aligned_mask.py

@@ -466,7 +466,7 @@ def __init__(self,
                  blur_kernel: int = 0,
                  blur_type: T.Literal["gaussian", "normalized"] | None = "gaussian",
                  blur_passes: int = 1) -> None:
-        logger.debug(parse_class_init(locals()))
+        logger.trace(parse_class_init(locals()))  # type:ignore[attr-defined]
         self._area = area
         self._landmark_type = landmark_type
         self._landmarks = landmarks

lib/model/autoclip.py

@@ -1,29 +1,28 @@
-""" Auto clipper for clipping gradients. """
+"""Auto clipper for clipping gradients."""
 from __future__ import annotations
 
 import logging
-import typing as T
+import math
+from collections import deque
 
 import numpy as np
 import torch
+from torch import nn
 
 from lib.logger import parse_class_init
 from lib.utils import get_module_objects
 
-if T.TYPE_CHECKING:
-    from keras import KerasTensor
-
 logger = logging.getLogger(__name__)
 
 
 class AutoClipper():
-    """ AutoClip: Adaptive Gradient Clipping for Source Separation Networks
+    """AutoClip: Adaptive Gradient Clipping for Source Separation Networks
 
     Parameters
     ----------
-    clip_percentile: int
+    clip_percentile
         The percentile to clip the gradients at
-    history_size: int, optional
+    history_size
         The number of iterations of data to use to calculate the norm Default: ``10000``
 
     References
@@ -33,32 +32,32 @@ class AutoClipper():
     """
     def __init__(self, clip_percentile: int, history_size: int = 10000) -> None:
         logger.debug(parse_class_init(locals()))
-
         self._clip_percentile = clip_percentile
-        self._history_size = history_size
-        self._grad_history: list[float] = []
+        self._grad_history: deque[float] = deque(maxlen=history_size)
 
-        logger.debug("Initialized %s", self.__class__.__name__)
-
-    def __call__(self, gradients: list[KerasTensor]) -> list[KerasTensor]:
-        """ Call the AutoClip function.
+    def __call__(self, parameters: list[nn.Parameter], *args) -> None:
+        """Call the AutoClip function.
 
         Parameters
         ----------
-        gradients: list[:class:`keras.KerasTensor`]
-            The list of gradient tensors for the optimizer
-
-        Returns
-        ----------
-        list[:class:`keras.KerasTensor`]
-            The autoclipped gradients
+        parameters
+            The parameters to clip
+        args
+            Unused but for compatibility
         """
-        self._grad_history.append(sum(g.data.norm(2).item() ** 2
-                                      for g in gradients if g is not None) ** (1. / 2))
-        self._grad_history = self._grad_history[-self._history_size:]
-        clip_value = np.percentile(self._grad_history, self._clip_percentile)
-        torch.nn.utils.clip_grad_norm_(gradients, T.cast(float, clip_value))
-        return gradients
+        with torch.no_grad():
+            norms = [p.grad.norm(2).item() for p in parameters if p.grad is not None]
+
+        if not norms:
+            return
+
+        global_norm = sum(n ** 2 for n in norms) ** 0.5
+        if not math.isfinite(global_norm):
+            return
+
+        self._grad_history.append(global_norm)
+        clip_value = float(np.percentile(self._grad_history, self._clip_percentile))
+        nn.utils.clip_grad_norm_(parameters, clip_value)
 
 
 __all__ = get_module_objects(__name__)

lib/model/losses/feature_loss.py

@@ -133,9 +133,9 @@ def _normalize_output(cls, inputs: torch.Tensor, epsilon: float = 1e-10) -> torc
 
         Parameters
         ----------
-        inputs: :class:`keras.KerasTensor`
+        inputs
             An output tensor from the trunk model
-        epsilon: float, optional
+        epsilon
             Epsilon to apply to the normalization operation. Default: `1e-10`
         """
         norm_factor = torch.sqrt(torch.sum(torch.square(inputs), dim=1, keepdim=True))

lib/model/optimizers/__init__.py

@@ -0,0 +1,5 @@
+#! /usr/env/bin/python3
+"""Custom Torch Optimizers"""
+from .adabelief import AdaBelief
+from .lion import Lion
+from .keras_legacy import AdaBelief as AdaBeliefKeras