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

torchzero/modules/misc/multistep.py

@@ -0,0 +1,194 @@
+from collections.abc import Iterable
+
+import torch
+
+from ...core import Chainable, Module, Var
+from ...utils import TensorList
+
+def _sequential_step(self: Module, var: Var, sequential: bool):
+    params = var.params
+    steps = self.settings[params[0]]['steps']
+
+    if sequential: modules = self.get_children_sequence() * steps
+    else: modules = [self.children['module']] * steps
+
+    if var.closure is None and len(modules) > 1: raise ValueError('Multistep and Sequential require closure')
+
+    # store original params unless this is last module and can update params directly
+    params_before_steps = None if (var.is_last and var.last_module_lrs is None) else [p.clone() for p in params]
+
+    # first step - pass var as usual
+    var = modules[0].step(var)
+    new_var = var
+
+    # subsequent steps - update parameters and create new var
+    if len(modules) > 1:
+        for m in modules[1:]:
+
+            # update params
+            if (not new_var.skip_update):
+                if new_var.last_module_lrs is not None:
+                    torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
+
+                torch._foreach_sub_(params, new_var.get_update())
+
+            # create new var since we are at a new point, that means grad, update and loss will be None
+            new_var = Var(params=new_var.params, closure=new_var.closure,
+                            model=new_var.model, current_step=new_var.current_step + 1)
+
+            # step
+            new_var = m.step(new_var)
+
+        # final parameter update
+        if (not new_var.skip_update):
+            if new_var.last_module_lrs is not None:
+                torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
+
+            torch._foreach_sub_(params, new_var.get_update())
+
+    # if last module, update is applied so return new var
+    if params_before_steps is None:
+        new_var.stop = True
+        new_var.skip_update = True
+        return new_var
+
+    # otherwise use parameter difference as update
+    var.update = list(torch._foreach_sub(params_before_steps, params))
+    for p, bef in zip(params, params_before_steps):
+        p.set_(bef) # pyright:ignore[reportArgumentType]
+    return var
+
+class Multistep(Module):
+    """Performs :code:`steps` inner steps with :code:`module` per each step.
+
+    The update is taken to be the parameter difference between parameters before and after the inner loop."""
+    def __init__(self, module: Chainable, steps: int):
+        defaults = dict(steps=steps)
+        super().__init__(defaults)
+        self.set_child('module', module)
+
+    @torch.no_grad
+    def step(self, var):
+        return _sequential_step(self, var, sequential=False)
+
+class Sequential(Module):
+    """On each step, this sequentially steps with :code:`modules` :code:`steps` times.
+
+    The update is taken to be the parameter difference between parameters before and after the inner loop."""
+    def __init__(self, modules: Iterable[Chainable], steps: int=1):
+        defaults = dict(steps=steps)
+        super().__init__(defaults)
+        self.set_children_sequence(modules)
+
+    @torch.no_grad
+    def step(self, var):
+        return _sequential_step(self, var, sequential=True)
+
+
+class NegateOnLossIncrease(Module):
+    """Uses an extra forward pass to evaluate loss at :code:`parameters+update`,
+    if loss is larger than at :code:`parameters`,
+    the update is set to 0 if :code:`backtrack=False` and to :code:`-update` otherwise"""
+    def __init__(self, backtrack=False):
+        defaults = dict(backtrack=backtrack)
+        super().__init__(defaults=defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+        if closure is None: raise RuntimeError('NegateOnLossIncrease requires closure')
+        backtrack = self.defaults['backtrack']
+
+        update = var.get_update()
+        f_0 = var.get_loss(backward=False)
+
+        torch._foreach_sub_(var.params, update)
+        f_1 = closure(False)
+
+        if f_1 <= f_0:
+            if var.is_last and var.last_module_lrs is None:
+                var.stop = True
+                var.skip_update = True
+                return var
+
+            torch._foreach_add_(var.params, update)
+            return var
+
+        torch._foreach_add_(var.params, update)
+        if backtrack:
+            torch._foreach_neg_(var.update)
+        else:
+            torch._foreach_zero_(var.update)
+        return var
+
+
+class Online(Module):
+    """Allows certain modules to be used for mini-batch optimization.
+
+    Examples:
+
+    Online L-BFGS with Backtracking line search
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Online(tz.m.LBFGS()),
+        tz.m.Backtracking()
+    )
+    ```
+
+    Online L-BFGS trust region
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.TrustCG(tz.m.Online(tz.m.LBFGS()))
+    )
+    ```
+
+    """
+    def __init__(self, *modules: Module,):
+        super().__init__()
+
+        self.set_child('module', modules)
+
+    @torch.no_grad
+    def update(self, var):
+        closure = var.closure
+        if closure is None: raise ValueError("Closure must be passed for Online")
+
+        step = self.global_state.get('step', 0) + 1
+        self.global_state['step'] = step
+
+        params = TensorList(var.params)
+        p_cur = params.clone()
+        p_prev = self.get_state(params, 'p_prev', cls=TensorList)
+
+        module = self.children['module']
+        var_c = var.clone(clone_update=False)
+
+        # on 1st step just step and store previous params
+        if step == 1:
+            p_prev.copy_(params)
+
+            module.update(var_c)
+            var.update_attrs_from_clone_(var_c)
+            return
+
+        # restore previous params and update
+        var_prev = Var(params=params, closure=closure, model=var.model, current_step=var.current_step)
+        params.set_(p_prev)
+        module.reset_for_online()
+        module.update(var_prev)
+
+        # restore current params and update
+        params.set_(p_cur)
+        p_prev.copy_(params)
+        module.update(var_c)
+        var.update_attrs_from_clone_(var_c)
+
+    @torch.no_grad
+    def apply(self, var):
+        module = self.children['module']
+        return module.apply(var.clone(clone_update=False))
+
+    def get_H(self, var):
+        return self.children['module'].get_H(var)

torchzero/modules/misc/regularization.py

@@ -0,0 +1,167 @@
+import torch
+
+from ...core import Chainable, Module, Target, Transform
+from ...core.reformulation import Reformulation
+from ...utils import Distributions, NumberList, TensorList
+
+
+class Dropout(Transform):
+    """Applies dropout to the update.
+
+    For each weight the update to that weight has :code:`p` probability to be set to 0.
+    This can be used to implement gradient dropout or update dropout depending on placement.
+
+    Args:
+        p (float, optional): probability that update for a weight is replaced with 0. Defaults to 0.5.
+        graft (bool, optional):
+            if True, update after dropout is rescaled to have the same norm as before dropout. Defaults to False.
+        target (Target, optional): what to set on var, refer to documentation. Defaults to 'update'.
+
+
+    Examples:
+        Gradient dropout.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Dropout(0.5),
+                tz.m.Adam(),
+                tz.m.LR(1e-3)
+            )
+
+        Update dropout.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.Dropout(0.5),
+                tz.m.LR(1e-3)
+            )
+
+    """
+    def __init__(self, p: float = 0.5, graft: bool=False, target: Target = 'update'):
+        defaults = dict(p=p, graft=graft)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        p = NumberList(s['p'] for s in settings)
+        graft = settings[0]['graft']
+
+        if graft:
+            target_norm = tensors.global_vector_norm()
+            tensors.mul_(tensors.rademacher_like(1-p).add_(1).div_(2))
+            return tensors.mul_(target_norm / tensors.global_vector_norm()) # graft
+
+        return tensors.mul_(tensors.rademacher_like(1-p).add_(1).div_(2))
+
+def _bernoulli_like(tensor, p = 0.5, generator = None):
+    """p is probability of a 1, other values will be 0."""
+    return torch.bernoulli(torch.full_like(tensor, p), generator = generator)
+
+class WeightDropout(Module):
+    """
+    Changes the closure so that it evaluates loss and gradients with random weights replaced with 0.
+
+    Dropout can be disabled for a parameter by setting :code:`use_dropout=False` in corresponding parameter group.
+
+    Args:
+        p (float, optional): probability that any weight is replaced with 0. Defaults to 0.5.
+        graft (bool, optional):
+            if True, parameters after dropout are rescaled to have the same norm as before dropout. Defaults to False.
+    """
+    def __init__(self, p: float = 0.5, graft: bool = True):
+        defaults = dict(p=p, graft=graft, use_dropout=True)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+        if closure is None: raise RuntimeError('WeightDropout requires closure')
+        params = TensorList(var.params)
+        p = NumberList(self.settings[p]['p'] for p in params)
+
+        # create masks
+        mask = []
+        for p, m in zip(params, mask):
+            prob = self.settings[p]['p']
+            use_dropout = self.settings[p]['use_dropout']
+            if use_dropout: mask.append(_bernoulli_like(p, prob))
+            else: mask.append(torch.ones_like(p))
+
+        @torch.no_grad
+        def dropout_closure(backward=True):
+            orig_params = params.clone()
+            params.mul_(mask)
+            if backward:
+                with torch.enable_grad(): loss = closure()
+            else:
+                loss = closure(False)
+            params.copy_(orig_params)
+            return loss
+
+        var.closure = dropout_closure
+        return var
+
+
+class PerturbWeights(Module):
+    """
+    Changes the closure so that it evaluates loss and gradients at weights perturbed by a random perturbation.
+
+    Can be disabled for a parameter by setting :code:`perturb=False` in corresponding parameter group.
+
+    Args:
+        alpha (float, optional): multiplier for perturbation magnitude. Defaults to 0.1.
+        relative (bool, optional): whether to multiply perturbation by mean absolute value of the parameter. Defaults to True.
+        distribution (bool, optional):
+            distribution of the random perturbation. Defaults to False.
+    """
+    def __init__(self, alpha: float = 0.1, relative:bool=True, distribution:Distributions = 'normal'):
+        defaults = dict(alpha=alpha, relative=relative, distribution=distribution, perturb=True)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+        if closure is None: raise RuntimeError('WeightDropout requires closure')
+        params = TensorList(var.params)
+
+        # create perturbations
+        perts = []
+        for p in params:
+            settings = self.settings[p]
+            if not settings['perturb']:
+                perts.append(torch.zeros_like(p))
+                continue
+
+            alpha = settings['alpha']
+            if settings['relative']:
+                alpha *= p.abs().mean()
+
+            distribution = self.settings[p]['distribution'].lower()
+            if distribution in ('normal', 'gaussian'):
+                perts.append(torch.randn_like(p).mul_(alpha))
+            elif distribution == 'uniform':
+                perts.append(torch.empty_like(p).uniform_(-alpha,alpha))
+            elif distribution == 'sphere':
+                r = torch.randn_like(p)
+                perts.append((r * alpha) / torch.linalg.vector_norm(r)) # pylint:disable=not-callable
+            else:
+                raise ValueError(distribution)
+
+        @torch.no_grad
+        def perturbed_closure(backward=True):
+            params.add_(perts)
+            if backward:
+                with torch.enable_grad(): loss = closure()
+            else:
+                loss = closure(False)
+            params.sub_(perts)
+            return loss
+
+        var.closure = perturbed_closure
+        return var

torchzero/modules/misc/split.py

@@ -0,0 +1,123 @@
+import warnings
+from collections.abc import Callable, Sequence, Iterable
+from typing import cast
+
+import torch
+
+from ...core import Chainable, Module, Var
+
+
+def _split(
+    module: Module,
+    idxs,
+    params,
+    var: Var,
+):
+    split_params = [p for i,p in enumerate(params) if i in idxs]
+
+    split_grad = None
+    if var.grad is not None:
+        split_grad = [g for i,g in enumerate(var.grad) if i in idxs]
+
+    split_update = None
+    if var.update is not None:
+        split_update = [u for i,u in enumerate(var.update) if i in idxs]
+
+    split_var = var.clone(clone_update=False, parent=var)
+    split_var.params = split_params
+    split_var.grad = split_grad
+    split_var.update = split_update
+
+    split_var = module.step(split_var)
+
+    # those should be set due to var being parent
+    if split_var.grad is not None:
+        assert var.grad is not None
+
+    if split_var.loss is not None:
+        assert var.loss is not None
+
+    if split_var.update is not None:
+
+        # make sure update is set, it will be filled with ``true`` and ``false`` tensors
+        if var.update is None:
+            if var.grad is None: var.update = [cast(torch.Tensor, None) for _ in var.params]
+            else: var.update = [g.clone() for g in var.grad]
+
+        # set all tensors from this split
+        for idx, u in zip(idxs, split_var.update):
+            var.update[idx] = u
+
+    return var
+
+_SingleFilter = Callable[[torch.Tensor], bool] | torch.Tensor | Iterable[torch.Tensor] | torch.nn.Module | Iterable[torch.nn.Module]
+Filter = _SingleFilter | Iterable[_SingleFilter]
+
+def _make_filter(filter: Filter):
+    if callable(filter): return filter
+    if isinstance(filter, torch.Tensor):
+        return lambda x: x is filter
+    if isinstance(filter, torch.nn.Module):
+        return _make_filter(filter.parameters())
+
+    # iterable
+    filters = [_make_filter(f) for f in filter]
+    return lambda x: any(f(x) for f in filters)
+
+class Split(Module):
+    """Apply ``true`` modules to all parameters filtered by ``filter``, apply ``false`` modules to all other parameters.
+
+    Args:
+        filter (Filter, bool]):
+            a filter that selects tensors to be optimized by ``true``.
+            - tensor or iterable of tensors (e.g. ``encoder.parameters()``).
+            - function that takes in tensor and outputs a bool (e.g. ``lambda x: x.ndim >= 2``).
+            - a sequence of above (acts as "or", so returns true if any of them is true).
+
+        true (Chainable | None): modules that are applied to tensors where ``filter`` is ``True``.
+        false (Chainable | None): modules that are applied to tensors where ``filter`` is ``False``.
+
+    ### Examples:
+
+    Muon with Adam fallback using same hyperparams as https://github.com/KellerJordan/Muon
+
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NAG(0.95),
+        tz.m.Split(
+            lambda p: p.ndim >= 2,
+            true = tz.m.Orthogonalize(),
+            false = [tz.m.Adam(0.9, 0.95), tz.m.Mul(1/66)],
+        ),
+        tz.m.LR(1e-2),
+    )
+    ```
+    """
+    def __init__(self, filter: Filter, true: Chainable | None, false: Chainable | None):
+        defaults = dict(filter=filter)
+        super().__init__(defaults)
+
+        if true is not None: self.set_child('true', true)
+        if false is not None: self.set_child('false', false)
+
+    def step(self, var):
+
+        params = var.params
+        filter = _make_filter(self.settings[params[0]]['filter'])
+
+        true_idxs = []
+        false_idxs = []
+        for i,p in enumerate(params):
+            if filter(p): true_idxs.append(i)
+            else: false_idxs.append(i)
+
+        if 'true' in self.children and len(true_idxs) > 0:
+            true = self.children['true']
+            var = _split(true, idxs=true_idxs, params=params, var=var)
+
+        if 'false' in self.children and len(false_idxs) > 0:
+            false = self.children['false']
+            var = _split(false, idxs=false_idxs, params=params, var=var)
+
+        return var

torchzero/modules/{ops → misc}/switch.py

@@ -7,7 +7,28 @@ from ...core import Chainable, Module
 
 
 class Alternate(Module):
-    """alternate between stepping with `modules`"""
+    """Alternates between stepping with :code:`modules`.
+
+    That is, first step is performed with 1st module, second step with second module, etc.
+
+    Args:
+        steps (int | Iterable[int], optional): number of steps to perform with each module. Defaults to 1.
+
+    Examples:
+        Alternate between Adam, SignSGD and RMSprop
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Alternate(
+                    tz.m.Adam(),
+                    [tz.m.SignSGD(), tz.m.Mul(0.5)],
+                    tz.m.RMSprop(),
+                ),
+                tz.m.LR(1e-3),
+            )
+    """
     LOOP = True
     def __init__(self, *modules: Chainable, steps: int | Iterable[int] = 1):
         if isinstance(steps, Iterable):
@@ -32,7 +53,7 @@ class Alternate(Module):
         var = module.step(var.clone(clone_update=False))
 
         # number of steps until next module
-        steps = self.settings[var.params[0]]['steps']
+        steps = self.defaults['steps']
         if isinstance(steps, int): steps = [steps]*len(self.children)
 
         if 'steps_to_next' not in self.global_state:
@@ -54,14 +75,34 @@ class Alternate(Module):
         return var
 
 class Switch(Alternate):
-    """switch to next module after some steps"""
+    """After :code:`steps` steps switches to the next module.
+
+    Args:
+        steps (int | Iterable[int]): Number of steps to perform with each module.
+
+    Examples:
+        Start with Adam, switch to L-BFGS after 1000th step and Truncated Newton on 2000th step.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Switch(
+                    [tz.m.Adam(), tz.m.LR(1e-3)],
+                    [tz.m.LBFGS(), tz.m.Backtracking()],
+                    [tz.m.NewtonCG(maxiter=20), tz.m.Backtracking()],
+                    steps = (1000, 2000)
+                )
+            )
+    """
+
     LOOP = False
     def __init__(self, *modules: Chainable, steps: int | Iterable[int]):
 
         if isinstance(steps, Iterable):
             steps = list(steps)
             if len(steps) != len(modules) - 1:
-                raise ValueError(f"steps must be the same length as modules, got {len(modules) = }, {len(steps) = }")
+                raise ValueError(f"steps must be the same length as modules minus 1, got {len(modules) = }, {len(steps) = }")
 
             steps.append(1)
 

torchzero/modules/momentum/__init__.py

@@ -6,9 +6,5 @@ from .cautious import (
     ScaleModulesByCosineSimilarity,
     UpdateGradientSignConsistency,
 )
-from .ema import EMA, Debias, Debias2, EMASquared, SqrtEMASquared, CenteredEMASquared, CenteredSqrtEMASquared
-from .experimental import CoordinateMomentum
-# from .matrix_momentum import MatrixMomentum
 
-from .momentum import NAG, HeavyBall
-from .matrix_momentum import MatrixMomentum, AdaptiveMatrixMomentum
+from .momentum import NAG, HeavyBall, EMA

torchzero/modules/momentum/averaging.py

@@ -10,7 +10,7 @@ from ...utils import tolist
 
 
 class Averaging(TensorwiseTransform):
-    """Average of past :code:`history_size` updates.
+    """Average of past ``history_size`` updates.
 
     Args:
         history_size (int): Number of past updates to average
@@ -21,8 +21,8 @@ class Averaging(TensorwiseTransform):
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
-        history_size = settings['history_size']
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        history_size = setting['history_size']
         if 'history' not in state:
             state['history'] = deque(maxlen=history_size)
             state['average'] = torch.zeros_like(tensor)
@@ -35,7 +35,7 @@ class Averaging(TensorwiseTransform):
         return average / len(history)
 
 class WeightedAveraging(TensorwiseTransform):
-    """Weighted average of past :code:`len(weights)` updates.
+    """Weighted average of past ``len(weights)`` updates.
 
     Args:
         weights (Sequence[float]): a sequence of weights from oldest to newest.
@@ -46,8 +46,8 @@ class WeightedAveraging(TensorwiseTransform):
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
-        weights = settings['weights']
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        weights = setting['weights']
 
         if 'history' not in state:
             state['history'] = deque(maxlen=len(weights))
@@ -69,7 +69,7 @@ class WeightedAveraging(TensorwiseTransform):
 
 
 class MedianAveraging(TensorwiseTransform):
-    """Median of past :code:`history_size` updates.
+    """Median of past ``history_size`` updates.
 
     Args:
         history_size (int): Number of past updates to average
@@ -80,8 +80,8 @@ class MedianAveraging(TensorwiseTransform):
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
-        history_size = settings['history_size']
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        history_size = setting['history_size']
 
         if 'history' not in state:
             state['history'] = deque(maxlen=history_size)