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

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)

torchzero/modules/adaptive/adahessian.py

@@ -0,0 +1,224 @@
+import math
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, Transform, apply_transform
+from ...utils import NumberList, TensorList, as_tensorlist
+from ..functional import debiased_step_size
+
+def _full_average(hvp: torch.Tensor):
+    if hvp.ndim >= 3:  # Conv kernel
+        return torch.mean(hvp.abs(), dim=[2, *range(3,hvp.ndim)], keepdim=True)
+    return hvp
+
+def _block_average(x: torch.Tensor, block_size: int | None, enable: bool):
+    """averages x over first dimension in blocks"""
+    if enable and x.ndim >= 2:
+        if math.prod(x.shape[1:]) <= 1: return x
+        if block_size is None: return _full_average(x)
+        size = x.size(0)
+
+        n_blocks = size // block_size
+        if n_blocks <= 1: return x.abs().mean(0, keepdim = True)
+
+        n_remaining = size - n_blocks * block_size
+        remaining = None
+        if n_remaining > 0:
+            remaining = x[-n_remaining:].abs().mean(0, keepdim=True).repeat_interleave(n_remaining, 0)
+            x = x[:-n_remaining]
+
+        x = x.view(block_size, n_blocks, *x.shape[1:])
+        x_mean = x.abs().mean(0).repeat_interleave(block_size, 0)
+
+        if remaining is None: return x_mean
+        return torch.cat([x_mean, remaining], 0)
+
+    return x
+
+
+def _rademacher_like(tensor, p = 0.5, generator = None):
+    """p is probability of a 1, other values will be -1."""
+    return torch.bernoulli(torch.full_like(tensor, p), generator = generator).mul_(2).sub_(1)
+
+def adahessian(
+    tensors: TensorList,
+    D: TensorList | None,
+    exp_avg_: TensorList,
+    D_exp_avg_sq_: TensorList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    update_freq: int,
+    eps: float | NumberList,
+    hessian_power: float | NumberList,
+    step: int,
+):
+    # momentum
+    exp_avg_.lerp_(tensors, 1-beta1)
+
+    # update preconditioner
+    if step % update_freq == 0:
+        assert D is not None
+        D_exp_avg_sq_.mul_(beta2).addcmul_(D, D, 1-beta2)
+
+    else:
+        assert D is None
+
+
+    denom = D_exp_avg_sq_.sqrt().pow_(hessian_power).add_(eps)
+    num = exp_avg_ * debiased_step_size(step+1, beta1, beta2)
+
+    return num.div_(denom)
+
+
+class AdaHessian(Module):
+    """AdaHessian: An Adaptive Second Order Optimizer for Machine Learning (https://arxiv.org/abs/2006.00719)
+
+    This is similar to Adam, but the second momentum is replaced by square root of an exponential moving average of random hessian-vector products.
+
+    Notes:
+        - In most cases AdaHessian should be the first module in the chain because it relies on autograd. Use the ``inner`` argument if you wish to apply AdaHessian preconditioning to another module's output.
+
+        - If you are using gradient estimators or reformulations, set ``hvp_method`` to "forward" or "central".
+
+        - This module requires a closure passed to the optimizer step, as it needs to re-evaluate the loss and gradients for calculating HVPs. The closure must accept a ``backward`` argument (refer to documentation).
+
+    Args:
+        beta1 (float, optional): first momentum. Defaults to 0.9.
+        beta2 (float, optional): second momentum for squared hessian diagonal estimates. Defaults to 0.999.
+        averaging (bool, optional):
+            whether to enable block diagonal averaging over 1st dimension on parameters that have 2+ dimensions.
+            This can be set per-parameter in param groups.
+        block_size (int, optional):
+            size of block in the block-diagonal averaging.
+        update_freq (int, optional):
+            frequency of updating hessian diagonal estimate via a hessian-vector product.
+            This value can be increased to reduce computational cost. Defaults to 1.
+        eps (float, optional):
+            division stability epsilon. Defaults to 1e-8.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        fd_h (float, optional): finite difference step size if ``hvp_method`` is "forward" or "central". Defaults to 1e-3.
+        n_samples (int, optional):
+            number of hessian-vector products with random vectors to evaluate each time when updating
+            the preconditioner. Larger values may lead to better hessian diagonal estimate. Defaults to 1.
+        seed (int | None, optional): seed for random vectors. Defaults to None.
+        inner (Chainable | None, optional):
+            Inner module. If this is specified, operations are performed in the following order.
+            1. compute hessian diagonal estimate.
+            2. pass inputs to ``inner``.
+            3. momentum and preconditioning are applied to the ouputs of ``inner``.
+
+    ## Examples:
+
+    Using AdaHessian:
+
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.AdaHessian(),
+        tz.m.LR(0.1)
+    )
+    ```
+
+    AdaHessian preconditioner can be applied to any other module by passing it to the ``inner`` argument.
+    Turn off AdaHessian's first momentum to get just the preconditioning. Here is an example of applying
+    AdaHessian preconditioning to nesterov momentum (``tz.m.NAG``):
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.AdaHessian(beta1=0, inner=tz.m.NAG(0.9)),
+        tz.m.LR(0.1)
+    )
+    ```
+
+    """
+    def __init__(
+        self,
+        beta1: float = 0.9,
+        beta2: float = 0.999,
+        averaging: bool = True,
+        block_size: int | None = None,
+        update_freq: int = 1,
+        eps: float = 1e-8,
+        hessian_power: float = 1,
+        hvp_method: Literal['autograd', 'forward', 'central'] = 'autograd',
+        fd_h: float = 1e-3,
+        n_samples = 1,
+        seed: int | None = None,
+        inner: Chainable | None = None
+    ):
+        defaults = dict(beta1=beta1, beta2=beta2, update_freq=update_freq, averaging=averaging, block_size=block_size, eps=eps, hessian_power=hessian_power, hvp_method=hvp_method, n_samples=n_samples, fd_h=fd_h, seed=seed)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, var):
+        params = var.params
+        settings = self.settings[params[0]]
+        hvp_method = settings['hvp_method']
+        fd_h = settings['fd_h']
+        update_freq = settings['update_freq']
+        n_samples = settings['n_samples']
+
+        seed = settings['seed']
+        generator = self.get_generator(params[0].device, seed)
+
+        beta1, beta2, eps, averaging, block_size, hessian_power = self.get_settings(params,
+            'beta1', 'beta2', 'eps', 'averaging', 'block_size', "hessian_power", cls=NumberList)
+
+        exp_avg, D_exp_avg_sq = self.get_state(params, 'exp_avg', 'h_exp_avg', cls=TensorList)
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        closure = var.closure
+        assert closure is not None
+
+        D = None
+        if step % update_freq == 0:
+
+            rgrad=None
+            for i in range(n_samples):
+                u = [_rademacher_like(p, generator=generator) for p in params]
+
+                Hvp, rgrad = self.Hvp(u, at_x0=True, var=var, rgrad=rgrad, hvp_method=hvp_method,
+                                     h=fd_h, normalize=True, retain_grad=i < n_samples-1)
+                Hvp = tuple(Hvp)
+
+                if D is None: D = Hvp
+                else: torch._foreach_add_(D, Hvp)
+
+            assert D is not None
+            if n_samples > 1: torch._foreach_div_(D, n_samples)
+
+            D = TensorList(D).zipmap_args(_block_average, block_size, averaging)
+
+        update = var.get_update()
+        if 'inner' in self.children:
+            update = apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var)
+
+        var.update = adahessian(
+            tensors=TensorList(update),
+            D=TensorList(D) if D is not None else None,
+            exp_avg_=exp_avg,
+            D_exp_avg_sq_=D_exp_avg_sq,
+            beta1=beta1,
+            beta2=beta2,
+            update_freq=update_freq,
+            eps=eps,
+            hessian_power=hessian_power,
+            step=step,
+        )
+        return var

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

@@ -10,9 +10,6 @@ from ..functional import (
     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_(
@@ -33,7 +30,7 @@ def adam_(
     params: list[torch.Tensor] | None = None,
     grads: list[torch.Tensor] | None = None,
 ):
-    """Returns new tensors or updates params in-place."""
+    """Returns new tensors."""
     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)
 
@@ -43,11 +40,12 @@ def adam_(
 
     exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
     if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
-    return (exp_avg_ / sqrt_exp_avg_sq.add_(eps)).lazy_mul(alpha)
+    return (exp_avg_.lazy_mul(alpha) / sqrt_exp_avg_sq.add_(eps))
 
 class Adam(Transform):
-    """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.
+    """Adam. Divides gradient EMA by EMA of gradient squares with debiased step size.
+
+    This implementation is identical to :code:`torch.optim.Adam`.
 
     Args:
         beta1 (float, optional): momentum. Defaults to 0.9.
@@ -75,7 +73,7 @@ class Adam(Transform):
         if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
         beta1,beta2,eps,alpha=unpack_dicts(settings, 'beta1','beta2','eps','alpha', cls=NumberList)