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

torchzero/modules/misc/basic.py

@@ -1,115 +0,0 @@
-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, vars, 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, vars, ascent): return ascent.clone()
-
-class Identity(OptimizerModule):
-    """Does nothing."""
-    def __init__(self, *args, **kwargs):
-        super().__init__({})
-
-    @torch.no_grad
-    def _update(self, vars, 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, vars, 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, vars, ascent):
-        ascent = vars.ascent = vars.maybe_compute_grad_(self.get_params())
-        return ascent
-
-class Zeros(OptimizerModule):
-    def __init__(self):
-        super().__init__({})
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        return ascent.zeros_like()
-
-class Fill(OptimizerModule):
-    def __init__(self, value):
-        super().__init__({"value": value})
-
-    @torch.no_grad
-    def _update(self, vars, 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, vars, ascent):
-        vars.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, vars):
-        if vars.closure is None: raise ValueError("MakeClosure requires a closure")
-
-        params = self.get_params()
-        orig_closure = vars.closure
-        orig_state = vars.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)
-
-        vars.closure = new_closure # type:ignore
-        return self._update_params_or_step_with_next(vars)
-

torchzero/modules/misc/lr.py

@@ -1,96 +0,0 @@
-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, vars, 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:
-                    vars.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

@@ -1,51 +0,0 @@
-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, vars):
-        # no next module, just perform multiple steps
-        if self.next_module is None:
-            ret = None
-            for step in range(self.num_steps):
-                state_copy = vars.copy(clone_ascent=True) if step != self.num_steps - 1 else vars
-                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
-                vars.grad = None; vars.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 = vars.copy(clone_ascent=True) if step != self.num_steps - 1 else vars
-            self.children['modules'].step(state_copy)
-
-            # since parameters are updated after stepping, grad and fx0 must be erased as they are no longer correct
-            vars.grad = None; vars.fx0 = None
-
-        p1 = self.get_params()
-        vars.ascent = p0 - p1
-
-        # undo ascent
-        p1.set_(p0)
-
-        return self._update_params_or_step_with_next(vars, p1)

torchzero/modules/misc/on_increase.py

@@ -1,53 +0,0 @@
-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, vars):
-        if vars.closure is None: raise ValueError('NegateOnLossIncrease requires closure.')
-        if vars.fx0 is None: vars.fx0 = vars.closure(False)
-
-        # subtract ascent direction to params and see if loss decreases
-        params = self.get_params()
-        ascent_direction = vars.maybe_use_grad_(params)
-        params -= ascent_direction
-        vars.fx0_approx = vars.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 vars.fx0_approx > vars.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 vars.get_loss()
-
-        # otherwise undo the ascent direction because it is passed to the child
-        params += ascent_direction
-
-        # if loss increases, negate ascent direction
-        if vars.fx0_approx > vars.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(vars)
-
-

torchzero/modules/operations/__init__.py

@@ -1,29 +0,0 @@
-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_,
-)

torchzero/modules/operations/multi.py

@@ -1,298 +0,0 @@
-from collections.abc import Iterable
-import torch
-
-from ...core import OptimizerModule
-
-_Value = int | float | OptimizerModule | Iterable[OptimizerModule]
-
-class Add(OptimizerModule):
-    """add `value` to update. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
-    def __init__(self, value: _Value):
-        super().__init__({})
-
-        if not isinstance(value, (int, float)):
-            self._set_child_('value', value)
-
-        self.value = value
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.value, (int, float)):
-            return ascent.add_(self.value)
-
-        state_copy = vars.copy(clone_ascent = True)
-        v = self.children['value'].return_ascent(state_copy)
-        return ascent.add_(v)
-
-
-class Sub(OptimizerModule):
-    """subtracts `value` from update. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
-    def __init__(self, subtrahend: _Value):
-        super().__init__({})
-
-        if not isinstance(subtrahend, (int, float)):
-            self._set_child_('subtrahend', subtrahend)
-
-        self.subtrahend = subtrahend
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.subtrahend, (int, float)):
-            return ascent.sub_(self.subtrahend)
-
-        state_copy = vars.copy(clone_ascent = True)
-        subtrahend = self.children['subtrahend'].return_ascent(state_copy)
-        return ascent.sub_(subtrahend)
-
-class RSub(OptimizerModule):
-    """subtracts update from `value`. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
-    def __init__(self, minuend: _Value):
-        super().__init__({})
-
-        if not isinstance(minuend, (int, float)):
-            self._set_child_('minuend', minuend)
-
-        self.minuend = minuend
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.minuend, (int, float)):
-            return ascent.sub_(self.minuend).neg_()
-
-        state_copy = vars.copy(clone_ascent = True)
-        minuend = self.children['minuend'].return_ascent(state_copy)
-        return ascent.sub_(minuend).neg_()
-
-class Subtract(OptimizerModule):
-    """Calculates `minuend - subtrahend`"""
-    def __init__(
-        self,
-        minuend: OptimizerModule | Iterable[OptimizerModule],
-        subtrahend: OptimizerModule | Iterable[OptimizerModule],
-    ):
-        super().__init__({})
-        self._set_child_('minuend', minuend)
-        self._set_child_('subtrahend', subtrahend)
-
-    @torch.no_grad
-    def step(self, vars):
-        state_copy = vars.copy(clone_ascent = True)
-        minuend = self.children['minuend'].return_ascent(state_copy)
-        vars.update_attrs_(state_copy)
-        subtrahend = self.children['subtrahend'].return_ascent(vars)
-
-        vars.ascent = minuend.sub_(subtrahend)
-        return self._update_params_or_step_with_next(vars)
-
-class Mul(OptimizerModule):
-    """multiplies update by `value`. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
-    def __init__(self, value: _Value):
-        super().__init__({})
-
-        if not isinstance(value, (int, float)):
-            self._set_child_('value', value)
-
-        self.value = value
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.value, (int, float)):
-            return ascent.mul_(self.value)
-
-        state_copy = vars.copy(clone_ascent = True)
-        v = self.children['value'].return_ascent(state_copy)
-        return ascent.mul_(v)
-
-
-class Div(OptimizerModule):
-    """divides update by `value`. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
-    def __init__(self, denominator: _Value):
-        super().__init__({})
-
-        if not isinstance(denominator, (int, float)):
-            self._set_child_('denominator', denominator)
-
-        self.denominator = denominator
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.denominator, (int, float)):
-            return ascent.div_(self.denominator)
-
-        state_copy = vars.copy(clone_ascent = True)
-        denominator = self.children['denominator'].return_ascent(state_copy)
-        return ascent.div_(denominator)
-
-class RDiv(OptimizerModule):
-    """`value` by update. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
-    def __init__(self, numerator: _Value):
-        super().__init__({})
-
-        if not isinstance(numerator, (int, float)):
-            self._set_child_('numerator', numerator)
-
-        self.numerator = numerator
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.numerator, (int, float)):
-            return ascent.reciprocal_().mul_(self.numerator)
-
-        state_copy = vars.copy(clone_ascent = True)
-        numerator = self.children['numerator'].return_ascent(state_copy)
-        return ascent.reciprocal_().mul_(numerator)
-
-class Divide(OptimizerModule):
-    """calculates *numerator / denominator*"""
-    def __init__(
-        self,
-        numerator: OptimizerModule | Iterable[OptimizerModule],
-        denominator: OptimizerModule | Iterable[OptimizerModule],
-    ):
-        super().__init__({})
-        self._set_child_('numerator', numerator)
-        self._set_child_('denominator', denominator)
-
-    @torch.no_grad
-    def step(self, vars):
-        state_copy = vars.copy(clone_ascent = True)
-        numerator = self.children['numerator'].return_ascent(state_copy)
-        vars.update_attrs_(state_copy)
-        denominator = self.children['denominator'].return_ascent(vars)
-
-        vars.ascent = numerator.div_(denominator)
-        return self._update_params_or_step_with_next(vars)
-
-
-class Pow(OptimizerModule):
-    """takes ascent to the power of `value`. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
-    def __init__(self, power: _Value):
-        super().__init__({})
-
-        if not isinstance(power, (int, float)):
-            self._set_child_('power', power)
-
-        self.power = power
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.power, (int, float)):
-            return ascent.pow_(self.power)
-
-        state_copy = vars.copy(clone_ascent = True)
-        power = self.children['power'].return_ascent(state_copy)
-        return ascent.pow_(power)
-
-class RPow(OptimizerModule):
-    """takes `value` to the power of ascent. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
-    def __init__(self, base: _Value):
-        super().__init__({})
-
-        if not isinstance(base, (int, float)):
-            self._set_child_('base', base)
-
-        self.base = base
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.base, (int, float)):
-            return self.base ** ascent
-
-        state_copy = vars.copy(clone_ascent = True)
-        base = self.children['base'].return_ascent(state_copy)
-        return base.pow_(ascent)
-
-class Power(OptimizerModule):
-    """calculates *base ^ power*"""
-    def __init__(
-        self,
-        base: OptimizerModule | Iterable[OptimizerModule],
-        power: OptimizerModule | Iterable[OptimizerModule],
-    ):
-        super().__init__({})
-        self._set_child_('base', base)
-        self._set_child_('power', power)
-
-    @torch.no_grad
-    def step(self, vars):
-        state_copy = vars.copy(clone_ascent = True)
-        base = self.children['base'].return_ascent(state_copy)
-        vars.update_attrs_(state_copy)
-        power = self.children['power'].return_ascent(vars)
-
-        vars.ascent = base.pow_(power)
-        return self._update_params_or_step_with_next(vars)
-
-
-class Lerp(OptimizerModule):
-    """Linear interpolation between update and `end` based on scalar `weight`.
-
-    `out = update + weight * (end - update)`"""
-    def __init__(self, end: OptimizerModule | Iterable[OptimizerModule], weight: float):
-        super().__init__({})
-
-        self._set_child_('end', end)
-        self.weight = weight
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-
-        state_copy = vars.copy(clone_ascent = True)
-        end = self.children['end'].return_ascent(state_copy)
-        return ascent.lerp_(end, self.weight)
-
-
-class Interpolate(OptimizerModule):
-    """Does a linear interpolation of two module's updates - `start` (given by input), and `end`, based on a scalar
-    `weight`.
-
-    `out = input + weight * (end - input)`"""
-    def __init__(
-        self,
-        input: OptimizerModule | Iterable[OptimizerModule],
-        end: OptimizerModule | Iterable[OptimizerModule],
-        weight: float,
-    ):
-        super().__init__({})
-        self._set_child_('input', input)
-        self._set_child_('end', end)
-        self.weight = weight
-
-    @torch.no_grad
-    def step(self, vars):
-        state_copy = vars.copy(clone_ascent = True)
-        input = self.children['input'].return_ascent(state_copy)
-        vars.update_attrs_(state_copy)
-        end = self.children['end'].return_ascent(vars)
-
-        vars.ascent = input.lerp_(end, weight = self.weight)
-
-        return self._update_params_or_step_with_next(vars)
-
-class AddMagnitude(OptimizerModule):
-    """Add `value` multiplied by sign of the ascent, i.e. this adds `value` to the magnitude of the update.
-
-    Args:
-        value (Value): value to add to magnitude, either a float or an OptimizerModule.
-        add_to_zero (bool, optional):
-            if True, adds `value` to 0s. Otherwise, zeros remain zero.
-            Only has effect if value is a float. Defaults to True.
-    """
-    def __init__(self, value: _Value, add_to_zero=True):
-        super().__init__({})
-
-        if not isinstance(value, (int, float)):
-            self._set_child_('value', value)
-
-        self.value = value
-        self.add_to_zero = add_to_zero
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if isinstance(self.value, (int, float)):
-            if self.add_to_zero: return ascent.add_(ascent.clamp_magnitude(min=1).sign_().mul_(self.value))
-            return ascent.add_(ascent.sign_().mul_(self.value))
-
-        state_copy = vars.copy(clone_ascent = True)
-        v = self.children['value'].return_ascent(state_copy)
-        return ascent.add_(v.abs_().mul_(ascent.sign()))