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

torchzero/modules/line_search/scipy_minimize_scalar.py

@@ -1,62 +0,0 @@
-import typing
-
-import torch
-try:
-    import scipy.optimize as scopt
-except ModuleNotFoundError:
-    scopt = typing.cast(typing.Any, None)
-
-from ...tensorlist import TensorList
-from ...core import OptimizationVars
-
-from .base_ls import LineSearchBase, MaxIterReached
-
-if typing.TYPE_CHECKING:
-    import scipy.optimize as scopt
-
-class ScipyMinimizeScalarLS(LineSearchBase):
-    """Line search via `scipy.optimize.minimize_scalar`. All args except maxiter are the same as for it.
-
-    Args:
-        method (Optional[str], optional): 'brent', 'golden' or 'bounded'. Defaults to None.
-        maxiter (Optional[int], optional): hard limit on maximum number of function evaluations. Defaults to None.
-        bracket (optional): bracket. Defaults to None.
-        bounds (optional): bounds. Defaults to None.
-        tol (Optional[float], optional): some kind of tolerance. Defaults to None.
-        options (optional): options for method. Defaults to None.
-        log_lrs (bool, optional): logs lrs and values into `_lrs`. Defaults to False.
-    """
-    def __init__(
-        self,
-        method: str | None = None,
-        maxiter: int | None = None,
-        bracket = None,
-        bounds = None,
-        tol: float | None = None,
-        options = None,
-        log_lrs = False,
-    ):
-        if scopt is None: raise ModuleNotFoundError("scipy is not installed")
-        super().__init__({}, maxiter=maxiter, log_lrs=log_lrs)
-        self.method = method
-        self.tol = tol
-        self.bracket = bracket
-        self.bounds = bounds
-        self.options = options
-
-    @torch.no_grad
-    def _find_best_lr(self, vars: OptimizationVars, params: TensorList) -> float:
-        try:
-            res = scopt.minimize_scalar(
-                self._evaluate_lr_ensure_float,
-                args = (vars.closure, vars.ascent, params),
-                method = self.method,
-                tol = self.tol,
-                bracket = self.bracket,
-                bounds = self.bounds,
-                options = self.options,
-            ) # type:ignore
-        except MaxIterReached:
-            pass
-
-        return float(self._best_lr)

torchzero/modules/meta/__init__.py

@@ -1,12 +0,0 @@
-"""Modules that use other modules."""
-# from .chain import Chain, ChainReturn
-import sys
-
-from .alternate import Alternate
-from .grafting import Graft, IntermoduleCautious, SignGrafting
-from .return_overrides import ReturnAscent, ReturnClosure, SetGrad
-
-# if sys.version_info[1] < 12:
-from .optimizer_wrapper import Wrap, WrapClosure
-# else:
-#     from .optimizer_wrapper import Wrap, WrapClosure

torchzero/modules/meta/alternate.py

@@ -1,65 +0,0 @@
-import random
-from collections.abc import Iterable
-from typing import Any, Literal
-
-from ...core import OptimizerModule, _Chainable
-
-
-class Alternate(OptimizerModule):
-    """Alternates stepping with multiple modules.
-
-    Args:
-        modules (Iterable[OptimizerModule  |  Iterable[OptimizerModule]]): modules to alternate between.
-        mode (int | list[int] | tuple[int] | "random"], optional):
-            can be integer - number of repeats for all modules;
-            list or tuple of integers per each module with number of repeats;
-            "random" to pick module randomly each time. Defaults to 1.
-        seed (int | None, optional): seed for "random" mode. Defaults to None.
-    """
-    def __init__(
-        self,
-        modules: Iterable[_Chainable],
-        mode: int | list[int] | tuple[int] | Literal["random"] = 1,
-        seed: int | None = None
-    ):
-        super().__init__({})
-        modules = list(modules)
-
-        for i,m in enumerate(modules):
-            self._set_child_(i, m)
-
-        self.random = random.Random(seed)
-
-        if isinstance(mode, int): mode = [mode for _ in modules]
-        self.mode: list[int] | tuple[int] | Literal['random'] = mode
-
-        self.cur = 0
-        if self.mode == 'random': self.remaining = 0
-        else:
-            self.remaining = self.mode[0]
-            if len(self.mode) != len(self.children):
-                raise ValueError(f"got {len(self.children)} modules but {len(mode)} repeats, they should be the same")
-
-    def step(self, vars):
-        if self.mode == 'random':
-            module = self.random.choice(list(self.children.values()))
-
-        else:
-            if self.remaining == 0:
-                self.cur += 1
-
-            if self.cur >= len(self.mode):
-                self.cur = 0
-
-            if self.remaining == 0: self.remaining = self.mode[self.cur]
-
-            module = self.children[self.cur]
-
-            self.remaining -= 1
-
-        if self.next_module is None:
-            return module.step(vars)
-
-        vars.ascent = module.return_ascent(vars)
-        return self._update_params_or_step_with_next(vars)
-

torchzero/modules/meta/grafting.py

@@ -1,195 +0,0 @@
-from collections.abc import Iterable
-from typing import Literal
-import torch
-
-from ...core import OptimizerModule
-from ...tensorlist import TensorList
-
-
-class Graft(OptimizerModule):
-    """
-    Optimizer grafting (magnitude#direction).
-    Takes update of one optimizer and makes its norm same as update of another optimizer.
-    Can be applied to all weights or layerwise.
-
-    Args:
-        magnitude (OptimizerModule | Iterable[OptimizerModule]):
-            module to use magnitude from.
-            If sequence of modules is provided, they will be chained.
-        direction (OptimizerModule | Iterable[OptimizerModule]):
-            module/modules to use direction from.
-            If sequence of modules is provided, they will be chained.
-        ord (int, optional): norm type. Defaults to 2.
-        eps (float, optional): epsilon for numerical stability. Defaults to 1e-8.
-        layerwise (bool, optional): whether to apply grafting layerwise. Defaults to False.
-
-    reference
-        *Agarwal, N., Anil, R., Hazan, E., Koren, T., & Zhang, C.
-        Learning Rate Grafting: Transferability of Optimizer Tuning.*
-    """
-    def __init__(
-        self,
-        magnitude: OptimizerModule | Iterable[OptimizerModule],
-        direction: OptimizerModule | Iterable[OptimizerModule],
-        ord: float = 2,
-        eps: float = 1e-8,
-        layerwise: bool = False,
-        # TODO: channelwise
-    ):
-        super().__init__({})
-        self._set_child_('magnitude', magnitude)
-        self._set_child_('direction', direction)
-        self.ord = ord
-        self.eps = eps
-        self.layerwise = layerwise
-
-
-    @torch.no_grad
-    def step(self, vars):
-        state_copy = vars.copy(clone_ascent=True)
-        magnitude = self.children['magnitude'].return_ascent(state_copy)
-
-        if state_copy.grad is not None: vars.grad = state_copy.grad
-        if state_copy.fx0 is not None: vars.fx0 = state_copy.fx0
-        if state_copy.fx0_approx is not None: vars.fx0_approx = state_copy.fx0_approx
-
-        direction = self.children['direction'].return_ascent(vars)
-
-        if self.layerwise:
-            M = magnitude.norm(self.ord)
-            D = direction.norm(self.ord)
-            D.select_set_(D == 0, M)
-
-        else:
-            M = magnitude.total_vector_norm(self.ord)
-            D = direction.total_vector_norm(self.ord)
-            if D == 0: D = M
-
-        vars.ascent = direction.mul_(M / (D + self.eps))
-        return self._update_params_or_step_with_next(vars)
-
-
-
-class SignGrafting(OptimizerModule):
-    """Weight-wise grafting-like operation where sign of the ascent is taken from first module
-    and magnitude from second module.
-
-    Args:
-        magnitude (OptimizerModule | Iterable[OptimizerModule]):
-            module to take magnitude from.
-            If sequence of modules is provided, they will be chained.
-        sign (OptimizerModule | Iterable[OptimizerModule]):
-            module to take sign from.
-            If sequence of modules is provided, they will be chained.
-    """
-    def __init__(
-        self,
-        magnitude: OptimizerModule | Iterable[OptimizerModule],
-        sign: OptimizerModule | Iterable[OptimizerModule],
-    ):
-        super().__init__({})
-
-        self._set_child_('magnitude', magnitude)
-        self._set_child_('sign', sign)
-
-
-    @torch.no_grad
-    def step(self, vars):
-        state_copy = vars.copy(clone_ascent=True)
-        magnitude = self.children['magnitude'].return_ascent(state_copy)
-
-        # make sure to store grad and fx0 if it was calculated
-        vars.update_attrs_(state_copy)
-
-        sign = self.children['sign'].return_ascent(vars)
-
-        vars.ascent = magnitude.copysign_(sign)
-        return self._update_params_or_step_with_next(vars)
-
-
-class IntermoduleCautious(OptimizerModule):
-    """Negates update for parameters where updates of two modules or module chains have inconsistent sign.
-    Optionally normalizes the update by the number of parameters that are not masked.
-
-    Args:
-        main_module (OptimizerModule | Iterable[OptimizerModule]):
-            main module or sequence of modules to chain, which update will be used with a consistency mask applied.
-        compare_module (OptimizerModule | Iterable[OptimizerModule]):
-            module or sequence of modules to chain, which update will be used to compute a consistency mask.
-            Can also be set to `ascent` to compare to update that is passed `main_module`, or `grad` to compare
-            to gradients.
-        normalize (bool, optional):
-            renormalize update after masking.
-            only has effect when mode is 'zero'. Defaults to False.
-        eps (float, optional): epsilon for normalization. Defaults to 1e-6.
-        mode (str, optional):
-            what to do with updates with inconsistent signs.
-
-            "zero" - set them to zero (as in paper)
-
-            "grad" - set them to the gradient
-
-            "compare_module" - set them to `compare_module`'s update
-
-            "negate" - negate them (same as using update magnitude and gradient sign)
-    """
-    def __init__(
-        self,
-        main_module: OptimizerModule | Iterable[OptimizerModule],
-        compare_module: OptimizerModule | Iterable[OptimizerModule] | Literal['ascent', 'grad'],
-        normalize=False,
-        eps=1e-6,
-        mode: Literal["zero", "grad", "backtrack", "compare_module"] = "zero",
-    ):
-        super().__init__({})
-
-        self._set_child_('main',main_module)
-        if isinstance(compare_module, str): self.compare_mode = compare_module
-        else:
-            self._set_child_('compare', compare_module)
-            self.compare_mode = 'module'
-        self.eps = eps
-        self.normalize = normalize
-        self.mode: Literal["zero", "grad", "backtrack", "compare_module"]  = mode
-
-    @torch.no_grad
-    def step(self, vars):
-        params = None
-        state_copy = vars.copy(clone_ascent=True)
-        ascent = self.children['main'].return_ascent(state_copy)
-        vars.update_attrs_(state_copy)
-
-        if self.compare_mode == 'module': compare = self.children['compare'].return_ascent(vars)
-        else:
-            params = self.get_params()
-            if self.compare_mode == 'ascent': compare: TensorList = vars.maybe_use_grad_(params)
-            elif self.compare_mode == 'grad': compare: TensorList = vars.maybe_compute_grad_(params)
-            else: raise ValueError(f'Invalid compare_module: {self.compare_mode}')
-
-        # mask will be > 0 for parameters where both signs are the same
-        mask = (ascent * compare) > 0
-
-        if self.mode == 'backtrack':
-            ascent -= ascent.mul(2).mul_(mask.logical_not_())
-
-        else:
-            # normalize if mode is `zero`
-            if self.normalize and self.mode == 'zero':
-                fmask = mask.to(ascent[0].dtype)
-                fmask /= fmask.total_mean() + self.eps
-            else:
-                fmask = mask
-
-            # apply the mask
-            ascent *= fmask
-
-            if self.mode == 'grad':
-                params = self.get_params()
-                ascent += vars.maybe_compute_grad_(params) * mask.logical_not_()
-
-            elif self.mode == 'compare_module':
-                ascent += compare * mask.logical_not_()
-
-        vars.ascent = ascent
-        return self._update_params_or_step_with_next(vars, params)
-

torchzero/modules/meta/optimizer_wrapper.py

@@ -1,173 +0,0 @@
-from collections.abc import Callable, Sequence
-from typing import Any, overload
-
-import torch
-from typing_extensions import Concatenate, ParamSpec
-
-from ...core import OptimizerModule
-from .return_overrides import SetGrad
-
-K = ParamSpec('K')
-
-class Wrap(OptimizerModule):
-    """
-    Wraps any torch.optim.Optimizer.
-
-    Sets .grad attribute to the current update and steps with the `optimizer`.
-
-    Additionally, if this is not the last module, this takes the update of `optimizer`,
-    undoes it and passes to the next module instead. That means you can chain multiple
-    optimizers together.
-
-    Args:
-        optimizer (torch.optim.Optimizer): optimizer to wrap,
-            or a callable (class) that constructs the optimizer.
-        kwargs:
-            if class is passed, kwargs are passed to the constructor.
-            parameters are passed separately and automatically
-            which is the point of passing a constructor
-            instead of an optimizer directly.
-
-    This can be constructed in two ways.
-    .. code-block:: python
-        wrapper = OptimizerWrapper(torch.optim.SGD(model.parameters(), lr = 0.1))
-        # or
-        wrapper = OptimizerWrapper(torch.optim.SGD, lr = 0.1)
-    """
-
-    @overload
-    def __init__(self, optimizer: torch.optim.Optimizer): ...
-    @overload
-    # def __init__[**K](
-    def __init__(
-        self,
-        optimizer: Callable[Concatenate[Any, K], torch.optim.Optimizer],
-        *args: K.args,
-        **kwargs: K.kwargs,
-        # optimizer: abc.Callable[..., torch.optim.Optimizer],
-        # *args,
-        # **kwargs,
-    ): ...
-    def __init__(self, optimizer, *args, **kwargs):
-
-        super().__init__({})
-        self._optimizer_cls: torch.optim.Optimizer | Callable[..., torch.optim.Optimizer] = optimizer
-        self._args = args
-        self._kwargs = kwargs
-
-    def _initialize_(self, params, set_passed_params):
-        """Initializes this optimizer and all children with the given parameters."""
-        super()._initialize_(params, set_passed_params=set_passed_params)
-        if isinstance(self._optimizer_cls, torch.optim.Optimizer) or not callable(self._optimizer_cls):
-            self.optimizer = self._optimizer_cls
-        else:
-            self.optimizer = self._optimizer_cls(params, *self._args, **self._kwargs)
-
-    @torch.no_grad
-    def step(self, vars):
-        # check attrs
-        # if self.pass_closure:
-        #     if state.closure is None: raise ValueError('ClosureOptimizerWrapper requires closure.')
-        #     if state.ascent is not None:
-        #         raise ValueError('pass_closure = True, means ascent must be None (not sure though)')
-
-        params = self.get_params()
-
-        if self.next_module is None:
-            # set grad to ascent and make a step with the optimizer
-            g = vars.maybe_use_grad_(params)
-            params.set_grad_(g)
-            vars.fx0 = self.optimizer.step()
-            return vars.get_loss()
-
-
-        params_before_step = params.clone()
-
-        g = vars.maybe_use_grad_(params)
-        params.set_grad_(g)
-        vars.fx0 = self.optimizer.step()
-
-        # calculate update as difference in params
-        vars.ascent = params_before_step - params
-        params.set_(params_before_step)
-        return self.next_module.step(vars)
-
-
-class WrapClosure(OptimizerModule):
-    """
-    Wraps any torch.optim.Optimizer. This only works with modules with :code:`target = "Closure"` argument.
-    The modified closure will be passed to the optimizer.
-
-    Alternative any module can be turned into a closure module by using :any:`MakeClosure` module,
-    in that case this should be placed after MakeClosure.
-
-    Args:
-        optimizer (torch.optim.Optimizer): optimizer to wrap,
-            or a callable (class) that constructs the optimizer.
-        kwargs:
-            if class is passed, kwargs are passed to the constructor.
-            parameters are passed separately and automatically
-            which is the point of passing a constructor
-            instead of an optimizer directly.
-
-    This can be constructed in two ways.
-
-    .. code-block:: python
-
-        wrapper = OptimizerWrapper(torch.optim.SGD(model.parameters(), lr = 0.1))
-        # or
-        wrapper = OptimizerWrapper(torch.optim.SGD, lr = 0.1)
-
-    """
-
-    @overload
-    def __init__(self, optimizer: torch.optim.Optimizer,): ...
-    @overload
-    def __init__(
-        self,
-        optimizer: Callable[Concatenate[Any, K], torch.optim.Optimizer],
-        *args: K.args,
-        **kwargs: K.kwargs,
-        # optimizer: abc.Callable[..., torch.optim.Optimizer],
-        # *args,
-        # **kwargs,
-    ): ...
-    def __init__(self, optimizer, *args, **kwargs):
-
-        super().__init__({})
-        self._optimizer_cls: torch.optim.Optimizer | Callable[..., torch.optim.Optimizer] = optimizer
-        self._args = args
-        self._kwargs = kwargs
-
-    def _initialize_(self, params, set_passed_params):
-        """Initializes this optimizer and all children with the given parameters."""
-        super()._initialize_(params, set_passed_params=set_passed_params)
-        if isinstance(self._optimizer_cls, torch.optim.Optimizer) or not callable(self._optimizer_cls):
-            self.optimizer = self._optimizer_cls
-        else:
-            self.optimizer = self._optimizer_cls(params, *self._args, **self._kwargs)
-
-    @torch.no_grad
-    def step(self, vars):
-        # check attrs
-        # if self.pass_closure:
-        #     if state.closure is None: raise ValueError('ClosureOptimizerWrapper requires closure.')
-        #     if state.ascent is not None:
-        #         raise ValueError('pass_closure = True, means ascent must be None (not sure though)')
-
-        params = self.get_params()
-
-        if self.next_module is None:
-            # set grad to ascent and make a step with the optimizer
-            vars.fx0 = self.optimizer.step(vars.closure) # type:ignore
-            return vars.get_loss()
-
-
-        params_before_step = params.clone()
-        vars.fx0 = self.optimizer.step(vars.closure) # type:ignore
-
-        # calculate update as difference in params
-        vars.ascent = params_before_step - params
-        params.set_(params_before_step)
-        return self.next_module.step(vars)
-

torchzero/modules/meta/return_overrides.py

@@ -1,46 +0,0 @@
-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, vars):
-        if self.next_module is not None: raise ValueError("SetGrad can't have children")
-        params = self.get_params()
-        g = vars.maybe_use_grad_(params) # this may execute the closure which might be modified
-        params.set_grad_(g)
-        return vars.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, vars) -> TensorList: # type:ignore
-        if self.next_module is not None: raise ValueError("ReturnAscent can't have children")
-        params = self.get_params()
-        update = vars.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, vars) -> _ClosureType: # type:ignore
-        if self.next_module is not None: raise ValueError("ReturnClosure can't have children")
-        if vars.closure is None:
-            raise ValueError("MakeClosure requires closure")
-        return vars.closure
-

torchzero/modules/misc/__init__.py

@@ -1,10 +0,0 @@
-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

@@ -1,43 +0,0 @@
-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, vars):
-        self.cur_step += 1
-
-        params = self.get_params()
-        accumulated_update = self.get_state_key('accumulated_grads')
-        accumulated_update += vars.maybe_use_grad_(params)
-
-        if self.cur_step % self.n_steps == 0:
-            vars.ascent = accumulated_update.clone()
-            if self.mean: vars.ascent /= self.n_steps
-            accumulated_update.zero_()
-            return self._update_params_or_step_with_next(vars)
-
-
-        return vars.get_loss()