torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/modules/optimizers/adagrad.py

@@ -1,49 +1,146 @@
-from collections import abc
-
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizerModule
-
-def _adagrad_step_(ascent: TensorList, grad_sum: TensorList, alpha: TensorList, lr_decay: TensorList, eps: TensorList, step: int):
-    clr = alpha / (1 + step * lr_decay)
-    grad_sum.addcmul_(ascent, ascent)
-    return ascent.div_(grad_sum.sqrt().add_(eps)).mul_(clr)
-
-class Adagrad(OptimizerModule):
-    """
-    Divides ascent direction by mean square root of the sum of all past ascent directions.
-
-    Exactly matches `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): term added to the denominator to improve numerical stability. Defaults to 1e-10.
-        alpha (float, optional): learning rate. Defaults to 1.
-
-    reference
-        https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf
-    """
-    def __init__(self, lr_decay: float = 0, initial_accumulator_value: float = 0, eps: float = 1e-10, alpha: float = 1):
-        defaults = dict(alpha = alpha, lr_decay = lr_decay, initial_accumulator_value=initial_accumulator_value, eps = eps)
-        super().__init__(defaults)
-        self.cur_step = 0
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        settings = self.get_all_group_keys()
-        if self.cur_step == 0: init = ascent.full_like(settings['initial_accumulator_value'])
-        else: init = None
-        grad_sum = self.get_state_key('grad_sum', init = init) # type:ignore
-
-        updated_direction = _adagrad_step_(
-            ascent=ascent,
-            grad_sum=grad_sum,
-            alpha=settings["alpha"],
-            eps=settings["eps"],
-            lr_decay=settings["lr_decay"],
-            step=self.cur_step,
-        )
-        self.cur_step += 1
-        return updated_direction
+from operator import itemgetter
+
+import torch
+
+from ...core import (
+    Chainable,
+    Module,
+    Preconditioner,
+    Target,
+    TensorwisePreconditioner,
+    Transform,
+    Vars,
+    apply,
+)
+from ...utils import NumberList, TensorList
+from ...utils.linalg import matrix_power_eigh
+from ..functional import add_power_, lerp_power_, root
+
+
+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,
+
+    # inner args
+    inner: Module | None = None,
+    params: list[torch.Tensor] | None = None,
+    grads: list[torch.Tensor] | None = None,
+    vars: Vars | None = None,
+):
+    """returns `tensors_`"""
+    clr = alpha / (1 + step * lr_decay)
+
+    sq_sum_ = add_power_(tensors_, sum_=sq_sum_, pow=pow)
+
+    if inner is not None:
+        assert params is not None
+        tensors_ = TensorList(apply(inner, tensors_, params=params, grads=grads, vars=vars))
+
+    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, matches pytorch 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,
+        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)
+        super().__init__(defaults=defaults, uses_grad=False)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        tensors = TensorList(tensors)
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        lr_decay,alpha,eps = self.get_settings('lr_decay', 'alpha', 'eps', params=params, cls=NumberList)
+
+        pow, use_sqrt = itemgetter('pow', 'use_sqrt')(self.settings[params[0]])
+
+        sq_sum = self.get_state('sq_sum', params=params, cls=TensorList)
+
+        # initialize accumulator on 1st step
+        if step == 1:
+            sq_sum.set_(tensors.full_like(self.get_settings('initial_accumulator_value', params=params)))
+
+        return adagrad_(
+            tensors,
+            sq_sum_=sq_sum,
+            alpha=alpha,
+            lr_decay=lr_decay,
+            eps=eps,
+            step=self.global_state["step"],
+            pow=pow,
+            use_sqrt=use_sqrt,
+
+            # inner args
+            inner=self.children.get("inner", None),
+            params=params,
+            grads=grads,
+            vars=vars,
+        )
+
+
+
+class FullMatrixAdagrad(TensorwisePreconditioner):
+    def __init__(self, beta: float | None = None, decay: float | None = None, concat_params=False, update_freq=1, inner: Chainable | None = None):
+        defaults = dict(beta=beta, decay=decay)
+        super().__init__(defaults, uses_grad=False, concat_params=concat_params, update_freq=update_freq, inner=inner)
+
+    @torch.no_grad
+    def update_tensor(self, tensor, param, grad, state, settings):
+        G = tensor.ravel()
+        GG = torch.outer(G, G)
+        decay = settings['decay']
+        beta = settings['beta']
+
+        if 'GG' not in state: state['GG'] = torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype)
+        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)
+
+    @torch.no_grad
+    def apply_tensor(self, tensor, param, grad, state, settings):
+        GG = state['GG']
+
+        if tensor.numel() == 1:
+            return tensor / (GG**(1/2)).squeeze()
+
+        try:
+            B = matrix_power_eigh(GG, -1/2)
+        except torch.linalg.LinAlgError:
+            return tensor.div_(tensor.abs().max()) # conservative scaling
+
+        return (B @ tensor.ravel()).view_as(tensor)
+

torchzero/modules/optimizers/adam.py

@@ -1,118 +1,112 @@
-from collections import abc
-
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizerModule
-
-def _adam_step(ascent: TensorList, exp_avg: TensorList, exp_avg_sq: TensorList, alpha, beta1, beta2, eps, step:int, max_exp_avg_sqs: TensorList | None, params: TensorList | None = None):
-    # Decay the first and second moment running average coefficient
-    exp_avg.lerp_compat_(ascent, 1 - beta1)
-    exp_avg_sq.mul_(beta2).addcmul_(ascent, ascent.conj(), value=1 - beta2)
-
-    bias_correction1 = 1 - beta1**step
-    bias_correction2 = 1 - beta2**step
-
-    if max_exp_avg_sqs is not None:
-        max_exp_avg_sqs.maximum_(exp_avg_sq)
-        denom = max_exp_avg_sqs.sqrt().div_(bias_correction2**0.5).add_(eps)
-    else:
-        denom = exp_avg_sq.sqrt().div_(bias_correction2**0.5).add_(eps)
-
-    if params is None:
-        return (exp_avg / denom).mul_(alpha / bias_correction1)
-
-    # else directly apply the update to params
-    params.addcdiv_(exp_avg, denom, value = -(alpha / bias_correction1))
-    return None
-
-
-
-class Adam(OptimizerModule):
-    """Adam. Combines momentum and RMSProp. Exactly matches `torch.optim.Adam`.
-
-    Args:
-        beta1 (float, optional): exponential decay rate of gradient moving average. Defaults to 0.9.
-        beta2 (float, optional): exponential decay rate of squared gradient moving average. Defaults to 0.999.
-        eps (float, optional): epsilon for numerical stability. Defaults to 1e-8.
-        amsgrad (bool, optional):
-            whether to use the AMSGrad variant of this algorithm from
-            On the Convergence of Adam and Beyond (default: False).
-        alpha (float, optional): learning rate. Defaults to 1.
-    """
-    def __init__(self, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8, alpha: float = 1, amsgrad=False):
-        defaults = dict(alpha = alpha, beta1=beta1, beta2=beta2, eps=eps)
-        super().__init__(defaults)
-
-        self.cur_step = 1
-        self.amsgrad = amsgrad
-
-    @torch.no_grad
-    def step(self, vars):
-        # Adam step is a bit differet from other optimizer steps
-        # due to how common it is, I implemented two additional optimizations,
-
-        # 1st - if next module is None or if next module is LR and module after is None
-        # this will directly update parameters using `addcdiv_`
-
-        # 2nd - if next module is LR`, adam will "fuse" with it to avoid an additional add operation.
-
-        # the optimizations are quite verbose and seem to barely have any effect, so I probably won't implement
-        # this for other modules
-
-        settings = self.get_all_group_keys()
-
-        if self.amsgrad:
-            exp_avg, exp_avg_sq, max_exp_avg_sqs = self.get_state_keys('exp_avg', 'exp_avg_sq', 'max_exp_avg_sqs')
-        else:
-            exp_avg, exp_avg_sq = self.get_state_keys('exp_avg', 'exp_avg_sq')
-            max_exp_avg_sqs = None
-
-        params = None
-
-        # apply addcdiv if next module is None
-        if self.next_module is None: params = self.get_params()
-
-        # fuse with LR module if it is next
-        if self.next_module is not None and self.next_module.IS_LR_MODULE:
-            alpha = self.next_module.get_group_key('lr') * settings['alpha']
-            self.next_module._skip = True # type:ignore
-
-            # apply addcdiv if module after LR is None.
-            if self.next_module.next_module is None: params = self.get_params()
-
-        else:
-            alpha = settings['alpha']
-
-        # get params if ascent is None so we need params to access their gradient as initial ascent
-        if vars.ascent is None:
-            if params is None: pg = self.get_params()
-            else: pg = params
-        else:
-            pg = None
-
-        ret = _adam_step(
-            ascent=vars.maybe_use_grad_(pg),
-            exp_avg = exp_avg,
-            exp_avg_sq = exp_avg_sq,
-            alpha = alpha,
-            beta1 = settings['beta1'],
-            beta2 = settings['beta2'],
-            eps = settings['eps'],
-            step = self.cur_step,
-            max_exp_avg_sqs = max_exp_avg_sqs,
-            params = params
-        )
-
-        self.cur_step += 1
-        if params is None:
-            assert ret is not None
-            vars.ascent = ret
-            return self._update_params_or_step_with_next(vars)
-
-        # next module is either None or LR
-        if self.next_module is None: return vars.get_loss()
-
-        # step with LR, which has _skip = True so it won't apply lr, but may step with the scheduler
-        self.next_module._update(vars, None) # type:ignore
-        return vars.get_loss()
+from operator import itemgetter
+from functools import partial
+
+import torch
+
+from ...core import Module, Target, Transform
+from ...utils import NumberList, TensorList
+from ..functional import (
+    debias, debiased_step_size,
+    ema_,
+    sqrt_ema_sq_,
+)
+from ..lr.lr import lazy_lr
+from ..momentum.experimental import sqrt_nag_ema_sq_
+from ..momentum.momentum import nag_
+
+
+def adam_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    exp_avg_sq_: TensorList,
+    alpha: float | NumberList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    eps: float | NumberList,
+    step: int,
+    pow: float = 2,
+    debiased: bool = True,
+    max_exp_avg_sq_: TensorList | None = None,
+    params_: TensorList | None = None,
+):
+    """Returns new tensors or updates params in-place."""
+    exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
+
+    sqrt_exp_avg_sq = sqrt_ema_sq_(tensors, exp_avg_sq_=exp_avg_sq_, beta=beta2, max_exp_avg_sq_=max_exp_avg_sq_,
+                                   debiased=False,step=step,pow=pow)
+
+    if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
+
+    # params is None, return update
+    if params_ is None: return (exp_avg_ / sqrt_exp_avg_sq.add_(eps)).lazy_mul(alpha)
+
+    # update params in-place
+    params_.addcdiv_(exp_avg_, sqrt_exp_avg_sq.add_(eps), -alpha)
+    return None
+
+class Adam(Module):
+    """Adam. Divides gradient EMA by EMA of gradient squares with debiased step size. This implementation is slightly different from
+    pytorch in that debiasing is applied after adding epsilon.
+
+    Args:
+        beta1 (float, optional): momentum. Defaults to 0.9.
+        beta2 (float, optional): second momentum. Defaults to 0.999.
+        eps (float, optional): epsilon. Defaults to 1e-8.
+        alpha (float, optional): learning rate. Defaults to 1.
+        amsgrad (bool, optional): Whether to divide by maximum of EMA of gradient squares instead. Defaults to False.
+        pow (float, optional): power used in second momentum power and root. Defaults to 2.
+        debiased (bool, optional): whether to apply debiasing to momentums based on current step. Defaults to True.
+    """
+    def __init__(
+        self,
+        beta1: float = 0.9,
+        beta2: float = 0.999,
+        eps: float = 1e-8,
+        amsgrad: bool = False,
+        alpha: float = 1.,
+        pow: float = 2,
+        debiased: bool = True,
+    ):
+        defaults=dict(beta1=beta1,beta2=beta2,eps=eps,alpha=alpha,amsgrad=amsgrad,pow=pow,debiased=debiased)
+        super().__init__(defaults)
+        self.getter = itemgetter('amsgrad','pow','debiased')
+
+    @torch.no_grad
+    def step(self, vars):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        beta1,beta2,eps,alpha=self.get_settings('beta1','beta2','eps','alpha', params=vars.params, cls=NumberList)
+        amsgrad,pow,debiased = self.getter(self.settings[vars.params[0]])
+
+        if amsgrad:
+            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg','exp_avg_sq','max_exp_avg_sq', params=vars.params, cls=TensorList)
+        else:
+            exp_avg, exp_avg_sq = self.get_state('exp_avg','exp_avg_sq', params=vars.params, cls=TensorList)
+            max_exp_avg_sq = None
+
+        # if this is last module, update parameters in-place with slightly more efficient addcdiv_
+        if vars.is_last:
+            if vars.last_module_lrs is not None: alpha = alpha * vars.last_module_lrs
+            passed_params = TensorList(vars.params)
+            vars.stop = True
+            vars.skip_update = True
+
+        else:
+            passed_params = None
+
+        vars.update = adam_(
+            tensors=TensorList(vars.get_update()),
+            exp_avg_=exp_avg,
+            exp_avg_sq_=exp_avg_sq,
+            alpha=alpha,
+            beta1=beta1,
+            beta2=beta2,
+            eps=eps,
+            step=step,
+            pow=pow,
+            debiased=debiased,
+            max_exp_avg_sq_=max_exp_avg_sq,
+            params_=passed_params,
+        )
+
+        return vars

torchzero/modules/optimizers/lion.py

@@ -1,15 +1,21 @@
 import torch
 
-from ...core import OptimizerModule
-from ...tensorlist import TensorList
+from ...core import Module, Target, Transform
+from ...utils import NumberList, TensorList
 
 
-def _lion_step_(ascent: TensorList, ema: TensorList, beta1, beta2,):
-    update = ema.lerp_compat(ascent, 1-beta1).sign_()
-    ema.lerp_compat_(ascent, 1-beta2)
+def lion_(tensors: TensorList, exp_avg_: TensorList, beta1, beta2,):
+    """
+    Lion update rule.
+
+    Returns new tensors.
+    """
+    update = exp_avg_.lerp(tensors, 1-beta1).sign_()
+    exp_avg_.lerp_(tensors, 1-beta2)
     return update
 
-class Lion(OptimizerModule):
+
+class Lion(Transform):
     """Lion (EvoLved Sign Momentum) optimizer from https://arxiv.org/abs/2302.06675.
 
     Args:
@@ -19,10 +25,11 @@ class Lion(OptimizerModule):
 
     def __init__(self, beta1: float = 0.9, beta2: float = 0.99):
         defaults = dict(beta1=beta1, beta2=beta2)
-        super().__init__(defaults)
+        super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def _update(self, vars, ascent):
-        beta1, beta2 = self.get_group_keys('beta1', 'beta2')
-        ema = self.get_state_key('ema')
-        return _lion_step_(ascent,ema,beta1,beta2)
+    def transform(self, tensors, params, grads, vars):
+        beta1, beta2 = self.get_settings('beta1', 'beta2', params = params, cls=NumberList)
+        exp_avg = self.get_state('ema', params=params, cls=TensorList)
+        return lion_(TensorList(tensors),exp_avg,beta1,beta2)
+