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

torchzero/modules/optimizers/rprop.py

@@ -2,7 +2,7 @@
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList, as_tensorlist
+from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states
 
 
 def _bool_ones_like(x):
@@ -135,7 +135,8 @@ class Rprop(Transform):
     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.
+
+    This implementation is identical to :code:`torch.optim.Rprop` if :code:`backtrack` is set to False.
 
     Args:
         nplus (float): multiplicative increase factor for when ascent didn't change sign (default: 1.2).
@@ -161,20 +162,22 @@ class Rprop(Transform):
         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(
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        nplus, nminus, lb, ub, alpha = unpack_dicts(settings, 'nplus', 'nminus', 'lb', 'ub', 'alpha', cls=NumberList)
+        prev, allowed, magnitudes = unpack_states(
+            states, tensors,
             'prev','allowed','magnitudes',
-            params=params,
             init=[torch.zeros_like, _bool_ones_like, torch.zeros_like],
             cls = TensorList,
         )
 
-        target = rprop_(
+        tensors = rprop_(
             tensors_ = as_tensorlist(tensors),
             prev_ = prev,
             allowed_ = allowed,
@@ -184,12 +187,11 @@ class Rprop(Transform):
             lb = lb,
             ub = ub,
             alpha = alpha,
-            backtrack=self.settings[params[0]]['backtrack'],
-            step=self.current_step,
+            backtrack=settings[0]['backtrack'],
+            step=step,
         )
 
-        self.current_step += 1
-        return target
+        return tensors
 
 
 class ScaleLRBySignChange(Transform):
@@ -220,23 +222,25 @@ class ScaleLRBySignChange(Transform):
     ):
         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']
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        tensors = as_tensorlist(tensors)
+        use_grad = settings[0]['use_grad']
         if use_grad: cur = as_tensorlist(grads)
-        else: cur = target
+        else: cur = 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)
+        nplus, nminus, lb, ub = unpack_dicts(settings, 'nplus', 'nminus', 'lb', 'ub', cls=NumberList)
+        prev, lrs = unpack_states(states, tensors, 'prev', 'lrs', cls=TensorList)
 
-        if self.current_step == 0:
-            lrs.set_(target.full_like(self.get_settings('alpha', params=params)))
+        if step == 0:
+            lrs.set_(tensors.full_like([s['alpha'] for s in settings]))
 
-        target = scale_by_sign_change_(
-            tensors_ = target,
+        tensors = scale_by_sign_change_(
+            tensors_ = tensors,
             cur = cur,
             prev_ = prev,
             lrs_ = lrs,
@@ -244,10 +248,9 @@ class ScaleLRBySignChange(Transform):
             nminus = nminus,
             lb = lb,
             ub = ub,
-            step = self.current_step,
+            step = step,
         )
-        self.current_step += 1
-        return target
+        return tensors
 
 class BacktrackOnSignChange(Transform):
     """Negates or undoes update for parameters where where gradient or update sign changes.
@@ -268,44 +271,77 @@ class BacktrackOnSignChange(Transform):
     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']
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        tensors = as_tensorlist(tensors)
+        use_grad = settings[0]['use_grad']
+        backtrack = settings[0]['backtrack']
 
         if use_grad: cur = as_tensorlist(grads)
-        else: cur = target
+        else: cur = tensors
 
-        target = backtrack_on_sign_change_(
-            tensors_ = target,
+        tensors = backtrack_on_sign_change_(
+            tensors_ = tensors,
             cur = cur,
-            prev_ = self.get_state('prev', params=params, cls=TensorList),
+            prev_ = unpack_states(states, tensors, 'prev', cls=TensorList),
             backtrack = backtrack,
-            step = self.current_step,
+            step = step,
         )
 
-        self.current_step += 1
-        return target
+        return tensors
 
 class SignConsistencyMask(Transform):
-    """0 if sign changed 1 otherwise"""
+    """
+    Outputs a mask of sign consistency of current and previous inputs.
+
+    The output is 0 for weights where input sign changed compared to previous input, 1 otherwise.
+
+    Examples:
+
+        GD that skips update for weights where gradient sign changed compared to previous gradient.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Mul(tz.m.SignConsistencyMask()),
+                tz.m.LR(1e-2)
+            )
+
+    """
     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)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', cls=TensorList)
         mask = prev.mul_(tensors).gt_(0)
-        prev.set_(tensors)
+        prev.copy_(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."""
+    """Outputs per-weight learning rates based on consecutive sign consistency.
+
+    The learning rate for a weight is multiplied by :code:`nplus` when two consecutive update signs are the same, otherwise it is multiplied by :code:`nplus`. The learning rates are bounded to be in :code:`(lb, ub)` range.
+
+    Examples:
+
+        GD scaled by consecutive gradient sign consistency
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Mul(tz.m.SignConsistencyLRs()),
+                tz.m.LR(1e-2)
+            )
+
+    """
     def __init__(
         self,
         nplus: float = 1.2,
@@ -317,16 +353,18 @@ class SignConsistencyLRs(Transform):
     ):
         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):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
         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)
+        nplus, nminus, lb, ub = unpack_dicts(settings, 'nplus', 'nminus', 'lb', 'ub', cls=NumberList)
+        prev, lrs = unpack_states(states, tensors, 'prev', 'lrs', cls=TensorList)
 
-        if self.current_step == 0:
-            lrs.set_(target.full_like(self.get_settings('alpha', params=params)))
+        if step == 0:
+            lrs.set_(target.full_like([s['alpha'] for s in settings]))
 
         target = sign_consistency_lrs_(
             tensors = target,
@@ -336,7 +374,6 @@ class SignConsistencyLRs(Transform):
             nminus = nminus,
             lb = lb,
             ub = ub,
-            step = self.current_step,
+            step = step,
         )
-        self.current_step += 1
         return target.clone()

torchzero/modules/optimizers/sam.py

@@ -0,0 +1,163 @@
+from contextlib import nullcontext
+import torch
+from ...utils import TensorList, NumberList
+from ...core import Module
+
+
+class SAM(Module):
+    """Sharpness-Aware Minimization from https://arxiv.org/pdf/2010.01412
+
+    SAM functions by seeking parameters that lie in neighborhoods having uniformly low loss value.
+    It performs two forward and backward passes per step.
+
+    This implementation modifies the closure to return loss and calculate gradients
+    of the SAM objective. All modules after this will use the modified objective.
+
+    .. note::
+        This module requires a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients at two points on each step.
+
+    Args:
+        rho (float, optional): Neighborhood size. Defaults to 0.05.
+        p (float, optional): norm of the SAM objective. Defaults to 2.
+        asam (bool, optional):
+            enables ASAM variant which makes perturbation relative to weight magnitudes.
+            ASAM requires a much larger :code:`rho`, like 0.5 or 1.
+            The :code:`tz.m.ASAM` class is idential to setting this argument to True, but
+            it has larger :code:`rho` by default.
+
+    Examples:
+        SAM-SGD:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.SAM(),
+                tz.m.LR(1e-2)
+            )
+
+        SAM-Adam:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.SAM(),
+                tz.m.Adam(),
+                tz.m.LR(1e-2)
+            )
+
+    References:
+        Foret, P., Kleiner, A., Mobahi, H., & Neyshabur, B. (2020). Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412. https://arxiv.org/abs/2010.01412#page=3.16
+    """
+    def __init__(self, rho: float = 0.05, p: float = 2, eps=1e-10, asam=False):
+        defaults = dict(rho=rho, p=p, eps=eps, asam=asam)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+
+        params = var.params
+        closure = var.closure
+        zero_grad = var.zero_grad
+        if closure is None: raise RuntimeError("SAM requires a closure passed to the optimizer step")
+        p, rho = self.get_settings(var.params, 'p', 'rho', cls=NumberList)
+        s = self.settings[var.params[0]]
+        eps = s['eps']
+        asam = s['asam']
+
+        # 1/p + 1/q = 1
+        # okay, authors of SAM paper, I will manually solve your equation
+        # so q = -p/(1-p)
+        q = -p / (1-p)
+        # as a validation for 2 it is -2 / -1 = 2
+
+        @torch.no_grad
+        def sam_closure(backward=True):
+            orig_grads = None
+            if not backward:
+                # if backward is False, make sure this doesn't modify gradients
+                # to avoid issues
+                orig_grads = [p.grad for p in params]
+
+            # gradient at initial parameters
+            zero_grad()
+            with torch.enable_grad():
+                closure()
+
+            grad = TensorList(p.grad if p.grad is not None else torch.zeros_like(p) for p in params)
+            grad_abs = grad.abs()
+
+            # compute e
+            term1 = grad.sign().mul_(rho)
+            term2 = grad_abs.pow(q-1)
+
+            if asam:
+                grad_abs.mul_(torch._foreach_abs(params))
+
+            denom = grad_abs.pow_(q).sum().pow(1/p)
+
+            e = term1.mul_(term2).div_(denom.clip(min=eps))
+
+            if asam:
+                e.mul_(torch._foreach_pow(params, 2))
+
+            # calculate loss and gradient approximation of inner problem
+            torch._foreach_add_(params, e)
+            if backward:
+                zero_grad()
+                with torch.enable_grad():
+                    # this sets .grad attributes
+                    sam_loss = closure()
+
+            else:
+                sam_loss = closure(False)
+
+            # and restore initial parameters
+            torch._foreach_sub_(params, e)
+
+            if orig_grads is not None:
+                for param,orig_grad in zip(params, orig_grads):
+                    param.grad = orig_grad
+
+            return sam_loss
+
+        var.closure = sam_closure
+        return var
+
+# different class because defaults for SAM are bad for ASAM
+class ASAM(SAM):
+    """Adaptive Sharpness-Aware Minimization from https://arxiv.org/pdf/2102.11600#page=6.52
+
+    SAM functions by seeking parameters that lie in neighborhoods having uniformly low loss value.
+    It performs two forward and backward passes per step.
+
+    This implementation modifies the closure to return loss and calculate gradients
+    of the SAM objective. All modules after this will use the modified objective.
+
+    .. note::
+        This module requires a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients at two points on each step.
+
+    Args:
+        rho (float, optional): Neighborhood size. Defaults to 0.05.
+        p (float, optional): norm of the SAM objective. Defaults to 2.
+
+    Examples:
+        ASAM-Adam:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ASAM(),
+                tz.m.Adam(),
+                tz.m.LR(1e-2)
+            )
+
+    References:
+        Kwon, J., Kim, J., Park, H., & Choi, I. K. (2021, July). Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning (pp. 5905-5914). PMLR. https://arxiv.org/abs/2102.11600
+    """
+    def __init__(self, rho: float = 0.5, p: float = 2, eps=1e-10):
+        super().__init__(rho=rho, p=p, eps=eps, asam=True)

torchzero/modules/optimizers/shampoo.py

@@ -4,7 +4,7 @@ from functools import partial
 import numpy as np
 import torch
 
-from ...core import Chainable, Transform, apply
+from ...core import Chainable, Transform, apply_transform
 from ...utils.linalg import matrix_power_eigh
 from ...utils import set_storage_
 
@@ -59,7 +59,7 @@ def _merge_small_dims(tensor: torch.Tensor, max_dim: int):
     if tensor.shape[sort_idxs[0]] > max_dim:
         return tensor, None, None
 
-    tensor = tensor.permute(*sort_idxs)
+    tensor = tensor.permute(*sort_idxs.tolist())
     flatten_end_idx = 0
     flat_sizes = []
     flat_numel = 1
@@ -80,19 +80,28 @@ def _unmerge_small_dims(tensor: torch.Tensor, flat_sizes: 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))
+    return tensor.permute(*np.argsort(sort_idxs).tolist())
 
 
 class Shampoo(Transform):
     """Shampoo from Preconditioned Stochastic Tensor Optimization (https://arxiv.org/abs/1802.09568).
 
+    .. note::
+        Shampoo is usually grafted to another optimizer like Adam, otherwise it can be unstable. An example of how to do grafting is given below in the Examples section.
+
+    .. note::
+        Shampoo is a very computationally expensive optimizer, increase :code:`update_freq` if it is too slow.
+
+    .. note::
+        SOAP optimizer usually outperforms Shampoo and is also not as computationally expensive. SOAP implementation is available as :code:`tz.m.SOAP`.
+
     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.
+        exp_override (int | None, optional): matrix exponent override, if not set, uses 2*ndim. Defaults to 2.
         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.
@@ -101,39 +110,62 @@ class Shampoo(Transform):
             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.
+
+    Examples:
+        Shampoo grafted to Adam
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.GraftModules(
+                    direction = tz.m.Shampoo(),
+                    magnitude = tz.m.Adam(),
+                ),
+                tz.m.LR(1e-3)
+            )
+
+        Adam with Shampoo preconditioner
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Shampoo(beta=0.999, inner=tz.m.EMA(0.9)),
+                tz.m.Debias(0.9, 0.999),
+                tz.m.LR(1e-3)
+            )
     """
     def __init__(
         self,
         decay: float | None = None,
         beta: float | None = None,
-        reg: float = 1e-6,
         update_freq: int = 10,
-        exp_override: int | None = None,
+        exp_override: int | None = 2,
         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)
+        defaults = dict(decay=decay, beta=beta, 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
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        merged_tensors = [] # 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)
+        for i,(t,state, setting) in enumerate(zip(tensors, states, settings)):
+            beta, update_freq, exp_override, merge_small, max_dim, precondition_1d = itemgetter(
+                'beta', 'update_freq', 'exp_override', 'merge_small', 'max_dim', 'precondition_1d')(setting)
 
             if merge_small:
                 t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
-            merged_target.append(t)
+
+            merged_tensors.append(t)
 
             # initialize accumulators and preconditioners for each dim on 1st step
             if 'accumulators' not in state:
@@ -167,22 +199,18 @@ class Shampoo(Transform):
 
         # inner step
         if 'inner' in self.children:
-            tensors = apply(self.children['inner'], tensors, params=params, grads=grads, vars=vars)
+            tensors = apply_transform(self.children['inner'], tensors, params=params, grads=grads)
 
             # 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)
+            merged_tensors = [] # target with merged dims
+            for i,(t,state, setting) in enumerate(zip(tensors, states, settings)):
+                if setting['merge_small']:
+                    t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, setting['max_dim'])
+                merged_tensors.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)
+        for i,(t,state, setting) in enumerate(zip(merged_tensors, states, settings)):
+            decay, merge_small, adagrad_eps= itemgetter('decay', 'merge_small', 'adagrad_eps')(setting)
 
             if 'diagonal_accumulator' in state:
                 tensors[i] = apply_diagonal_(t, state['diagonal_accumulator'], decay=decay, eps=adagrad_eps)