torchzero 0.3.15__py3-none-any.whl → 0.4.0__py3-none-any.whl

torchzero/modules/adaptive/adagrad.py

@@ -1,62 +1,14 @@
-from operator import itemgetter
 from typing import Literal
-
 import torch
+
 from ...core import (
     Chainable,
-    Module,
-    Target,
-    TensorwiseTransform,
-    Transform,
-    Var,
-    apply_transform,
+    TensorTransform,
 )
-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)
+from ...utils import NumberList, TensorList, unpack_dicts
+from ...linalg.matrix_power import matrix_power as _matrix_power, MatrixPowerMethod
 
-    return tensors_
-
-
-
-class Adagrad(Transform):
+class Adagrad(TensorTransform):
     """Adagrad, divides by sum of past squares of gradients.
 
     This implementation is identical to ``torch.optim.Adagrad``.
@@ -72,103 +24,53 @@ class Adagrad(Transform):
     """
     def __init__(
         self,
+
+        # hyperparams
         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,
+
+        # tfms
         inner: Chainable | None = None,
+        accumulator_tfm: 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)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults["accumulator_tfm"]
+        super().__init__(defaults=defaults, inner=inner)
 
-        if inner is not None:
-            self.set_child('inner', inner)
+        self.set_child('accumulator', accumulator_tfm)
 
     @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
+    def single_tensor_initialize(self, tensor, param, grad, loss, state, setting):
+        state["accumulator"] = torch.full_like(tensor, fill_value=setting["initial_accumulator_value"])
 
-        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)
+    @torch.no_grad
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
+        torch._foreach_addcmul_([state["accumulator"] for state in states], tensors, tensors)
+        self.increment_counter("step", start=0)
 
-    gg = tensors_.dot(tensors_)
+    @torch.no_grad
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        tensors_ = TensorList(tensors)
+        step = self.global_state["step"] # 0 on first apply
+        eps, alpha, lr_decay = unpack_dicts(settings, "eps", "alpha", "lr_decay", cls=NumberList)
 
-    if beta is None or step == 1: accumulator += gg
-    else: accumulator = lerp(accumulator, gg, 1-beta)
+        accumulator = [state["accumulator"] for state in states]
+        accumulator = TensorList(self.inner_step_tensors(
+            "accumulator", tensors=accumulator, clone=True, params=params, grads=grads, loss=loss, must_exist=False))
 
-    if decay is not None:
-        accumulator *= 1-decay
+        denom = accumulator.sqrt().add_(eps)
+        tensors_ /= denom
 
-    if inner is not None:
-        assert params is not None
-        tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
+        clr = alpha / (1 + step * lr_decay)
+        tensors_.lazy_mul_(clr)
 
-    if divide: accumulator = accumulator / max(step, 1)
+        return tensors_
 
-    if use_sqrt: tensors_.div_(eps + accumulator.sqrt()).mul_(clr)
-    else: tensors_.div_(eps + accumulator).mul_(clr)
 
-    return tensors_, accumulator
 
-class AdagradNorm(Transform):
+class AdagradNorm(TensorTransform):
     """Adagrad-Norm, divides by sum of past means of squares of gradients.
 
     Args:
@@ -176,7 +78,6 @@ class AdagradNorm(Transform):
         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.
     """
@@ -185,71 +86,104 @@ class AdagradNorm(Transform):
         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,
+        beta_debias: bool = True,
+        layerwise: bool = True,
+        use_sqrt: bool = True,
+        alpha: float = 1,
         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)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner']
+        super().__init__(defaults=defaults, inner=inner)
+
+    @torch.no_grad
+    def multi_tensor_initialize(self, tensors, params, grads, loss, states, settings):
+
+        # layerwise initialize in each state
+        if settings[0]["layerwise"]:
+            for tensor, state, setting in zip(tensors, states, settings):
+
+                initial_accumulator_value = setting["initial_accumulator_value"]
+                state["accumulator"] = torch.tensor(initial_accumulator_value, device=tensor.device, dtype=tensor.dtype)
+
+        # global initialize in global state
+        else:
+            initial_accumulator_value = settings[0]["initial_accumulator_value"]
+            tensor = tensors[0]
+            self.global_state["accumulator"] = torch.tensor(initial_accumulator_value, device=tensor.device, dtype=tensor.dtype)
+
+    def _get_accumulator(self, states, settings) -> torch.Tensor | TensorList:
+        layerwise = settings[0]["layerwise"]
+        if layerwise:
+            return TensorList(s["accumulator"] for s in states)
+
+        return self.global_state["accumulator"]
 
-        if inner is not None:
-            self.set_child('inner', inner)
+    @torch.no_grad
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        accumulator = self._get_accumulator(states, settings)
+        self.increment_counter("step", start=0)
+
+        # compute squared gradient norm (gg)
+        if isinstance(accumulator, TensorList): gg = tensors.tensorwise_dot(tensors)
+        else: gg = tensors.dot(tensors)
+
+        # update the accumulator
+        beta = settings[0]["beta"]
+        if beta is None: accumulator.add_(gg) # pyright:ignore[reportArgumentType]
+        else: accumulator.lerp_(gg, weight=1-beta) # pyright:ignore[reportArgumentType, reportCallIssue]
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(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)
+        accumulator = self._get_accumulator(states, settings)
+        eps, alpha, lr_decay = unpack_dicts(settings, "eps", "alpha", "lr_decay", cls=NumberList)
+        step = self.global_state["step"] # 0 on 1st step
+        fs = settings[0]
+        beta = fs["beta"]
 
-        use_sqrt, divide, initial_accumulator_value = itemgetter('use_sqrt', 'divide', "initial_accumulator_value")(settings[0])
+        # ------------------------ debias if beta is not None ------------------------ #
+        if fs["beta_debias"] and beta is not None:
+            accumulator = accumulator / (1 - beta ** (step + 1))
 
-        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,
+        # ---------------------------- compute denominator --------------------------- #
+        if fs["use_sqrt"]:
+            denom = accumulator.sqrt().add_(eps) # pyright:ignore[reportArgumentType]
+        else:
+            denom = accumulator + eps # pyright:ignore[reportOperatorIssue]
 
-            beta = self.defaults["beta"],
-            decay = self.defaults["decay"],
-            # inner args
-            inner=self.children.get("inner", None),
-            params=params,
-            grads=grads,
-        )
 
-        return d
+        # ---------------------------- compute the update ---------------------------- #
+        tensors /= denom
+        clr = alpha / (1 + step * lr_decay) # lr decay
+        tensors.lazy_mul_(clr)
 
+        return tensors
 
-class FullMatrixAdagrad(TensorwiseTransform):
+
+
+class FullMatrixAdagrad(TensorTransform):
     """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.
+        reg (float, optional): regularization, scale of identity matrix added to accumulator. Defaults to 1e-12.
         precond_freq (int, optional): frequency of updating the inverse square root of the accumulator. Defaults to 1.
+        beta (float | None, optional): momentum for gradient outer product accumulators. if None, uses sum. Defaults to None.
+        beta_debias (bool, optional): whether to use debiasing, only has effect when ``beta`` is not ``None``. Defaults to True.
         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.
+        matrix_power (float, optional): accumulator matrix power. Defaults to -1/2.
+        concat_params (bool, optional): if False, each parameter will have it's own accumulator. Defaults to True.
         inner (Chainable | None, optional): inner modules to apply preconditioning to. Defaults to None.
 
     ## Examples:
@@ -284,73 +218,89 @@ class FullMatrixAdagrad(TensorwiseTransform):
     """
     def __init__(
         self,
+        reg: float = 1e-12,
+        precond_freq: int = 1,
         beta: float | None = None,
-        decay: float | None = None,
-        sqrt: bool = True,
+        beta_debias: bool=True,
+        init: Literal["identity", "zeros", "GGT"] = "identity",
+        matrix_power: float = -1/2,
+        matrix_power_method: MatrixPowerMethod = "eigh_abs",
         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,
+        accumulator_tfm: 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,)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults["concat_params"], defaults["accumulator_tfm"]
+        super().__init__(defaults=defaults, inner=inner, concat_params=concat_params)
+
+        self.set_child("accumulator", accumulator_tfm)
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, loss, state, setting):
+    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
+
         G = tensor.ravel()
-        GG = torch.outer(G, G)
-        decay = setting['decay']
+        GGᵀ = torch.outer(G, G)
+
+        # initialize
+        if "accumulator" not in state:
+            init = setting['init']
+            if init == 'identity': state['accumulator'] = torch.eye(GGᵀ.size(0), device=GGᵀ.device, dtype=GGᵀ.dtype)
+            elif init == 'zeros': state['accumulator'] =  torch.zeros_like(GGᵀ)
+            elif init == 'GGT': state['accumulator'] = GGᵀ.clone()
+            else: raise ValueError(init)
+
+        # update
         beta = setting['beta']
-        init = setting['init']
+        accumulator: torch.Tensor = state["accumulator"]
 
-        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 None: accumulator.add_(GGᵀ)
+        else: accumulator.lerp_(GGᵀ, 1-beta)
 
-        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
+        # update number of GGᵀ in accumulator for divide
+        state['num_GGTs'] = state.get('num_GGTs', 0) + 1
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+    def single_tensor_apply(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']
+        accumulator: torch.Tensor = state['accumulator']
+        accumulator = self.inner_step_tensors("accumulator", [accumulator], clone=True, must_exist=False)[0]
+
         precond_freq = setting['precond_freq']
         reg = setting['reg']
+        beta = setting["beta"]
 
-        if divide: GG = GG/state.get('i', 1)
-
+        # add regularizer
         if reg != 0:
-            GG = GG + torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype).mul_(reg)
+            device = accumulator.device; dtype = accumulator.dtype
+            accumulator = accumulator + torch.eye(accumulator.size(0), device=device, dtype=dtype).mul_(reg)
 
+        # for single value use sqrt
         if tensor.numel() == 1:
-            GG = GG.squeeze()
-            if sqrt: return tensor / GG.sqrt()
-            return tensor / GG
+            dir = tensor.mul_(accumulator.squeeze() ** setting["matrix_power"])
 
-        try:
-            if sqrt:
+        # otherwise use matrix inverse square root
+        else:
+
+            # compute inverse square root and store to state
+            try:
                 if "B" not in state or step % precond_freq == 0:
-                    B = state["B"] = matrix_power_eigh(GG, -1/2)
+                    B = state["B"] = _matrix_power(accumulator, setting["matrix_power"], method=setting["matrix_power_method"])
                 else:
                     B = state["B"]
 
-            else: return torch.linalg.solve(GG, tensor.ravel()).view_as(tensor) # pylint:disable = not-callable
+                dir = (B @ tensor.ravel()).view_as(tensor)
+
+            # fallback to diagonal Adagrad on fail
+            except torch.linalg.LinAlgError:
+                dir = tensor.mul_(accumulator.diagonal() ** setting["matrix_power"])
 
-        except torch.linalg.LinAlgError:
-            # fallback to diagonal AdaGrad
-            denom = GG.diagonal()
-            if sqrt: denom = denom.sqrt()
-            return tensor.div_(denom + max(reg, 1e-12))
+        # debias
+        if setting["beta_debias"] and beta is not None:
+            num_GGTs = state.get('num_GGTs', 1)
+            bias_correction = 1 - beta ** num_GGTs
+            dir *= bias_correction ** 0.5
 
-        return (B @ tensor.ravel()).view_as(tensor)
+        return dir