torchzero 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/modules/clipping/clipping.py

@@ -5,7 +5,7 @@ import math
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList, generic_eq
+from ...utils import NumberList, TensorList
 
 
 def clip_grad_value_(params: Iterable[torch.Tensor], value: float):
@@ -24,7 +24,7 @@ def _clip_norm_(
     min: float | NumberList | None,
     max: float | NumberList | None,
     norm_value: float | NumberList | None,
-    ord: float,
+    ord: float | Literal['mean_abs'],
     dim: int | Sequence[int] | Literal["global"] | None,
     inverse_dims: bool,
     min_size: int,
@@ -54,9 +54,13 @@ def _clip_norm_(
         size = math.prod(tensor.size(d) for d in real_dim)
         if size < min_size: continue
 
-        norm: torch.Tensor = torch.linalg.vector_norm(tensor, ord=ord, dim=real_dim, keepdim=True) # pylint:disable=not-callable
+        if ord == 'mean_abs':
+            norm = tensor.abs().mean(dim=real_dim, keepdim=True)
+        else:
+            norm: torch.Tensor = torch.linalg.vector_norm(tensor, ord=ord, dim=real_dim, keepdim=True) # pylint:disable=not-callable
+
         if norm.numel() == 1 and norm == 0: continue
-        norm = torch.where(norm == 0, 1, norm)
+        norm = torch.where(norm <= 1e-12, 1, norm)
 
         # normalize = True, perform normalization
         norm_v = norm_value[i] if isinstance(norm_value, (list,tuple)) else norm_value
@@ -90,7 +94,7 @@ def _clip_norm_(
 def clip_grad_norm_(
     params: Iterable[torch.Tensor],
     max_norm: float | None,
-    ord: float = 2,
+    ord: float | Literal['mean_abs'] = 2,
     dim: int | Sequence[int] | Literal["global"] | None = None,
     inverse_dims: bool = False,
     min_size: int = 2,
@@ -118,7 +122,7 @@ def clip_grad_norm_(
 def normalize_grads_(
     params: Iterable[torch.Tensor],
     norm_value: float,
-    ord: float = 2,
+    ord: float | Literal['mean_abs'] = 2,
     dim: int | Sequence[int] | Literal["global"] | None = None,
     inverse_dims: bool = False,
     min_size: int = 1,
@@ -145,14 +149,44 @@ def normalize_grads_(
 
 
 class ClipValue(Transform):
-    """Clips update magnitude to be within `(-value, value)` range."""
+    """Clips update magnitude to be within `(-value, value)` range.
+
+    Args:
+        value (float): value to clip to.
+        target (str): refer to :ref:`target argument` in documentation.
+
+    Examples:
+
+        Gradient clipping:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ClipValue(1),
+                tz.m.Adam(),
+                tz.m.LR(1e-2),
+            )
+
+        Update clipping:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.ClipValue(1),
+                tz.m.LR(1e-2),
+            )
+
+    """
     def __init__(self, value: float, target: Target = 'update'):
         defaults = dict(value=value)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        value = self.get_settings('value', params=params)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        value = [s['value'] for s in settings]
         return TensorList(tensors).clip_([-v for v in value], value)
 
 class ClipNorm(Transform):
@@ -172,23 +206,47 @@ class ClipNorm(Transform):
             minimal numer of elements in a parameter or slice to clip norm. Defaults to 1.
         target (str, optional):
             what this affects.
+
+    Examples:
+
+        Gradient norm clipping:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ClipNorm(1),
+                tz.m.Adam(),
+                tz.m.LR(1e-2),
+            )
+
+        Update norm clipping:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.ClipNorm(1),
+                tz.m.LR(1e-2),
+            )
     """
     def __init__(
         self,
         max_norm: float,
-        ord: float = 2,
+        ord: float | Literal['mean_abs'] = 2,
         dim: int | Sequence[int] | Literal["global"] | None = None,
         inverse_dims: bool = False,
         min_size: int = 1,
         target: Target = "update",
     ):
         defaults = dict(max_norm=max_norm,ord=ord,dim=dim,min_size=min_size,inverse_dims=inverse_dims)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        max_norm = self.get_settings('max_norm', params=params, cls=NumberList)
-        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        max_norm = NumberList(s['max_norm'] for s in settings)
+        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(settings[0])
         _clip_norm_(
             tensors_ = TensorList(tensors),
             min = 0,
@@ -218,23 +276,47 @@ class Normalize(Transform):
             minimal size of a dimension to normalize along it. Defaults to 1.
         target (str, optional):
             what this affects.
+
+    Examples:
+
+        Gradient normalization:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Normalize(1),
+                tz.m.Adam(),
+                tz.m.LR(1e-2),
+            )
+
+        Update normalization:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.Normalize(1),
+                tz.m.LR(1e-2),
+            )
     """
     def __init__(
         self,
         norm_value: float = 1,
-        ord: float = 2,
+        ord: float | Literal['mean_abs'] = 2,
         dim: int | Sequence[int] | Literal["global"] | None = None,
         inverse_dims: bool = False,
         min_size: int = 1,
         target: Target = "update",
     ):
         defaults = dict(norm_value=norm_value,ord=ord,dim=dim,min_size=min_size, inverse_dims=inverse_dims)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        norm_value = self.get_settings('norm_value', params=params, cls=NumberList)
-        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        norm_value = NumberList(s['norm_value'] for s in settings)
+        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(settings[0])
 
         _clip_norm_(
             tensors_ = TensorList(tensors),
@@ -299,6 +381,21 @@ class Centralize(Transform):
             if True, the `dims` argument is inverted, and all other dimensions are centralized.
         min_size (int, optional):
             minimal size of a dimension to normalize along it. Defaults to 1.
+
+    Examples:
+
+        Standard gradient centralization:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Centralize(dim=0),
+                tz.m.LR(1e-2),
+            )
+
+    References:
+        - Yong, H., Huang, J., Hua, X., & Zhang, L. (2020). Gradient centralization: A new optimization technique for deep neural networks. In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part I 16 (pp. 635-652). Springer International Publishing. https://arxiv.org/abs/2004.01461
     """
     def __init__(
         self,
@@ -308,11 +405,11 @@ class Centralize(Transform):
         target: Target = "update",
     ):
         defaults = dict(dim=dim,min_size=min_size,inverse_dims=inverse_dims)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        dim, min_size, inverse_dims = itemgetter('dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        dim, min_size, inverse_dims = itemgetter('dim', 'min_size', 'inverse_dims')(settings[0])
 
         _centralize_(tensors_ = TensorList(tensors), dim=dim, inverse_dims=inverse_dims, min_size=min_size)
 

torchzero/modules/clipping/ema_clipping.py

@@ -4,8 +4,8 @@ from collections.abc import Iterable, Sequence
 
 import torch
 
-from ...core import Module, Target, Transform, apply, Chainable
-from ...utils import NumberList, TensorList, generic_eq
+from ...core import Module, Target, Transform, apply_transform, Chainable
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 
 class ClipNormByEMA(Transform):
     """Clips norm to be no larger than the norm of an exponential moving average of past updates.
@@ -14,9 +14,10 @@ class ClipNormByEMA(Transform):
         beta (float, optional): beta for the exponential moving average. Defaults to 0.99.
         ord (float, optional): order of the norm. Defaults to 2.
         eps (float, optional): epsilon for division. Defaults to 1e-6.
-        tensorwise (bool, optional): whether to calculate norm separately for each layer, or global norm for all layers. Defaults to True.
+        tensorwise (bool, optional):
+            if True, norms are calculated parameter-wise, otherwise treats all parameters as single vector. Defaults to True.
         max_ema_growth (float | None, optional):
-            if specified, exponential moving average norm can grow but at most this value per step. Defaults to 1.5.
+            if specified, restricts how quickly exponential moving average norm can grow. The norm is allowed to grow by at most this value per step. Defaults to 1.5.
         ema_init (str, optional):
             How to initialize exponential moving average on first step, "update" to use the first update or "zeros". Defaults to 'zeros'.
     """
@@ -29,18 +30,20 @@ class ClipNormByEMA(Transform):
         tensorwise:bool=True,
         max_ema_growth: float | None = 1.5,
         ema_init: Literal['zeros', 'update'] = 'zeros',
+        inner: Chainable | None = None,
     ):
         defaults = dict(beta=beta, ord=ord, tensorwise=tensorwise, ema_init=ema_init, eps=eps, max_ema_growth=max_ema_growth)
-        super().__init__(defaults, uses_grad=False)
+        super().__init__(defaults, inner=inner)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        ord, tensorwise, ema_init, max_ema_growth = itemgetter('ord', 'tensorwise', 'ema_init', 'max_ema_growth')(self.settings[params[0]])
-
-        beta, eps = self.get_settings('beta', 'eps', params=params, cls=NumberList)
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
         tensors = TensorList(tensors)
+        ord, tensorwise, ema_init, max_ema_growth = itemgetter('ord', 'tensorwise', 'ema_init', 'max_ema_growth')(settings[0])
+
+        beta, eps = unpack_dicts(settings, 'beta', 'eps', cls=NumberList)
+
+        ema = unpack_states(states, tensors, 'ema', init = (torch.zeros_like if ema_init=='zeros' else tensors), cls=TensorList)
 
-        ema = self.get_state('ema', params=params, init = (torch.zeros_like if ema_init=='zeros' else tensors), cls=TensorList)
         ema.lerp_(tensors, 1-beta)
 
         if tensorwise:
@@ -48,7 +51,7 @@ class ClipNormByEMA(Transform):
 
             # clip ema norm growth
             if max_ema_growth is not None:
-                prev_ema_norm = self.get_state('prev_ema_norm', params=params, init=ema_norm, cls=TensorList)
+                prev_ema_norm = unpack_states(states, tensors, 'prev_ema_norm', init=ema_norm, cls=TensorList)
                 allowed_norm = (prev_ema_norm * max_ema_growth).clip(min=1e-6)
                 ema_denom = (ema_norm / allowed_norm).clip(min=1)
                 ema.div_(ema_denom)
@@ -77,7 +80,12 @@ class ClipNormByEMA(Transform):
             if self.NORMALIZE: denom.clip_(min=eps[0])
             else: denom.clip_(min=1)
 
-        tensors.div_(denom)
+        self.global_state['denom'] = denom
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        denom = self.global_state.pop('denom')
+        torch._foreach_div_(tensors, denom)
         return tensors
 
 class NormalizeByEMA(ClipNormByEMA):
@@ -87,9 +95,10 @@ class NormalizeByEMA(ClipNormByEMA):
         beta (float, optional): beta for the exponential moving average. Defaults to 0.99.
         ord (float, optional): order of the norm. Defaults to 2.
         eps (float, optional): epsilon for division. Defaults to 1e-6.
-        tensorwise (bool, optional): whether to calculate norm separately for each layer, or global norm for all layers. Defaults to True.
+        tensorwise (bool, optional):
+            if True, norms are calculated parameter-wise, otherwise treats all parameters as single vector. Defaults to True.
         max_ema_growth (float | None, optional):
-            if specified, exponential moving average norm can grow but at most this value per step. Defaults to 1.5.
+            if specified, restricts how quickly exponential moving average norm can grow. The norm is allowed to grow by at most this value per step. Defaults to 1.5.
         ema_init (str, optional):
             How to initialize exponential moving average on first step, "update" to use the first update or "zeros". Defaults to 'zeros'.
     """
@@ -98,38 +107,44 @@ class NormalizeByEMA(ClipNormByEMA):
 # TODO Centralize by EMA?
 
 class ClipValueByEMA(Transform):
-    """Clips magnitude of update to be no larger than magnitude of an exponential moving average of past (unclipped) updates.
+    """Clips magnitude of update to be no larger than magnitude of exponential moving average of past (unclipped) updates.
 
     Args:
         beta (float, optional): beta for the exponential moving average. Defaults to 0.99.
         ema_init (str, optional):
             How to initialize exponential moving average on first step, "update" to use the first update or "zeros". Defaults to 'zeros'.
-        ema_tfm (Chainable | None, optional): optional modules applied to exponential moving average before clipping by it. Defaults to None.
+        ema_tfm (Chainable | None, optional):
+            optional modules applied to exponential moving average before clipping by it. Defaults to None.
     """
     def __init__(
         self,
         beta=0.99,
         ema_init: Literal['zeros', 'update'] = 'zeros',
         ema_tfm:Chainable | None=None,
+        inner: Chainable | None = None,
     ):
         defaults = dict(beta=beta, ema_init=ema_init)
-        super().__init__(defaults, uses_grad=False)
+        super().__init__(defaults, inner=inner)
 
         if ema_tfm is not None:
             self.set_child('ema_tfm', ema_tfm)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        ema_init = itemgetter('ema_init')(self.settings[params[0]])
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        ema_init = itemgetter('ema_init')(settings[0])
 
-        beta = self.get_settings('beta', params=params, cls=NumberList)
+        beta = unpack_dicts(settings, 'beta', cls=NumberList)
         tensors = TensorList(tensors)
 
-        ema = self.get_state('ema', params=params, init = (torch.zeros_like if ema_init=='zeros' else lambda t: t.abs()), cls=TensorList)
+        ema = unpack_states(states, tensors, 'ema', init = (torch.zeros_like if ema_init=='zeros' else lambda t: t.abs()), cls=TensorList)
         ema.lerp_(tensors.abs(), 1-beta)
 
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        ema = unpack_states(states, tensors, 'ema', cls=TensorList)
+
         if 'ema_tfm' in self.children:
-            ema = TensorList(apply(self.children['ema_tfm'], ema, params, vars.grad, vars))
+            ema = TensorList(apply_transform(self.children['ema_tfm'], ema.clone(), params, grads, loss))
 
         tensors.clip_(-ema, ema)
         return tensors

torchzero/modules/clipping/growth_clipping.py

@@ -19,7 +19,7 @@ class ClipValueGrowth(TensorwiseTransform):
             bounds the tracked multiplicative clipping decay to prevent collapse to 0.
             Next update is at most :code:`max(previous update * mul, max_decay)`.
             Defaults to 2.
-        target (Target, optional): what to set on vars.. Defaults to "update".
+        target (Target, optional): what to set on var. Defaults to "update".
     """
     def __init__(
         self,
@@ -30,15 +30,13 @@ class ClipValueGrowth(TensorwiseTransform):
         target: Target = "update",
     ):
         defaults = dict(add=add, mul=mul, min_value=min_value, max_decay=max_decay)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
 
-    def transform(self, tensor, param, grad, vars):
-        add, mul, min_value, max_decay = itemgetter('add','mul','min_value','max_decay')(self.settings[param])
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        add, mul, min_value, max_decay = itemgetter('add','mul','min_value','max_decay')(setting)
         add: float | None
 
-        state = self.state[param]
-
         if add is None and mul is None:
             return tensor
 
@@ -122,7 +120,8 @@ class ClipNormGrowth(Transform):
 
     Args:
         add (float | None, optional): additive clipping, next update norm is at most `previous norm + add`. Defaults to None.
-        mul (float | None, optional): multiplicative clipping, next update norm is at most `previous norm * mul`. Defaults to 1.5.
+        mul (float | None, optional):
+            multiplicative clipping, next update norm is at most `previous norm * mul`. Defaults to 1.5.
         min_value (float | None, optional):
             minimum value for multiplicative clipping to prevent collapse to 0.
             Next norm is at most :code:`max(prev_norm, min_value) * mul`. Defaults to 1e-4.
@@ -133,7 +132,7 @@ class ClipNormGrowth(Transform):
         ord (float, optional): norm order. Defaults to 2.
         parameterwise (bool, optional):
             if True, norms are calculated parameter-wise, otherwise treats all parameters as single vector. Defaults to True.
-        target (Target, optional): what to set on vars. Defaults to "update".
+        target (Target, optional): what to set on var. Defaults to "update".
     """
     def __init__(
         self,
@@ -146,39 +145,39 @@ class ClipNormGrowth(Transform):
         target: Target = "update",
     ):
         defaults = dict(add=add, mul=mul, min_value=min_value, max_decay=max_decay, ord=ord, parameterwise=parameterwise)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
 
 
-    def transform(self, tensors, params, grads, vars):
-        parameterwise = self.settings[params[0]]['parameterwise']
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        parameterwise = settings[0]['parameterwise']
         tensors = TensorList(tensors)
 
         if parameterwise:
             ts = tensors
-            stts = [self.state[p] for p in params]
-            stns = [self.settings[p] for p in params]
+            stts = states
+            stns = settings
 
         else:
             ts = [tensors.to_vec()]
             stts = [self.global_state]
-            stns = [self.settings[params[0]]]
+            stns = [settings[0]]
 
 
-        for t,state, settings in zip(ts, stts, stns):
+        for t, state, setting in zip(ts, stts, stns):
             if 'prev_norm' not in state:
-                state['prev_norm'] = torch.linalg.vector_norm(t, ord=settings['ord']) # pylint:disable=not-callable
+                state['prev_norm'] = torch.linalg.vector_norm(t, ord=setting['ord']) # pylint:disable=not-callable
                 state['prev_denom'] = 1
                 continue
 
             _,  state['prev_norm'], state['prev_denom'] = norm_growth_clip_(
                 tensor_ = t,
                 prev_norm = state['prev_norm'],
-                add = settings['add'],
-                mul = settings['mul'],
-                min_value = settings['min_value'],
-                max_decay = settings['max_decay'],
-                ord = settings['ord'],
+                add = setting['add'],
+                mul = setting['mul'],
+                min_value = setting['min_value'],
+                max_decay = setting['max_decay'],
+                ord = setting['ord'],
             )
 
         if not parameterwise:

torchzero/modules/experimental/__init__.py

@@ -1,15 +1,41 @@
+"""This submodule contains various untested experimental modules, some of them are to be moved out of experimental when properly tested, some are to remain here forever or to be deleted depending on the degree of their usefulness."""
 from .absoap import ABSOAP
 from .adadam import Adadam
+from .adam_lambertw import AdamLambertW
 from .adamY import AdamY
+from .adaptive_step_size import AdaptiveStepSize
 from .adasoap import AdaSOAP
+from .cosine import (
+    AdaptiveDifference,
+    AdaptiveDifferenceEMA,
+    CosineDebounce,
+    CosineMomentum,
+    CosineStepSize,
+    ScaledAdaptiveDifference,
+)
+from .cubic_adam import CubicAdam
 from .curveball import CurveBall
-from .soapy import SOAPY
+
+# from dct import DCTProjection
+from .eigendescent import EigenDescent
+from .etf import (
+    ExponentialTrajectoryFit,
+    ExponentialTrajectoryFitV2,
+    PointwiseExponential,
+)
+from .exp_adam import ExpAdam
+from .expanded_lbfgs import ExpandedLBFGS
+from .fft import FFTProjection
 from .gradmin import GradMin
+from .hnewton import HNewton
+from .modular_lbfgs import ModularLBFGS
+from .newton_solver import NewtonSolver
+from .newtonnewton import NewtonNewton
+from .parabolic_search import CubicParabolaSearch, ParabolaSearch
 from .reduce_outward_lr import ReduceOutwardLR
-from .spectral import SpectralPreconditioner
+from .structural_projections import BlockPartition, TensorizeProjection
 from .subspace_preconditioners import (
     HistorySubspacePreconditioning,
     RandomSubspacePreconditioning,
 )
-from .tropical_newton import TropicalNewton
-from .newton_solver import NewtonSolver
+from .tensor_adagrad import TensorAdagrad