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

torchzero/core/transform.py

@@ -1,18 +1,36 @@
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Sequence, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any, Literal, final
 
 import torch
 
-from ..utils import set_storage_, TensorList, vec_to_tensors
-from .module import Module, Var, Chain, Chainable
+from ..utils import TensorList, set_storage_, vec_to_tensors
+from .module import Chain, Chainable, Module, Var
 
 Target = Literal['grad', 'update', 'closure', 'params_direct', 'params_difference', 'update_difference']
 
+
 class Transform(Module, ABC):
-    """Base class for a transform. This is an abstract class, to use it, subclass it and override `update` and `apply` methods.
+    """Base class for a transform.
+    This is an abstract class, to use it, subclass it and override ``update_tensors`` and ``apply_tensors`` methods.
 
     A transform is a module that can also be applied manually to an arbitrary sequence of tensors.
+    It has two methods:
+
+    - ``update_tensors`` updates the internal state of this transform, it doesn't modify tensors. \
+            It may be called multiple times before ``apply_tensors``.
+    - ``apply_tensors`` applies this transform to tensors, without modifying the internal state if possible.
+
+    Alternatively, if update-apply structure doesn't make sense for a transform, all logic can be defined within ``apply_tensors``.
+
+    Transform can be applied to tensors corresponding to custom parameters
+    by calling ``keyed_transform_update`` and ``keyed_transform_apply``,
+    parameters will be keys to store per-parameter states, so they should remain the same python objects.
+
+    Alternatively you can manually create a list of state dictionaries per each tensor and pass it to
+    ``transform_update`` and ``transform_apply``.
+
+    A transform can modify the closure instead of directly modifying update by passing ``target="closure"``.
 
     Args:
         defaults (dict[str,Any] | None): dict with default values.
@@ -21,6 +39,7 @@ class Transform(Module, ABC):
             `grad` is always computed and can't be None. Otherwise set to False.
         target (Target, optional):
             what to set on var. Defaults to 'update'.
+
     """
     def __init__(
         self,
@@ -29,7 +48,6 @@ class Transform(Module, ABC):
         uses_loss: bool = False,
         concat_params: bool = False,
         update_freq: int = 1,
-        scale_first: bool = False,
         inner: Chainable | None = None,
         target: Target = 'update',
     ):
@@ -39,8 +57,8 @@ class Transform(Module, ABC):
         self._uses_loss = uses_loss
         self._concat_params = concat_params
         self._update_freq = update_freq
-        self._scale_first = scale_first
         self._inner = inner
+        self._var = None
 
     def update_tensors(
         self,
@@ -93,14 +111,6 @@ class Transform(Module, ABC):
         states = states[:num]
         settings = settings[:num]
 
-        scale_factor = 1
-
-        # scaling factor for 1st step
-        if self._scale_first and step == 0:
-            # initial step size guess from pytorch LBFGS
-            scale_factor = 1 / TensorList(tensors).abs().global_sum().clip(min=1)
-            scale_factor = scale_factor.clip(min=torch.finfo(tensors[0].dtype).eps)
-
         # update transform
         if step % self._update_freq == 0:
             self.update_tensors(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
@@ -109,7 +119,6 @@ class Transform(Module, ABC):
         self.global_state["__tensors"] = tensors
         self.global_state["__params"] = params
         self.global_state["__grads"] = grads
-        self.global_state["__scale_factor"] = scale_factor
 
 
     @final
@@ -140,23 +149,19 @@ class Transform(Module, ABC):
         tensors = self.global_state.pop("__tensors")
         params  = self.global_state.pop("__params")
         grads   = self.global_state.pop("__grads")
-        scale_factor = self.global_state.pop("__scale_factor")
 
         # step with inner
         if self._inner is not None:
-            tensors = apply_transform(self._inner, tensors=un_tensors, params=un_params, grads=un_grads)
+            tensors = apply_transform(self._inner, tensors=un_tensors, params=un_params, grads=un_grads, var=self._var)
             if self._concat_params:
                 tensors = [torch.cat([t.ravel() for t in tensors])]
 
         # apply transform
         tensors = list(self.apply_tensors(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings))
 
-        # scale initial step, when preconditioner might not have been applied
-        if self._scale_first and self.global_state['__step'] == 1:
-            torch._foreach_mul_(tensors, scale_factor)
-
         if self._concat_params:
             tensors = vec_to_tensors(vec=tensors[0], reference=un_tensors)
+
         return tensors
 
     def _get_keyed_states_settings(self, params: list[torch.Tensor]):
@@ -220,7 +225,9 @@ class Transform(Module, ABC):
         self.pre_step(var)
 
         # update
+        self._var = var
         self.keyed_transform_update(update, params, var.grad, var.loss)
+        self._var = None
 
     def apply(self, var: Var):
         if self._target != 'update':
@@ -234,7 +241,10 @@ class Transform(Module, ABC):
         params=var.params
 
         # apply
+        self._var = var
         var.update = self.keyed_transform_apply(update, params, var.grad, var.loss)
+        self._var = None
+
         self.post_step(var)
         return var
 
@@ -246,12 +256,14 @@ class Transform(Module, ABC):
         if self._uses_loss: var.get_loss(False)
         params=var.params
         self.pre_step(var)
+        self._var = var
 
         # ---------------------------------- update ---------------------------------- #
         if self._target == 'update':
             update = var.get_update()
             self.keyed_transform_update(update, params, var.grad, var.loss)
             var.update = list(self.keyed_transform_apply(update, params, var.grad, var.loss))
+            self._var = None
             return var
 
         # ----------------------------------- grad ----------------------------------- #
@@ -259,6 +271,7 @@ class Transform(Module, ABC):
             grad = var.get_grad()
             self.keyed_transform_update(grad, params, grad, var.loss)
             var.grad = list(self.keyed_transform_apply(grad, params, grad, var.loss))
+            self._var = None
             return var
 
         # ------------------------------- params_direct ------------------------------ #
@@ -266,6 +279,7 @@ class Transform(Module, ABC):
             self.keyed_transform_update(var.params, params, var.grad, var.loss)
             new_params = self.keyed_transform_apply(var.params, params, var.grad, var.loss)
             for p, new_p in zip(var.params, new_params): set_storage_(p, new_p)
+            self._var = None
             return var
 
         # ----------------------------- params_differnce ----------------------------- #
@@ -274,6 +288,7 @@ class Transform(Module, ABC):
             self.keyed_transform_update(p_clone, params, var.grad, var.loss)
             new_params = tuple(self.keyed_transform_apply(p_clone, params, var.grad, var.loss))
             var.update = list(torch._foreach_sub(var.params, new_params))
+            self._var = None
             return var
 
         # ----------------------------- update_difference ---------------------------- #
@@ -283,6 +298,7 @@ class Transform(Module, ABC):
             self.keyed_transform_update(u_clone, params, var.grad, var.loss)
             new_update = tuple(self.keyed_transform_apply(u_clone, params, var.grad, var.loss))
             var.update = list(torch._foreach_sub(update, new_update))
+            self._var = None
             return var
 
         # ---------------------------------- closure --------------------------------- #
@@ -291,12 +307,17 @@ class Transform(Module, ABC):
             if original_closure is None: raise ValueError('Target = "closure", but closure is None')
 
             params = var.params
+            parent_var = self._var
             def transformed_closure(backward=True):
                 if backward:
                     loss = original_closure()
                     current_grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+
+                    self._var = parent_var
                     self.keyed_transform_update(current_grad, params, var.grad, var.loss)
                     transformed_grad = list(self.keyed_transform_apply(current_grad, params, var.grad, var.loss))
+                    self._var = None
+
                     for p, g in zip(params, transformed_grad):
                         p.grad = g
 
@@ -307,6 +328,7 @@ class Transform(Module, ABC):
 
             var.closure = transformed_closure
             self.post_step(var)
+            self._var = None
             return var
 
         # ---------------------------------- invalid --------------------------------- #
@@ -316,7 +338,7 @@ class Transform(Module, ABC):
 class TensorwiseTransform(Transform, ABC):
     """Base class for a parameter-wise transform.
 
-    This is an abstract class, to use it, subclass it and override `transform`.
+    This is an abstract class, to use it, subclass it and override `update_tensor` and `apply_tensor`.
 
     Args:
         defaults (dict[str,Any] | None): dict with default values.
@@ -333,7 +355,6 @@ class TensorwiseTransform(Transform, ABC):
         uses_loss: bool = False,
         concat_params: bool = False,
         update_freq: int = 1,
-        scale_first: bool = False,
         inner: Chainable | None = None,
         target: Target = 'update',
     ):
@@ -342,7 +363,6 @@ class TensorwiseTransform(Transform, ABC):
             uses_grad=uses_grad,
             concat_params=concat_params,
             update_freq=update_freq,
-            scale_first=scale_first,
             uses_loss=uses_loss,
             inner=inner,
             target=target,

torchzero/modules/__init__.py

@@ -1,15 +1,23 @@
+from . import experimental
 from .clipping import *
+from .conjugate_gradient import *
 from .grad_approximation import *
+from .higher_order import *
+from .least_squares import *
 from .line_search import *
-from .step_size import *
+from .misc import *
 from .momentum import *
 from .ops import *
-from .optimizers import *
+from .adaptive import *
 from .projections import *
 from .quasi_newton import *
+from .second_order import *
 from .smoothing import *
+from .step_size import *
+from .termination import *
+from .trust_region import *
+from .variance_reduction import *
 from .weight_decay import *
 from .wrappers import *
-from .second_order import *
-from .higher_order import *
-from .misc import *
+from .restarts import *
+from .zeroth_order import *

torchzero/modules/{optimizers → adaptive}/__init__.py

@@ -1,4 +1,4 @@
-from .adagrad import Adagrad, FullMatrixAdagrad
+from .adagrad import Adagrad, FullMatrixAdagrad, AdagradNorm
 
 # from .curveball import CurveBall
 # from .spectral import SpectralPreconditioner
@@ -6,12 +6,15 @@ from .adahessian import AdaHessian
 from .adam import Adam
 from .adan import Adan
 from .adaptive_heavyball import AdaptiveHeavyBall
+from .aegd import AEGD
 from .esgd import ESGD
-from .ladagrad import LMAdagrad
+from .lmadagrad import LMAdagrad
 from .lion import Lion
 from .mars import MARSCorrection
+from .matrix_momentum import MatrixMomentum
 from .msam import MSAM, MSAMObjective
 from .muon import DualNormCorrection, MuonAdjustLR, Orthogonalize, orthogonalize_grads_
+from .natural_gradient import NaturalGradient
 from .orthograd import OrthoGrad, orthograd_
 from .rmsprop import RMSprop
 from .rprop import (

torchzero/modules/adaptive/adagrad.py

@@ -0,0 +1,356 @@
+from operator import itemgetter
+from typing import Literal
+
+import torch
+from ...core import (
+    Chainable,
+    Module,
+    Target,
+    TensorwiseTransform,
+    Transform,
+    Var,
+    apply_transform,
+)
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+from ...utils.linalg import matrix_power_eigh
+from ..functional import add_power_, lerp_power_, root, epsilon_step_size
+from ...utils.linalg.linear_operator import Dense
+
+def adagrad_(
+    tensors_: TensorList,
+    sq_sum_: TensorList,
+    alpha: float | NumberList,
+    lr_decay: float | NumberList,
+    eps: float | NumberList,
+    step: int,
+    pow: float = 2,
+    use_sqrt: bool = True,
+    divide: bool = False,
+
+    decay: float | None = None,
+    beta: float | None = None,
+
+    # inner args
+    inner: Module | None = None,
+    params: list[torch.Tensor] | None = None,
+    grads: list[torch.Tensor] | None = None,
+):
+    """returns `tensors_`"""
+    clr = alpha / (1 + step * lr_decay)
+
+    if beta is None or step == 1: sq_sum_ = add_power_(tensors_, sum_=sq_sum_, pow=pow)
+    else: sq_sum_ = lerp_power_(tensors_, exp_avg_pow_=sq_sum_, beta=beta, pow=pow)
+    if decay is not None:
+        sq_sum_.mul_(1-decay)
+
+    if inner is not None:
+        assert params is not None
+        tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
+
+    if divide: sq_sum_ = sq_sum_ / max(step, 1)
+
+    if use_sqrt: tensors_.div_(root(sq_sum_, p=pow, inplace=False).add_(eps)).mul_(clr)
+    else: tensors_.div_(sq_sum_.add(eps)).mul_(clr)
+
+    return tensors_
+
+
+
+class Adagrad(Transform):
+    """Adagrad, divides by sum of past squares of gradients.
+
+    This implementation is identical to ``torch.optim.Adagrad``.
+
+    Args:
+        lr_decay (float, optional): learning rate decay. Defaults to 0.
+        initial_accumulator_value (float, optional): initial value of the sum of squares of gradients. Defaults to 0.
+        eps (float, optional): division epsilon. Defaults to 1e-10.
+        alpha (float, optional): step size. Defaults to 1.
+        pow (float, optional): power for gradients and accumulator root. Defaults to 2.
+        use_sqrt (bool, optional): whether to take the root of the accumulator. Defaults to True.
+        inner (Chainable | None, optional): Inner modules that are applied after updating accumulator and before preconditioning. Defaults to None.
+    """
+    def __init__(
+        self,
+        lr_decay: float = 0,
+        initial_accumulator_value: float = 0,
+        eps: float = 1e-10,
+        alpha: float = 1,
+        pow: float = 2,
+        use_sqrt: bool = True,
+        divide: bool=False,
+        beta:float | None = None,
+        decay: float | None = None,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(alpha = alpha, lr_decay = lr_decay, initial_accumulator_value=initial_accumulator_value,
+                        eps = eps, pow=pow, use_sqrt = use_sqrt, divide=divide, beta=beta, decay=decay)
+        super().__init__(defaults=defaults, uses_grad=False)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        lr_decay,alpha,eps = unpack_dicts(settings, 'lr_decay', 'alpha', 'eps', cls=NumberList)
+
+        pow, use_sqrt, divide = itemgetter('pow', 'use_sqrt', 'divide')(settings[0])
+
+        sq_sum = unpack_states(states, tensors, 'sq_sum', cls=TensorList)
+
+        # initialize accumulator on 1st step
+        if step == 1:
+            sq_sum.set_(tensors.full_like([s['initial_accumulator_value'] for s in settings]))
+
+        return adagrad_(
+            tensors,
+            sq_sum_=sq_sum,
+            alpha=alpha,
+            lr_decay=lr_decay,
+            eps=eps,
+            step=step,
+            pow=pow,
+            use_sqrt=use_sqrt,
+            divide=divide,
+
+            beta = self.defaults["beta"],
+            decay = self.defaults["decay"],
+            # inner args
+            inner=self.children.get("inner", None),
+            params=params,
+            grads=grads,
+        )
+
+
+def lerp(start, end, weight):
+    return start + weight * (end - start)
+
+def adagrad_norm_(
+    tensors_: TensorList,
+    accumulator: float | torch.Tensor,
+    alpha: float | NumberList,
+    lr_decay: float | NumberList,
+    eps: float | NumberList,
+    step: int,
+    use_sqrt: bool = True,
+    divide: bool = False,
+
+    decay: float | None = None,
+    beta: float | None = None,
+
+    # inner args
+    inner: Module | None = None,
+    params: list[torch.Tensor] | None = None,
+    grads: list[torch.Tensor] | None = None,
+):
+    """returns `tensors_`"""
+    clr = alpha / (1 + step * lr_decay)
+
+    gg = tensors_.dot(tensors_)
+
+    if beta is None or step == 1: accumulator += gg
+    else: accumulator = lerp(accumulator, gg, 1-beta)
+
+    if decay is not None:
+        accumulator *= 1-decay
+
+    if inner is not None:
+        assert params is not None
+        tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
+
+    if divide: accumulator = accumulator / max(step, 1)
+
+    if use_sqrt: tensors_.div_(eps + accumulator.sqrt()).mul_(clr)
+    else: tensors_.div_(eps + accumulator).mul_(clr)
+
+    return tensors_, accumulator
+
+class AdagradNorm(Transform):
+    """Adagrad-Norm, divides by sum of past means of squares of gradients.
+
+    Args:
+        lr_decay (float, optional): learning rate decay. Defaults to 0.
+        initial_accumulator_value (float, optional): initial value of the sum of squares of gradients. Defaults to 0.
+        eps (float, optional): division epsilon. Defaults to 1e-10.
+        alpha (float, optional): step size. Defaults to 1.
+        pow (float, optional): power for gradients and accumulator root. Defaults to 2.
+        use_sqrt (bool, optional): whether to take the root of the accumulator. Defaults to True.
+        inner (Chainable | None, optional): Inner modules that are applied after updating accumulator and before preconditioning. Defaults to None.
+    """
+    def __init__(
+        self,
+        lr_decay: float = 0,
+        initial_accumulator_value: float = 0,
+        eps: float = 1e-10,
+        alpha: float = 1,
+        pow: float = 2,
+        use_sqrt: bool = True,
+        divide: bool=False,
+        beta:float | None = None,
+        decay: float | None = None,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(alpha = alpha, lr_decay = lr_decay, initial_accumulator_value=initial_accumulator_value,
+                        eps = eps, pow=pow, use_sqrt = use_sqrt, divide=divide, beta=beta, decay=decay)
+        super().__init__(defaults=defaults, uses_grad=False)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+        lr_decay,alpha,eps = unpack_dicts(settings, 'lr_decay', 'alpha', 'eps', cls=NumberList)
+
+        use_sqrt, divide, initial_accumulator_value = itemgetter('use_sqrt', 'divide', "initial_accumulator_value")(settings[0])
+
+        accumulator = self.global_state.get("accumulator", initial_accumulator_value)
+
+        d, self.global_state["accumulator"] = adagrad_norm_(
+            tensors,
+            accumulator=accumulator,
+            alpha=alpha,
+            lr_decay=lr_decay,
+            eps=eps,
+            step=step,
+            use_sqrt=use_sqrt,
+            divide=divide,
+
+            beta = self.defaults["beta"],
+            decay = self.defaults["decay"],
+            # inner args
+            inner=self.children.get("inner", None),
+            params=params,
+            grads=grads,
+        )
+
+        return d
+
+
+class FullMatrixAdagrad(TensorwiseTransform):
+    """Full-matrix version of Adagrad, can be customized to make RMSprop or Adam (see examples).
+
+    Note:
+        A more memory-efficient version equivalent to full matrix Adagrad on last n gradients is implemented in ``tz.m.LMAdagrad``.
+
+    Args:
+        beta (float | None, optional): momentum for gradient outer product accumulators. if None, uses sum. Defaults to None.
+        decay (float | None, optional): decay for gradient outer product accumulators. Defaults to None.
+        sqrt (bool, optional): whether to take the square root of the accumulator. Defaults to True.
+        concat_params (bool, optional): if False, each parameter will have it's own accumulator. Defaults to True.
+        precond_freq (int, optional): frequency of updating the inverse square root of the accumulator. Defaults to 1.
+        init (Literal[str], optional):
+            how to initialize the accumulator.
+            - "identity" - with identity matrix (default).
+            - "zeros" - with zero matrix.
+            - "ones" - with matrix of ones.
+             -"GGT" - with the first outer product
+        divide (bool, optional): whether to divide the accumulator by number of gradients in it. Defaults to False.
+        inner (Chainable | None, optional): inner modules to apply preconditioning to. Defaults to None.
+
+    ## Examples:
+
+    Plain full-matrix adagrad
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.FullMatrixAdagrd(),
+        tz.m.LR(1e-2),
+    )
+    ```
+
+    Full-matrix RMSprop
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.FullMatrixAdagrad(beta=0.99),
+        tz.m.LR(1e-2),
+    )
+    ```
+
+    Full-matrix Adam
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.FullMatrixAdagrad(beta=0.999, inner=tz.m.EMA(0.9)),
+        tz.m.Debias(0.9, 0.999),
+        tz.m.LR(1e-2),
+    )
+    ```
+    """
+    def __init__(
+        self,
+        beta: float | None = None,
+        decay: float | None = None,
+        sqrt: bool = True,
+        concat_params=True,
+        precond_freq: int = 1,
+        init: Literal["identity", "zeros", "ones", "GGT"] = "identity",
+        reg: float = 1e-12,
+        divide: bool = False,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(beta=beta, decay=decay, sqrt=sqrt, precond_freq=precond_freq, init=init, divide=divide, reg=reg)
+        super().__init__(defaults, uses_grad=False, concat_params=concat_params, inner=inner,)
+
+    @torch.no_grad
+    def update_tensor(self, tensor, param, grad, loss, state, setting):
+        G = tensor.ravel()
+        GG = torch.outer(G, G)
+        decay = setting['decay']
+        beta = setting['beta']
+        init = setting['init']
+
+        if 'GG' not in state:
+            if init == 'identity': state['GG'] = torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype)
+            elif init == 'zeros': state['GG'] =  torch.zeros_like(GG)
+            elif init == 'ones': state['GG'] = torch.ones_like(GG)
+            elif init == 'GGT': state['GG'] = GG.clone()
+            else: raise ValueError(init)
+        if decay is not None: state['GG'].mul_(decay)
+
+        if beta is not None: state['GG'].lerp_(GG, 1-beta)
+        else: state['GG'].add_(GG)
+        state['i'] = state.get('i', 0) + 1 # number of GGTs in sum
+
+    @torch.no_grad
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        step = state.get('step', 0)
+        state['step'] = step + 1
+
+        GG: torch.Tensor = state['GG']
+        sqrt = setting['sqrt']
+        divide = setting['divide']
+        precond_freq = setting['precond_freq']
+        reg = setting['reg']
+
+        if divide: GG = GG/state.get('i', 1)
+
+        if reg != 0:
+            GG = GG + torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype).mul_(reg)
+
+        if tensor.numel() == 1:
+            GG = GG.squeeze()
+            if sqrt: return tensor / GG.sqrt()
+            return tensor / GG
+
+        try:
+            if sqrt:
+                if "B" not in state or step % precond_freq == 0:
+                    B = state["B"] = matrix_power_eigh(GG, -1/2)
+                else:
+                    B = state["B"]
+
+            else: return torch.linalg.solve(GG, tensor.ravel()).view_as(tensor) # pylint:disable = not-callable
+
+        except torch.linalg.LinAlgError:
+            # fallback to diagonal AdaGrad
+            denom = GG.diagonal()
+            if sqrt: denom = denom.sqrt()
+            return tensor.div_(denom + max(reg, 1e-12))
+
+        return (B @ tensor.ravel()).view_as(tensor)