torchzero 0.0.1__py3-none-any.whl

torchzero/modules/meta/return_overrides.py

@@ -0,0 +1,46 @@
+import torch
+from ...tensorlist import TensorList
+from ...core import OptimizerModule, _get_loss, _ClosureType
+
+class SetGrad(OptimizerModule):
+    """Doesn't update parameters, instead replaces all parameters `.grad` attribute with the current update.
+    You can now step with any pytorch optimizer that utilises the `.grad` attribute."""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def step(self, state):
+        if self.next_module is not None: raise ValueError("SetGrad can't have children")
+        params = self.get_params()
+        g = state.maybe_use_grad_(params) # this may execute the closure which might be modified
+        params.set_grad_(g)
+        return state.get_loss()
+
+
+class ReturnAscent(OptimizerModule):
+    """Doesn't update parameters, instead returns the update as a TensorList of tensors."""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def step(self, state) -> TensorList: # type:ignore
+        if self.next_module is not None: raise ValueError("ReturnAscent can't have children")
+        params = self.get_params()
+        update = state.maybe_use_grad_(params) # this will execute the closure which might be modified
+        return update
+
+class ReturnClosure(OptimizerModule):
+    """Doesn't update parameters, instead returns the current modified closure.
+    For example, if you put this after :code:`torchzero.modules.FDM(target = "closure")`,
+    the closure will set `.grad` attribute to gradients approximated via finite difference.
+    You can then pass that closure to something that requires closure like `torch.optim.LBFGS`."""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def step(self, state) -> _ClosureType: # type:ignore
+        if self.next_module is not None: raise ValueError("ReturnClosure can't have children")
+        if state.closure is None:
+            raise ValueError("MakeClosure requires closure")
+        return state.closure
+

torchzero/modules/misc/__init__.py

@@ -0,0 +1,10 @@
+r"""
+This module includes various basic operators, notable LR for setting the learning rate,
+as well as gradient/update clipping and normalization.
+"""
+
+from .basic import Clone, Fill, Grad, Identity, Lambda, Zeros, Alpha, GradToUpdate, MakeClosure
+from .lr import LR
+from .on_increase import NegateOnLossIncrease
+from .multistep import Multistep
+from .accumulate import Accumulate

torchzero/modules/misc/accumulate.py

@@ -0,0 +1,43 @@
+from collections.abc import Callable, Iterable
+
+import torch
+
+from ...tensorlist import TensorList
+
+from ...core import OptimizerModule
+
+
+class Accumulate(OptimizerModule):
+    """Accumulates update over n steps, and steps once updates have been accumulated.
+    Put this as the first module to get gradient accumulation.
+
+    Args:
+        n_steps (int): number of steps (batches) to accumulate the update over.
+        mean (bool, optional):
+            If True, divides accumulated gradients by number of step,
+            since most loss functions calculate the mean of all samples
+            over batch. Defaults to True.
+    """
+    def __init__(self, n_steps: int, mean = True):
+
+        super().__init__({})
+        self.n_steps = n_steps
+        self.mean = mean
+        self.cur_step = 0
+
+    @torch.no_grad
+    def step(self, state):
+        self.cur_step += 1
+
+        params = self.get_params()
+        accumulated_update = self.get_state_key('accumulated_grads')
+        accumulated_update += state.maybe_use_grad_(params)
+
+        if self.cur_step % self.n_steps == 0:
+            state.ascent = accumulated_update.clone()
+            if self.mean: state.ascent /= self.n_steps
+            accumulated_update.zero_()
+            return self._update_params_or_step_with_next(state)
+
+
+        return state.get_loss()

torchzero/modules/misc/basic.py

@@ -0,0 +1,115 @@
+from collections.abc import Callable, Iterable
+
+import torch
+
+from ...tensorlist import TensorList
+
+from ...core import OptimizerModule, _Chainable
+
+
+class Alpha(OptimizerModule):
+    """Multiplies update by the learning rate, won't get picked up by learning rate schedulers."""
+    def __init__(self, alpha = 1e-3):
+        defaults = dict(alpha = alpha)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        # multiply ascent direction by lr in-place
+        lr = self.get_group_key('alpha')
+        ascent *= lr
+        return ascent
+
+class Clone(OptimizerModule):
+    """Clones the update. Some modules update ascent in-place, so this may be
+    useful if you need to preserve it."""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def _update(self, state, ascent): return ascent.clone()
+
+class Identity(OptimizerModule):
+    """Does nothing."""
+    def __init__(self, *args, **kwargs):
+        super().__init__({})
+
+    @torch.no_grad
+    def _update(self, state, ascent): return ascent
+
+class Lambda(OptimizerModule):
+    """Applies a function to the ascent direction.
+    The function must take a TensorList as the argument, and return the modified tensorlist.
+
+    Args:
+        f (Callable): function
+    """
+    def __init__(self, f: Callable[[TensorList], TensorList]):
+        super().__init__({})
+        self.f = f
+
+    @torch.no_grad()
+    def _update(self, state, ascent): return self.f(ascent)
+
+class Grad(OptimizerModule):
+    """Uses gradient as the update. This is useful for chains."""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        ascent = state.ascent = state.maybe_compute_grad_(self.get_params())
+        return ascent
+
+class Zeros(OptimizerModule):
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        return ascent.zeros_like()
+
+class Fill(OptimizerModule):
+    def __init__(self, value):
+        super().__init__({"value": value})
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        return ascent.fill(self.get_group_key('value'))
+
+
+class GradToUpdate(OptimizerModule):
+    """sets gradient and .grad attributes to current update"""
+    def __init__(self):
+        super().__init__({})
+
+    def _update(self, state, ascent):
+        state.set_grad_(ascent, self.get_params())
+        return ascent
+
+class MakeClosure(OptimizerModule):
+    """Makes a closure that sets `.grad` attribute to the update generated by `modules`"""
+    def __init__(self, modules: _Chainable):
+        super().__init__({})
+        self._set_child_('modules', modules)
+
+    def step(self, state):
+        if state.closure is None: raise ValueError("MakeClosure requires a closure")
+
+        params = self.get_params()
+        orig_closure = state.closure
+        orig_state = state.copy(True)
+
+        def new_closure(backward = True):
+            if backward:
+                cloned_state = orig_state.copy(True)
+                g = self.children['modules'].return_ascent(cloned_state)
+                params.set_grad_(g)
+                return cloned_state.get_loss()
+
+            else:
+                return orig_closure(False)
+
+        state.closure = new_closure # type:ignore
+        return self._update_params_or_step_with_next(state)
+

torchzero/modules/misc/lr.py

@@ -0,0 +1,96 @@
+import random
+from collections.abc import Callable, Iterable
+from functools import partial
+from typing import TYPE_CHECKING, Any, overload
+
+import torch
+
+from ...tensorlist import TensorList
+
+from ...core import OptimizerModule
+
+if TYPE_CHECKING:
+    from ...optim import Modular
+
+def _init_scheduler_hook(opt: "Modular", module: "LR", scheduler_cls, **kwargs):
+    """post init hook that initializes the lr scheduler to the LR module and sets `_scheduler_step_fn`."""
+    scheduler = scheduler_cls(module, **kwargs)
+    module._scheduler_step_fn = scheduler.step
+
+def _set_momentum_hook(optimizer, state, momentum):
+    for module in optimizer.unrolled_modules:
+        if 'momentum' in module.defaults:
+            for g in module.param_groups:
+                g['momentum'] = momentum
+        elif 'beta1' in module.defaults:
+            for g in module.param_groups:
+                g['beta1'] = momentum
+
+class LR(OptimizerModule):
+    """Multiplies update by the learning rate. Optionally uses an lr scheduler.
+
+    Args:
+        lr (float, optional): learning rate. Defaults to 1e-3.
+        scheduler (Callable[..., torch.optim.lr_scheduler.LRScheduler  |  Any] | None, optional):
+            A scheduler class, for example `torch.optim.lr_scheduler.OneCycleLR`. Defaults to None.
+        cycle_momentum (bool, optional):
+            enables schedulers that support it to affect momentum (like OneCycleLR).
+            The momentum will be cycled on ALL modules that have `momentum` or `beta1` setting.
+            This does not support external optimizers, wrapped with `Wrap`. Defaults to True.
+        sheduler_step_every (int, optional):
+            step with scheduler every n optimizer steps.
+            Useful when the scheduler steps once per epoch. Defaults to 1.
+        **kwargs:
+            kwargs to pass to `scheduler`.
+    """
+    IS_LR_MODULE = True
+    def __init__(
+        self,
+        lr: float = 1e-3,
+        scheduler_cls: Callable[..., torch.optim.lr_scheduler.LRScheduler | Any] | None = None,
+        cycle_momentum: bool = True,
+        sheduler_step_every: int = 1,
+        # *args,
+        **kwargs,
+    ):
+
+        defaults = dict(lr = lr)
+
+        if (scheduler_cls is not None) and cycle_momentum:
+            defaults['momentum'] = 0
+        super().__init__(defaults)
+
+        self._scheduler_step_fn = None
+        self.sheduler_step_every = sheduler_step_every
+        self.cycle_momentum = cycle_momentum
+        self.cur = 0
+
+        if scheduler_cls is not None:
+            self.post_init_hooks.append(lambda opt, module: _init_scheduler_hook(opt, module, scheduler_cls, **kwargs))
+
+        self._skip = False
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        # step with scheduler
+        if self._scheduler_step_fn is not None:
+            if self.cur != 0 and self.cur % self.sheduler_step_every == 0:
+                self._scheduler_step_fn()
+
+                # add a hook to cycle momentum
+                if self.cycle_momentum:
+                    state.add_post_step_hook(_set_momentum_hook)
+
+            # remove init hook to delete reference to scheduler
+            if self.cur == 0 and len(self.post_init_hooks) == 1:
+                del self.post_init_hooks[0]
+
+        # skip if lr was applied by previous module (LR fusing)
+        if not self._skip:
+            # multiply ascent direction by lr in-place
+            lr = self.get_group_key('lr')
+            ascent *= lr
+
+        self.cur += 1
+        self._skip = False
+        return ascent

torchzero/modules/misc/multistep.py

@@ -0,0 +1,51 @@
+from collections.abc import Callable, Iterable
+
+import torch
+
+from ...tensorlist import TensorList
+
+from ...core import OptimizerModule, _Chainable
+
+
+class Multistep(OptimizerModule):
+    """Performs multiple steps (per batch), passes total update to the next module.
+
+    Args:
+        modules (_Chainable): modules to perform multiple steps with.
+        num_steps (int, optional): number of steps to perform. Defaults to 2.
+    """
+    def __init__(self, modules: _Chainable, num_steps: int = 2):
+        super().__init__({})
+        self.num_steps = num_steps
+
+        self._set_child_('modules', modules)
+
+    def step(self, state):
+        # no next module, just perform multiple steps
+        if self.next_module is None:
+            ret = None
+            for step in range(self.num_steps):
+                state_copy = state.copy(clone_ascent=True) if step != self.num_steps - 1 else state
+                ret = self.children['modules'].step(state_copy)
+
+                # since parameters are updated after stepping, grad and fx0 must be erased as they are no longer correct
+                state.grad = None; state.fx0 = None
+
+            return ret
+
+        # accumulate steps and pass to next module
+        p0 = self.get_params().clone()
+        for step in range(self.num_steps):
+            state_copy = state.copy(clone_ascent=True) if step != self.num_steps - 1 else state
+            self.children['modules'].step(state_copy)
+
+            # since parameters are updated after stepping, grad and fx0 must be erased as they are no longer correct
+            state.grad = None; state.fx0 = None
+
+        p1 = self.get_params()
+        state.ascent = p0 - p1
+
+        # undo ascent
+        p1.set_(p0)
+
+        return self._update_params_or_step_with_next(state, p1)

torchzero/modules/misc/on_increase.py

@@ -0,0 +1,53 @@
+import torch
+
+from ...core import OptimizerModule
+
+
+class NegateOnLossIncrease(OptimizerModule):
+    """Performs an additional evaluation to check if update increases the loss. If it does,
+    negates or backtracks the update.
+
+    Args:
+        backtrack (bool, optional):
+            if True, sets update to minus update, otherwise sets it to zero. Defaults to True.
+    """
+    def __init__(self, backtrack = True):
+        super().__init__({})
+        self.backtrack = backtrack
+
+    @torch.no_grad()
+    def step(self, state):
+        if state.closure is None: raise ValueError('NegateOnLossIncrease requires closure.')
+        if state.fx0 is None: state.fx0 = state.closure(False)
+
+        # subtract ascent direction to params and see if loss decreases
+        params = self.get_params()
+        ascent_direction = state.maybe_use_grad_(params)
+        params -= ascent_direction
+        state.fx0_approx = state.closure(False)
+
+        # if this has no children, update params and return loss
+        if self.next_module is None:
+            if params is None: params = self.get_params()
+
+            if state.fx0_approx > state.fx0:
+                # loss increased, so we negate thea scent direction
+                # we are currently at params - ascent direction
+                # so we add twice the ascent direction
+                params.add_(ascent_direction, alpha = 2 if self.backtrack else 1)
+
+            # else: we are already at a lower loss point
+            return state.get_loss()
+
+        # otherwise undo the ascent direction because it is passed to the child
+        params += ascent_direction
+
+        # if loss increases, negate ascent direction
+        if state.fx0_approx > state.fx0:
+            if self.backtrack: ascent_direction.neg_()
+            else: ascent_direction.zero_()
+
+        # otherwise undo the ascent direction and pass the updated ascent direction to the child
+        return self.next_module.step(state)
+
+

torchzero/modules/momentum/__init__.py

@@ -0,0 +1,4 @@
+"""
+Modules that implement momentum.
+"""
+from .momentum import HeavyBall, NesterovMomentum, RandomCoordinateMomentum, GradientAveraging

torchzero/modules/momentum/momentum.py

@@ -0,0 +1,106 @@
+from collections import abc
+
+import torch
+
+from ...tensorlist import TensorList
+from ...core import OptimizerModule
+
+def _heavyball_step(ascent, velocity: TensorList, momentum, dampening: TensorList):
+    velocity.mul_(momentum).add_(ascent * (1 - dampening))
+    return velocity.clone()
+
+class HeavyBall(OptimizerModule):
+    """Polyak's (heavy ball) momentum. Exactly matches pytorch SGD `momentum` option.
+
+    Args:
+        decay (float, optional): momentum decay. Defaults to 0.9.
+        dampening (float, optional): momentum dampening. Defaults to 0.
+    """
+    def __init__(self, momentum: float = 0.9, dampening: float = 0, ):
+        defaults = dict(momentum = momentum, dampening = dampening)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        velocity = self.get_state_key('velocity', init = ascent)
+        settings = self.get_all_group_keys()
+        updated_direction = _heavyball_step(ascent, velocity, settings['momentum'], settings['dampening'])
+        return updated_direction
+
+
+def _nesterov_step_(ascent, velocity: TensorList, momentum, dampening,):
+    # update velocity with the ascent direction
+    velocity += ascent
+
+    # decay velocity (this can be moved before previous line for slightly different results)
+    velocity *= momentum
+
+    # update ascent direction with velocity
+    ascent += velocity * (1 - dampening)
+
+
+class NesterovMomentum(OptimizerModule):
+    """Nesterov momentum. Exactly matches pytorch SGD with `nesterov=True`,
+    except this also supports dampening.
+
+    Args:
+        decay (float, optional): momentum decay. Defaults to 0.9.
+        dampening (float, optional): momentum dampening. Defaults to 0.
+    """
+    def __init__(self, decay: float = 0.9, dampening: float = 0, ):
+        defaults = dict(momentum = decay, dampening = dampening)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        velocity = self.get_state_key('velocity')
+        settings = self.get_all_group_keys()
+        _nesterov_step_(ascent, velocity, settings['momentum'], settings['dampening'])
+        return ascent
+
+class GradientAveraging(OptimizerModule):
+    """Averages last 2 gradients (TODO)"""
+    def __init__(self, dampening: float = 0, ):
+        defaults = dict(dampening = dampening)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        velocity = self.get_state_key('velocity')
+        dampening = self.get_group_key('dampening')
+
+        new_direction = ascent + velocity * (1-dampening)
+        velocity.copy_(ascent)
+
+        return new_direction
+
+
+class RandomCoordinateMomentum(OptimizerModule):
+    """Only uses `p` random coordinates of the new update. Other coordinates remain from previous update.
+    This works but I don't know if it is any good.
+
+    Args:
+        p (float, optional): probability to update velocity with a new weigh value. Defaults to 0.1.
+        nesterov (bool, optional): if False, update uses delayed momentum. Defaults to True.
+    """
+    def __init__(self, p: float = 0.1, nesterov=True):
+        defaults = dict(p=p)
+        super().__init__(defaults)
+        self.nesterov = nesterov
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        velocity = self.get_state_key('velocity', init = ascent)
+        settings = self.get_all_group_keys()
+
+        # pick p veclocity indexes to update with the new ascent direction
+        indexes = ascent.bernoulli_like(settings['p']).as_bool()
+
+        if self.nesterov:
+            # update the velocity at those indexes
+            velocity.masked_set_(mask = indexes, value = ascent)
+            return velocity.clone()
+
+        new_ascent = velocity.clone()
+        velocity.masked_set_(mask = indexes, value = ascent)
+        return new_ascent

torchzero/modules/operations/__init__.py

@@ -0,0 +1,29 @@
+from .multi import (
+    Add,
+    AddMagnitude,
+    Div,
+    Divide,
+    Interpolate,
+    Lerp,
+    Mul,
+    Pow,
+    Power,
+    RDiv,
+    RPow,
+    RSub,
+    Sub,
+    Subtract,
+)
+from .reduction import Mean, Product, Sum
+from .singular import (
+    Abs,
+    Cos,
+    MagnitudePower,
+    NanToNum,
+    Negate,
+    Operation,
+    Reciprocal,
+    Sign,
+    Sin,
+    sign_grad_,
+)