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

torchzero/modules/optimizers/rprop.py

@@ -1,99 +1,342 @@
-from collections import abc
-
-import torch
-
-from ...tensorlist import TensorList, where
-from ...core import OptimizerModule
-
-
-def _bool_ones_like(x):
-    return torch.ones_like(x, dtype=torch.bool)
-
-class Rprop(OptimizerModule):
-    """
-    Resilient propagation. The update magnitude gets multiplied by `nplus` if gradient didn't change the sign,
-    or `nminus` if it did. Then the update is applied with the sign of the current gradient.
-
-    Additionally, if gradient changes sign, the update for that weight is reverted.
-    Next step, magnitude for that weight won't change.
-
-    Compared to pytorch this also implements backtracking update when sign changes.
-    To make this behave exactly the same as `torch.optim.Rprop`, set `backtrack` to False.
-
-    Args:
-        nplus (float): multiplicative increase factor for when ascent didn't change sign (default: 1.2).
-        nminus (float): multiplicative decrease factor for when ascent changed sign (default: 0.5).
-        lb (float): minimum step size, can be None (default: 1e-6)
-        ub (float): maximum step size, can be None (default: 50)
-        backtrack (float):
-            if True, when ascent sign changes, undoes last weight update, otherwise sets update to 0.
-            When this is False, this exactly matches pytorch Rprop. (default: True)
-        alpha (float): learning rate (default: 1).
-
-    reference
-        *Riedmiller, M., & Braun, H. (1993, March). A direct adaptive method for faster backpropagation learning:
-        The RPROP algorithm. In IEEE international conference on neural networks (pp. 586-591). IEEE.*
-    """
-    def __init__(
-        self,
-        nplus: float = 1.2,
-        nminus: float = 0.5,
-        lb: float | None = 1e-6,
-        ub: float | None = 50,
-        backtrack=True,
-        alpha: float = 1,
-    ):
-        defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub)
-        super().__init__(defaults)
-        self.current_step = 0
-        self.backtrack = backtrack
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-
-        sign = ascent.sign_()
-        nplus, nminus, lb, ub = self.get_group_keys('nplus', 'nminus', 'lb', 'ub')
-        prev, allowed, magnitudes = self.get_state_keys(
-            'prev', 'allowed', 'magnitudes',
-            inits = [torch.zeros_like, _bool_ones_like, torch.zeros_like],
-            params=params
-        )
-
-        # initialize on 1st step
-        if self.current_step == 0:
-            magnitudes.fill_(self.get_group_key('alpha')).clamp_(lb, ub)
-            ascent = magnitudes * sign
-            prev.copy_(ascent)
-            self.current_step += 1
-            return ascent
-
-        mul = (sign * prev).mul_(allowed)
-
-        sign_changed = mul < 0
-        sign_same = mul > 0
-        zeroes = mul == 0
-
-        mul.fill_(1)
-        mul.masked_fill_(sign_changed, nminus)
-        mul.masked_fill_(sign_same, nplus)
-
-        # multiply magnitudes based on sign change and clamp to bounds
-        magnitudes.mul_(mul).clamp_(lb, ub)
-
-        # revert update if sign changed
-        if self.backtrack:
-            ascent = sign.mul_(magnitudes)
-            ascent.masked_set_(sign_changed, prev.neg_())
-        else:
-            ascent = sign.mul_(magnitudes * ~sign_changed)
-
-        # update allowed to only have weights where last update wasn't reverted
-        allowed.set_(sign_same | zeroes)
-
-        prev.copy_(ascent)
-        self.current_step += 1
-        return ascent
-
-
-
+
+import torch
+
+from ...core import Module, Target, Transform
+from ...utils import NumberList, TensorList, as_tensorlist
+
+
+def _bool_ones_like(x):
+    return torch.ones_like(x, dtype=torch.bool)
+
+def sign_consistency_lrs_(
+    tensors: TensorList,
+    prev_: TensorList,
+    lrs_: TensorList,
+    nplus: float | NumberList,
+    nminus: float | NumberList,
+    lb: float | NumberList,
+    ub: float | NumberList,
+    step: int,
+):
+    """returns `lrs_`"""
+    sign = tensors.sign()
+    if step == 0:
+        prev_.set_(sign)
+        return lrs_.clamp_(lb, ub)
+
+    mul = sign * prev_
+    prev_.set_(sign)
+
+    sign_changed = mul < 0
+    sign_same = mul > 0
+
+    mul.fill_(1)
+    mul.masked_fill_(sign_changed, nminus)
+    mul.masked_fill_(sign_same, nplus)
+
+    # multiply magnitudes based on sign change and clamp to bounds
+    lrs_.mul_(mul).clamp_(lb, ub)
+    return lrs_
+
+def scale_by_sign_change_(
+    tensors_: TensorList,
+    cur: TensorList,
+    prev_: TensorList,
+    lrs_: TensorList,
+    nplus: float | NumberList,
+    nminus: float | NumberList,
+    lb: float | NumberList,
+    ub: float | NumberList,
+    step: int,
+):
+    """returns `tensors_`"""
+    lrs_ = sign_consistency_lrs_(cur,prev_=prev_,lrs_=lrs_,nplus=nplus,nminus=nminus,
+                             lb=lb,ub=ub,step=step)
+    return tensors_.mul_(lrs_)
+
+def backtrack_on_sign_change_(
+    tensors_: TensorList,
+    cur: TensorList,
+    prev_: TensorList,
+    backtrack: bool,
+    step: int
+):
+    """returns `tensors_`."""
+    if step == 0:
+        prev_.set_(cur)
+        return tensors_
+
+    # mask will be > 0 for parameters where both signs are the same
+    mask = (cur * prev_) < 0
+    if backtrack: tensors_.masked_set_(mask, prev_)
+    else: tensors_.select_set_(mask, 0)
+
+    prev_.set_(cur)
+    return tensors_
+
+def rprop_(
+    tensors_: TensorList,
+    prev_: TensorList,
+    allowed_: TensorList,
+    magnitudes_: TensorList,
+    nplus: float | NumberList,
+    nminus: float | NumberList,
+    lb: float | NumberList,
+    ub: float | NumberList,
+    alpha: float | NumberList,
+    backtrack: bool,
+    step: int,
+):
+    """returns new tensors."""
+
+    sign = tensors_.sign_()
+
+    # initialize on 1st step
+    if step == 0:
+        magnitudes_.fill_(alpha).clamp_(lb, ub)
+        new_tensors = magnitudes_ * sign
+        prev_.copy_(new_tensors)
+        return new_tensors
+
+    mul = (sign * prev_).mul_(allowed_)
+
+    sign_changed = mul < 0
+    sign_same = mul > 0
+    zeroes = mul == 0
+
+    mul.fill_(1)
+    mul.masked_fill_(sign_changed, nminus)
+    mul.masked_fill_(sign_same, nplus)
+
+    # multiply magnitudes based on sign change and clamp to bounds
+    magnitudes_.mul_(mul).clamp_(lb, ub)
+
+    # revert update if sign changed
+    if backtrack:
+        new_tensors = sign.mul_(magnitudes_)
+        new_tensors.masked_set_(sign_changed, prev_.neg_())
+    else:
+        new_tensors = sign.mul_(magnitudes_ * ~sign_changed)
+
+    # update allowed to only have weights where last update wasn't reverted
+    allowed_.set_(sign_same | zeroes)
+
+    prev_.copy_(new_tensors)
+    return new_tensors
+
+
+
+class Rprop(Transform):
+    """
+    Resilient propagation. The update magnitude gets multiplied by `nplus` if gradient didn't change the sign,
+    or `nminus` if it did. Then the update is applied with the sign of the current gradient.
+
+    Additionally, if gradient changes sign, the update for that weight is reverted.
+    Next step, magnitude for that weight won't change.
+
+    Compared to pytorch this also implements backtracking update when sign changes.
+    To make this behave exactly the same as `torch.optim.Rprop`, set `backtrack` to False.
+
+    Args:
+        nplus (float): multiplicative increase factor for when ascent didn't change sign (default: 1.2).
+        nminus (float): multiplicative decrease factor for when ascent changed sign (default: 0.5).
+        lb (float): minimum step size, can be None (default: 1e-6)
+        ub (float): maximum step size, can be None (default: 50)
+        backtrack (float):
+            if True, when ascent sign changes, undoes last weight update, otherwise sets update to 0.
+            When this is False, this exactly matches pytorch Rprop. (default: True)
+        alpha (float): initial per-parameter learning rate (default: 1).
+
+    reference
+        *Riedmiller, M., & Braun, H. (1993, March). A direct adaptive method for faster backpropagation learning:
+        The RPROP algorithm. In IEEE international conference on neural networks (pp. 586-591). IEEE.*
+    """
+    def __init__(
+        self,
+        nplus: float = 1.2,
+        nminus: float = 0.5,
+        lb: float = 1e-6,
+        ub: float = 50,
+        backtrack=True,
+        alpha: float = 1,
+    ):
+        defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub, backtrack=backtrack)
+        self.current_step = 0
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        nplus, nminus, lb, ub, alpha = self.get_settings('nplus', 'nminus', 'lb', 'ub', 'alpha', params=params, cls=NumberList)
+        prev, allowed, magnitudes = self.get_state(
+            'prev','allowed','magnitudes',
+            params=params,
+            init=[torch.zeros_like, _bool_ones_like, torch.zeros_like],
+            cls = TensorList,
+        )
+
+        target = rprop_(
+            tensors_ = as_tensorlist(tensors),
+            prev_ = prev,
+            allowed_ = allowed,
+            magnitudes_ = magnitudes,
+            nplus = nplus,
+            nminus = nminus,
+            lb = lb,
+            ub = ub,
+            alpha = alpha,
+            backtrack=self.settings[params[0]]['backtrack'],
+            step=self.current_step,
+        )
+
+        self.current_step += 1
+        return target
+
+
+class ScaleLRBySignChange(Transform):
+    """
+    learning rate gets multiplied by `nplus` if ascent/gradient didn't change the sign,
+    or `nminus` if it did.
+
+    This is part of RProp update rule.
+
+    Args:
+        nplus (float): learning rate gets multiplied by `nplus` if ascent/gradient didn't change the sign
+        nminus (float): learning rate gets multiplied by `nminus` if ascent/gradient changed the sign
+        lb (float): lower bound for lr.
+        ub (float): upper bound for lr.
+        alpha (float): initial learning rate.
+
+    """
+
+    def __init__(
+        self,
+        nplus: float = 1.2,
+        nminus: float = 0.5,
+        lb=1e-6,
+        ub=50.0,
+        alpha=1.0,
+        use_grad=False,
+        target: Target = "update",
+    ):
+        defaults = dict(nplus=nplus, nminus=nminus, alpha=alpha, lb=lb, ub=ub, use_grad=use_grad)
+        super().__init__(defaults, uses_grad=use_grad, target=target)
+        self.current_step = 0
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        target = as_tensorlist(tensors)
+        use_grad = self.settings[params[0]]['use_grad']
+        if use_grad: cur = as_tensorlist(grads)
+        else: cur = target
+
+        nplus, nminus, lb, ub = self.get_settings('nplus', 'nminus', 'lb', 'ub', params=params, cls=NumberList)
+        prev, lrs = self.get_state('prev', 'lrs', params=params, cls=TensorList)
+
+        if self.current_step == 0:
+            lrs.set_(target.full_like(self.get_settings('alpha', params=params)))
+
+        target = scale_by_sign_change_(
+            tensors_ = target,
+            cur = cur,
+            prev_ = prev,
+            lrs_ = lrs,
+            nplus = nplus,
+            nminus = nminus,
+            lb = lb,
+            ub = ub,
+            step = self.current_step,
+        )
+        self.current_step += 1
+        return target
+
+class BacktrackOnSignChange(Transform):
+    """Negates or undoes update for parameters where where gradient or update sign changes.
+
+    This is part of RProp update rule.
+
+    Args:
+        normalize (bool, optional): renormalize update after masking. Defaults to False.
+        eps (_type_, optional): epsilon for normalization. Defaults to 1e-6.
+        use_grad (bool, optional):
+            if True, tracks sign change of the gradient,
+            otherwise track sign change of the update. Defaults to True.
+        backtrack (bool, optional):
+            if True, undoes the update when sign changes, otherwise negates it.
+            Defaults to True.
+
+    """
+    def __init__(self, use_grad = False, backtrack = True, target: Target = 'update'):
+        defaults = dict(use_grad=use_grad, backtrack=backtrack, target=target)
+        super().__init__(defaults, uses_grad=use_grad)
+        self.current_step = 0
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        target = as_tensorlist(tensors)
+        settings = self.settings[params[0]]
+        use_grad = settings['use_grad']
+        backtrack = settings['backtrack']
+
+        if use_grad: cur = as_tensorlist(grads)
+        else: cur = target
+
+        target = backtrack_on_sign_change_(
+            tensors_ = target,
+            cur = cur,
+            prev_ = self.get_state('prev', params=params, cls=TensorList),
+            backtrack = backtrack,
+            step = self.current_step,
+        )
+
+        self.current_step += 1
+        return target
+
+class SignConsistencyMask(Transform):
+    """0 if sign changed 1 otherwise"""
+    def __init__(self,target: Target = 'update'):
+        super().__init__({}, uses_grad=False, target = target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        prev = self.get_state('prev', params=params, cls=TensorList)
+        mask = prev.mul_(tensors).gt_(0)
+        prev.set_(tensors)
+        return mask
+
+
+class SignConsistencyLRs(Transform):
+    """LR for each weight is increased when two consequtive update signs are the same, decreased otherwise. This returns the LRs themselves."""
+    def __init__(
+        self,
+        nplus: float = 1.2,
+        nminus: float = 0.5,
+        lb: float | None = 1e-6,
+        ub: float | None = 50,
+        alpha: float = 1,
+        target: Target = 'update'
+    ):
+        defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub)
+        super().__init__(defaults, uses_grad=False, target = target)
+        self.current_step = 0
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        target = as_tensorlist(tensors)
+        nplus, nminus, lb, ub = self.get_settings('nplus', 'nminus', 'lb', 'ub', params=params, cls=NumberList)
+        prev, lrs = self.get_state('prev', 'lrs', params=params, cls=TensorList)
+
+        if self.current_step == 0:
+            lrs.set_(target.full_like(self.get_settings('alpha', params=params)))
+
+        target = sign_consistency_lrs_(
+            tensors = target,
+            prev_ = prev,
+            lrs_ = lrs,
+            nplus = nplus,
+            nminus = nminus,
+            lb = lb,
+            ub = ub,
+            step = self.current_step,
+        )
+        self.current_step += 1
+        return target.clone()

torchzero/modules/optimizers/shampoo.py

@@ -0,0 +1,197 @@
+from collections.abc import Sequence
+from operator import itemgetter
+from functools import partial
+import numpy as np
+import torch
+
+from ...core import Chainable, Transform, apply
+from ...utils.linalg import matrix_power_eigh
+from ...utils import set_storage_
+
+
+def update_shampoo_preconditioner_(
+    grad: torch.Tensor,
+    accumulators_: list[torch.Tensor | None],
+    preconditioners_: list[torch.Tensor | None],
+    step: int,
+    update_freq: int,
+    exp_override: int | None,
+    beta: float | None,
+):
+    for i, (accumulator, preconditioner) in enumerate(zip(accumulators_, preconditioners_)):
+        if accumulator is None: continue
+        assert preconditioner is not None
+
+        axes = list(range(i)) + list(range(i + 1, grad.ndim))
+        if beta is None: accumulator.add_(torch.tensordot(grad, grad, (axes, axes))) # pyright:ignore[reportArgumentType]
+        else: accumulator.lerp_(torch.tensordot(grad, grad, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
+
+        if step % update_freq == 0:
+            matrix_exp = -1/(grad.ndim*2) if exp_override is None else -1/exp_override
+            set_storage_(preconditioner, matrix_power_eigh(accumulator, matrix_exp))
+
+
+def apply_shampoo_preconditioner(
+    tensor: torch.Tensor,
+    preconditioners_: list[torch.Tensor | None],
+    decay: float | None,
+):
+    for i, preconditioner in enumerate(preconditioners_):
+        if preconditioner is None: continue
+        tensor = torch.tensordot(tensor, preconditioner, ([0], [0])) # pyright:ignore[reportArgumentType]
+        if decay is not None: preconditioner.mul_(decay)
+    return tensor
+
+
+def update_diagonal_(grad: torch.Tensor, diagonal_accumulator_: torch.Tensor, beta: float | None):
+    if beta is None: diagonal_accumulator_.add_(grad.pow(2))
+    else: diagonal_accumulator_.mul_(beta).addcmul_(grad, grad, value=1-beta)
+
+def apply_diagonal_(grad_: torch.Tensor, diagonal_accumulator_: torch.Tensor, decay: float | None, eps: float):
+    grad_.div_(diagonal_accumulator_.sqrt() + eps)
+    if decay is not None: diagonal_accumulator_.mul_(decay)
+    return grad_
+
+def _merge_small_dims(tensor: torch.Tensor, max_dim: int):
+    """a safer merger"""
+    if tensor.ndim == 0: return tensor, None, None
+    sort_idxs = np.argsort(tensor.shape)
+    if tensor.shape[sort_idxs[0]] > max_dim:
+        return tensor, None, None
+
+    tensor = tensor.permute(*sort_idxs)
+    flatten_end_idx = 0
+    flat_sizes = []
+    flat_numel = 1
+    for i, size in enumerate(tensor.shape):
+        if flat_numel * size <= max_dim:
+            flatten_end_idx = i
+            flat_numel *= size
+            flat_sizes.append(size)
+        else:
+            break
+
+    if flatten_end_idx != 0:
+        tensor = tensor.flatten(end_dim=flatten_end_idx)
+
+    return tensor, flat_sizes, sort_idxs
+
+def _unmerge_small_dims(tensor: torch.Tensor, flat_sizes: Sequence[int] | None, sort_idxs: np.ndarray | Sequence[int] | None):
+    if flat_sizes is None: return tensor
+    assert sort_idxs is not None
+    tensor = tensor.unflatten(0, flat_sizes)
+    return tensor.permute(*np.argsort(sort_idxs))
+
+
+class Shampoo(Transform):
+    """Shampoo from Preconditioned Stochastic Tensor Optimization (https://arxiv.org/abs/1802.09568).
+
+    Args:
+        decay (float | None, optional): slowly decays preconditioners. Defaults to None.
+        beta (float | None, optional):
+            if None calculates sum as in standard shampoo, otherwise uses EMA of preconditioners. Defaults to None.
+        matrix_eps (float, optional): epsilon for matrix operations. Defaults to 1e-10.
+        update_freq (int, optional): preconditioner update frequency. Defaults to 10.
+        exp_override (int | None, optional): matrix exponent override, if not set, uses 2*ndim. Defaults to None.
+        merge_small (bool, optional): whether to merge small dims on tensors. Defaults to True.
+        max_dim (int, optional): maximum dimension size for preconditioning. Defaults to 2_000.
+        precondition_1d (bool, optional): whether to precondition 1d tensors. Defaults to True.
+        adagrad_eps (float, optional): epsilon for adagrad division for tensors where shampoo can't be applied. Defaults to 1e-8.
+        inner (Chainable | None, optional):
+            module applied after updating preconditioners and before applying preconditioning.
+            For example if beta≈0.999 and `inner=tz.m.EMA(0.9)`, this becomes Adam with shampoo preconditioner (ignoring debiasing).
+            Defaults to None.
+    """
+    def __init__(
+        self,
+        decay: float | None = None,
+        beta: float | None = None,
+        reg: float = 1e-6,
+        update_freq: int = 10,
+        exp_override: int | None = None,
+        merge_small: bool = True,
+        max_dim: int = 2_000,
+        precondition_1d: bool = True,
+        adagrad_eps: float = 1e-8,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(decay=decay, beta=beta, reg=reg, update_freq=update_freq, exp_override=exp_override, merge_small=merge_small, max_dim=max_dim, precondition_1d=precondition_1d,adagrad_eps=adagrad_eps)
+        super().__init__(defaults, uses_grad=False)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    def transform(self, tensors, params, grads, vars):
+        merged_target = [] # target with merged dims
+
+        # update preconditioners
+        for i,(p,t) in enumerate(zip(params, tensors)):
+            state = self.state[p]
+            settings = self.settings[p]
+            beta, reg, update_freq, exp_override, merge_small, max_dim, precondition_1d = itemgetter(
+                'beta', 'reg', 'update_freq', 'exp_override', 'merge_small', 'max_dim', 'precondition_1d')(settings)
+
+            if merge_small:
+                t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
+            merged_target.append(t)
+
+            # initialize accumulators and preconditioners for each dim on 1st step
+            if 'accumulators' not in state:
+
+                if not precondition_1d and t.ndim <= 1:
+                    state['accumulators'] = []
+
+                else:
+                    state['accumulators'] = [torch.eye(s, dtype=t.dtype, device=t.device) if 1<s<max_dim else None for s in t.shape]
+                    state['preconditioners'] = [torch.eye(s, dtype=t.dtype, device=t.device) if 1<s<max_dim else None for s in t.shape]
+
+                # either scalar parameter, 1d with precondition_1d=False, or too big, then basic diagonal preconditioner is used.
+                if len([i is not None for i in state['accumulators']]) == 0:
+                    state['diagonal_accumulator'] = torch.zeros_like(t)
+
+                state['step'] = 0
+
+            # update preconditioners
+            if 'diagonal_accumulator' in state:
+                update_diagonal_(t, state['diagonal_accumulator'], beta)
+            else:
+                update_shampoo_preconditioner_(
+                    t,
+                    accumulators_=state['accumulators'],
+                    preconditioners_=state['preconditioners'],
+                    step=state['step'],
+                    update_freq=update_freq,
+                    exp_override=exp_override,
+                    beta=beta,
+                )
+
+        # inner step
+        if 'inner' in self.children:
+            tensors = apply(self.children['inner'], tensors, params=params, grads=grads, vars=vars)
+
+            # have to merge small dims again
+            merged_target = [] # target with merged dims
+            for i,(p,t) in enumerate(zip(params, tensors)):
+                state = self.state[p]
+                settings = self.settings[p]
+                if settings['merge_small']:
+                    t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, settings['max_dim'])
+                merged_target.append(t)
+
+        # precondition
+        for i, (p, t) in enumerate(zip(params, merged_target)):
+            state = self.state[p]
+            settings = self.settings[p]
+            decay, merge_small, adagrad_eps= itemgetter('decay', 'merge_small', 'adagrad_eps')(settings)
+
+            if 'diagonal_accumulator' in state:
+                tensors[i] = apply_diagonal_(t, state['diagonal_accumulator'], decay=decay, eps=adagrad_eps)
+            else:
+                tensors[i] = apply_shampoo_preconditioner(t, preconditioners_=state['preconditioners'], decay=decay)
+
+            if merge_small:
+                tensors[i] = _unmerge_small_dims(tensors[i], state['flat_sizes'], state['sort_idxs'])
+
+            state['step'] += 1
+
+        return tensors