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

torchzero/optim/modular.py

@@ -1,148 +0,0 @@
-from collections import abc
-import warnings
-from inspect import cleandoc
-import torch
-from typing import Any
-
-from ..core import OptimizerModule, TensorListOptimizer, OptimizationVars, _Chain, _Chainable
-from ..utils.python_tools import flatten
-
-def _unroll_modules(flat_modules: list[OptimizerModule], nested) -> list[OptimizerModule]:
-    """returns a list of all modules, including all nested ones"""
-    unrolled = []
-    for m in flat_modules:
-        unrolled.append(m)
-        if len(m.children) > 0:
-            unrolled.extend(_unroll_modules(list(m.children.values()), nested=True))
-        if nested:
-            if m.next_module is not None:
-                unrolled.extend(_unroll_modules([m.next_module], nested=True))
-    return unrolled
-
-
-class Modular(TensorListOptimizer):
-    """Creates a modular optimizer by chaining together a sequence of optimizer modules.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        *modules (Iterable[OptimizerModule] | OptimizerModule):
-            A sequence of optimizer modules to chain together. This argument will be flattened."""
-    def __init__(self, params, *modules: _Chainable):
-        flat_modules = flatten(modules)
-        self.modules: list[OptimizerModule] = flat_modules
-        self.chain = _Chain(flat_modules)
-
-        # save unrolled modules and make sure there is only 1 LR module.
-        self.unrolled_modules = _unroll_modules(flat_modules, nested=False)
-        num_lr_modules = len([m for m in self.unrolled_modules if m.IS_LR_MODULE])
-        if num_lr_modules > 1:
-            warnings.warn(cleandoc(
-                f"""More then 1 lr modules have been added.
-                This may lead to incorrect behaviour with learning rate scheduling and per-parameter learning rates.
-                Make sure there is a single `LR` module, use `Alpha` module instead of it where needed.
-                \nList of modules: {self.unrolled_modules}; \nlist of lr modules: {[m for m in self.unrolled_modules if m.IS_LR_MODULE]}"""
-            ))
-
-        if isinstance(params, torch.nn.Module):
-            self.model = params
-            params = list(params.parameters())
-        else:
-            self.model = None
-            params = list(params)
-
-        # if there is an `lr` setting, make sure there is an LR module that can use it
-        for p in params:
-            if isinstance(p, dict):
-                if 'lr' in p:
-                    if num_lr_modules == 0:
-                        warnings.warn(cleandoc(
-                            """Passed "lr" setting in a parameter group, but there is no LR module that can use that setting.
-                            Add an `LR` module to make per-layer "lr" setting work."""
-                        ))
-
-        super().__init__(params, {})
-        self.chain._initialize_(params, set_passed_params=True)
-
-        # run post-init hooks
-        for module in self.unrolled_modules:
-            for hook in module.post_init_hooks:
-                hook(self, module)
-
-    def state_dict(self):
-        state_dict = {}
-        state_dict['__self__'] = super().state_dict()
-        for i,v in enumerate(self.unrolled_modules):
-            state_dict[str(i)] = v.state_dict()
-        return state_dict
-
-    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
-        super().load_state_dict(state_dict['__self__'])
-        for i,v in enumerate(self.unrolled_modules):
-            if str(i) in state_dict:
-                v.load_state_dict(state_dict[str(i)])
-            else:
-                warnings.warn(f"Tried to load state dict for {i}th module: {v.__class__.__name__}, but it is not present in state_dict with {list(state_dict.keys()) = }")
-
-    def get_lr_module(self, last=True) -> OptimizerModule:
-        """
-        Retrieves the module in the chain that controls the learning rate.
-
-        This method is useful for setting up a learning rate scheduler. By default, it retrieves the last module in the chain
-        that has an `lr` group parameter.
-
-        Args:
-            last (bool, optional):
-                If multiple modules have an `lr` parameter, this argument controls which one is returned.
-                - If `True` (default), the last module is returned.
-                - If `False`, the first module is returned.
-
-        Returns:
-            OptimizerModule: The module that controls the learning rate.
-
-        Raises:
-            ValueError: If no modules in the chain have an `lr` parameter. To fix this, add an `LR` module.
-
-        Example:
-
-        .. code:: py
-            from torch.optim.lr_scheduler import OneCycleLR
-            import torchzero as tz
-
-            opt = tz.Modular(model.parameters(), [tz.m.RMSProp(), tz.m.LR(1e-2), tz.m.DirectionalNewton()])
-            lr_scheduler = OneCycleLR(opt.get_lr_module(), max_lr = 1e-1, total_steps = 1000, cycle_momentum=False)
-
-        """
-        modules = list(reversed(self.unrolled_modules)) if last else self.unrolled_modules
-        for m in modules:
-            if 'lr' in m.param_groups[0]: return m
-
-        raise ValueError(f'No modules out of {", ".join(m.__class__.__name__ for m in modules)} support and `lr` parameter. The easiest way to fix is is to add an `LR(1)` module at the end.')
-
-    def get_module_by_name(self, name: str | type, last=True) -> OptimizerModule:
-        """Returns the first or last module in the chain that matches the provided name or type.
-
-        Args:
-            name (str | type): the name (as a string) or the type of the module to search for.
-            last (bool, optional):
-                If multiple modules match, this argument controls which one is returned.
-                - If `True` (default), the last matching module is returned.
-                - If `False`, the first matching module is returned.
-
-        Returns:
-            OptimizerModule: The matching optimizer module.
-
-        Raises:
-            ValueError: If no modules in the chain match the provided name or type.
-        """
-        modules = list(reversed(self.unrolled_modules)) if last else self.unrolled_modules
-        for m in modules:
-            if isinstance(name, str) and m.__class__.__name__ == name: return m
-            if isinstance(name, type) and isinstance(m, name): return m
-
-        raise ValueError(f'No modules out of {", ".join(m.__class__.__name__ for m in modules)} match "{name}".')
-
-    def step(self, closure=None): # type:ignore
-        vars = OptimizationVars(closure, self.model)
-        res = self.chain.step(vars)
-        for hook in vars.post_step_hooks: hook(self, vars)
-        return res

torchzero/optim/quasi_newton/__init__.py

@@ -1 +0,0 @@
-from .directional_newton import DirectionalNewton

torchzero/optim/quasi_newton/directional_newton.py

@@ -1,58 +0,0 @@
-from ...modules import (
-    SGD,
-)
-from ...modules import DirectionalNewton as _DirectionalNewton, LR
-from ..modular import Modular
-
-
-class DirectionalNewton(Modular):
-    """Minimizes a parabola in the direction of the gradient (or update if momentum or weight decay is enabled)
-    via one additional forward pass, and uses another forward pass to make sure it didn't overstep.
-    So in total this performs three forward passes and one backward.
-
-    First forward and backward pass is used to calculate the value and gradient at initial parameters.
-    Then a gradient descent step is performed with `lr` learning rate, and loss is recalculated
-    with new parameters. A quadratic is fitted to two points and gradient,
-    if it has positive curvature, this makes a step towards the minimum, and checks if lr decreased
-    with an additional forward pass.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float, optional):
-            learning rate. Since you shouldn't put this module after LR(), you have to specify
-            the learning rate in this argument. Defaults to 1e-2.
-        max_dist (float | None, optional):
-            maximum distance to step when minimizing quadratic.
-            If minimum is further than this distance, minimization is not performed. Defaults to 1e4.
-        validate_step (bool, optional):
-            uses an additional forward pass to check
-            if step towards the minimum actually decreased the loss. Defaults to True.
-        momentum (float, optional): momentum. Defaults to 0.
-        dampening (float, optional): momentum dampening. Defaults to 0.
-        weight_decay (float, optional): weight decay (L2 regularization). Defaults to 0.
-        nesterov (bool, optional):
-            enables nesterov momentum, otherwise uses heavyball momentum. Defaults to False.
-
-    Note:
-        While lr scheduling is supported, this uses lr of the first parameter for all parameters.
-    """
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-4,
-        max_dist: float | None = 1e5,
-        validate_step: bool = True,
-        momentum: float = 0,
-        dampening: float = 0,
-        weight_decay: float = 0,
-        nesterov: bool = False,
-
-    ):
-
-        modules = [
-            SGD(momentum=momentum,dampening=dampening,weight_decay=weight_decay,nesterov=nesterov),
-            LR(lr),
-            _DirectionalNewton(max_dist, validate_step)
-        ]
-        super().__init__(params, modules)
-

torchzero/optim/second_order/__init__.py

@@ -1 +0,0 @@
-from .newton import ExactNewton

torchzero/optim/second_order/newton.py

@@ -1,94 +0,0 @@
-from typing import Any, Literal
-
-import torch
-
-from ...modules import (
-    LR,
-    ClipNorm,
-    FallbackLinearSystemSolvers,
-    LinearSystemSolvers,
-    LineSearches,
-    get_line_search,
-)
-from ...modules import ExactNewton as _ExactNewton
-from ..modular import Modular
-
-
-class ExactNewton(Modular):
-    """Peforms an exact Newton step using batched autograd. Note that torch.func would be way more efficient
-    but much more restrictive to what operations are allowed (I will add it at some point).
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float, optional): learning rate. Defaults to 1.
-        tikhonov (float, optional):
-            tikhonov regularization (constant value added to the diagonal of the hessian). Defaults to 0.
-        solver (LinearSystemSolvers, optional):
-            solver for Hx = g. Defaults to "cholesky_lu" (cholesky or LU if it fails).
-        fallback (FallbackLinearSystemSolvers, optional):
-            what to do if solver fails. Defaults to "safe_diag"
-            (takes nonzero diagonal elements, or fallbacks to gradient descent if all elements are 0).
-        max_norm (float, optional):
-            clips the newton step to L2 norm to avoid instability by giant steps.
-            A mauch better way is to use trust region methods. I haven't implemented any
-            but you can use `tz.optim.wrappers.scipy.ScipyMinimize` with one of the trust region methods.
-            Defaults to None.
-        validate (bool, optional):
-            validate if the step didn't increase the loss by `loss * tol` with an additional forward pass.
-            If not, undo the step and perform a gradient descent step.
-        tol (float, optional):
-            only has effect if `validate` is enabled.
-            If loss increased by `loss * tol`, perform gradient descent step.
-            Set this to 0 to guarantee that loss always decreases. Defaults to 1.
-        gd_lr (float, optional):
-            only has effect if `validate` is enabled.
-            Gradient descent step learning rate. Defaults to 1e-2.
-        line_search (OptimizerModule | None, optional): line search module, can be None. Defaults to None.
-        batched_hessian (bool, optional):
-            whether to use experimental pytorch vmap-vectorized hessian calculation. As per pytorch docs,
-            should be faster, but this feature being experimental, there may be performance cliffs.
-            Defaults to True.
-        diag (False, optional):
-            only use the diagonal of the hessian. This will still calculate the full hessian!
-            This is mainly useful for benchmarking.
-    """
-    def __init__(
-        self,
-        params,
-        lr: float = 1,
-        tikhonov: float | Literal['eig'] = 0.0,
-        solver: LinearSystemSolvers = "cholesky_lu",
-        fallback: FallbackLinearSystemSolvers = "safe_diag",
-        max_norm: float | None = None,
-        validate=False,
-        tol: float = 1,
-        gd_lr = 1e-2,
-        line_search: LineSearches | None = None,
-        batched_hessian = True,
-
-        diag: bool = False,
-    ):
-        modules: list[Any] = [
-            _ExactNewton(
-                tikhonov=tikhonov,
-                batched_hessian=batched_hessian,
-                solver=solver,
-                fallback=fallback,
-                validate=validate,
-                tol = tol,
-                gd_lr=gd_lr,
-                diag = diag,
-            ),
-        ]
-
-        if max_norm is not None:
-            modules.append(ClipNorm(max_norm))
-
-        modules.append(LR(lr))
-
-        if line_search is not None:
-            modules.append(get_line_search(line_search))
-
-        super().__init__(params, modules)
-
-

torchzero/optim/zeroth_order/__init__.py

@@ -1,4 +0,0 @@
-from .fdm import FDM, FDMWrapper
-from .newton_fdm import NewtonFDM, RandomSubspaceNewtonFDM
-from .rfdm import RandomGaussianSmoothing, RandomizedFDM, RandomizedFDMWrapper, SPSA
-from .rs import RandomSearch, CyclicRS

torchzero/optim/zeroth_order/fdm.py

@@ -1,87 +0,0 @@
-from typing import Literal
-
-import torch
-
-from ...modules import FDM as _FDM, WrapClosure, SGD, WeightDecay, LR
-from ...modules.gradient_approximation._fd_formulas import _FD_Formulas
-from ..modular import Modular
-
-
-class FDM(Modular):
-    """Gradient approximation via finite difference.
-
-    This performs `n + 1` evaluations per step with `forward` and `backward` formulas,
-    and `2 * n` with `central` formula, where n is the number of parameters.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float, optional): learning rate. Defaults to 1e-3.
-        eps (float, optional): finite difference epsilon. Defaults to 1e-3.
-        formula (_FD_Formulas, optional): finite difference formula. Defaults to "forward".
-        n_points (T.Literal[2, 3], optional): number of points for finite difference formula, 2 or 3. Defaults to 2.
-        momentum (float, optional): momentum. Defaults to 0.
-        dampening (float, optional): momentum dampening. Defaults to 0.
-        nesterov (bool, optional):
-            enables nesterov momentum, otherwise uses heavyball momentum. Defaults to False.
-        weight_decay (float, optional): weight decay (L2 regularization). Defaults to 0.
-        decoupled (bool, optional):
-            decouples weight decay from gradient. If True, weight decay doesn't depend on learning rate.
-    """
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-3,
-        eps: float = 1e-3,
-        formula: _FD_Formulas = "forward",
-        n_points: Literal[2, 3] = 2,
-        momentum: float = 0,
-        dampening: float = 0,
-        nesterov: bool = False,
-        weight_decay: float = 0,
-        decoupled=False,
-
-    ):
-        modules: list = [
-            _FDM(eps = eps, formula=formula, n_points=n_points),
-            SGD(momentum = momentum, dampening = dampening, weight_decay = weight_decay if not decoupled else 0, nesterov = nesterov),
-            LR(lr),
-
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-
-class FDMWrapper(Modular):
-    """Gradient approximation via finite difference. This wraps any other optimizer.
-    This also supports optimizers that perform multiple gradient evaluations per step, like LBFGS.
-
-    Exaple:
-    ```
-    lbfgs = torch.optim.LBFGS(params, lr = 1)
-    fdm = FDMWrapper(optimizer = lbfgs)
-    ```
-
-    This performs n+1 evaluations per step with `forward` and `backward` formulas,
-    and 2*n with `central` formula.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        optimizer (torch.optim.Optimizer): optimizer that will perform optimization using FDM-approximated gradients.
-        eps (float, optional): finite difference epsilon. Defaults to 1e-3.
-        formula (_FD_Formulas, optional): finite difference formula. Defaults to "forward".
-        n_points (T.Literal[2, 3], optional): number of points for finite difference formula, 2 or 3. Defaults to 2.
-    """
-    def __init__(
-        self,
-        optimizer: torch.optim.Optimizer,
-        eps: float = 1e-3,
-        formula: _FD_Formulas = "forward",
-        n_points: Literal[2, 3] = 2,
-    ):
-        modules = [
-            _FDM(eps = eps, formula=formula, n_points=n_points, target = 'closure'),
-            WrapClosure(optimizer)
-        ]
-        # some optimizers have `eps` setting in param groups too.
-        # it should not be passed to FDM
-        super().__init__([p for g in optimizer.param_groups.copy() for p in g['params']], modules)

torchzero/optim/zeroth_order/newton_fdm.py

@@ -1,146 +0,0 @@
-from typing import Any, Literal
-import torch
-
-from ...modules import (LR, FallbackLinearSystemSolvers,
-                        LinearSystemSolvers, LineSearches, ClipNorm)
-from ...modules import NewtonFDM as _NewtonFDM, get_line_search
-from ...modules.experimental.subspace import Proj2Masks, ProjRandom, Subspace
-from ..modular import Modular
-
-
-class NewtonFDM(Modular):
-    """Newton method with gradient and hessian approximated via finite difference.
-
-    This performs approximately `4 * n^2 + 1` evaluations per step;
-    if `diag` is True, performs `n * 2 + 1` evaluations per step.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float, optional): learning rate.
-        eps (float, optional): epsilon for finite difference.
-            Note that with float32 this needs to be quite high to avoid numerical instability. Defaults to 1e-2.
-        diag (bool, optional): whether to only approximate diagonal elements of the hessian.
-            This also ignores `solver` if True. Defaults to False.
-        solver (LinearSystemSolvers, optional):
-            solver for Hx = g. Defaults to "cholesky_lu" (cholesky or LU if it fails).
-        fallback (FallbackLinearSystemSolvers, optional):
-            what to do if solver fails. Defaults to "safe_diag"
-            (takes nonzero diagonal elements, or fallbacks to gradient descent if all elements are 0).
-        validate (bool, optional):
-            validate if the step didn't increase the loss by `loss * tol` with an additional forward pass.
-            If not, undo the step and perform a gradient descent step.
-        tol (float, optional):
-            only has effect if `validate` is enabled.
-            If loss increased by `loss * tol`, perform gradient descent step.
-            Set this to 0 to guarantee that loss always decreases. Defaults to 1.
-        gd_lr (float, optional):
-            only has effect if `validate` is enabled.
-            Gradient descent step learning rate. Defaults to 1e-2.
-        line_search (OptimizerModule | None, optional): line search module, can be None. Defaults to 'brent'.
-    """
-    def __init__(
-        self,
-        params,
-        lr: float = 1,
-        eps: float = 1e-2,
-        diag=False,
-        solver: LinearSystemSolvers = "cholesky_lu",
-        fallback: FallbackLinearSystemSolvers = "safe_diag",
-        max_norm: float | None = None,
-        validate=False,
-        tol: float = 2,
-        gd_lr = 1e-2,
-        line_search: LineSearches | None = 'brent',
-    ):
-        modules: list[Any] = [
-            _NewtonFDM(eps = eps, diag = diag, solver=solver, fallback=fallback, validate=validate, tol=tol, gd_lr=gd_lr),
-        ]
-
-        if max_norm is not None:
-            modules.append(ClipNorm(max_norm))
-
-        modules.append(LR(lr))
-
-        if line_search is not None:
-            modules.append(get_line_search(line_search))
-
-        super().__init__(params, modules)
-
-
-class RandomSubspaceNewtonFDM(Modular):
-    """This projects the parameters into a smaller dimensional subspace,
-    making approximating the hessian via finite difference feasible.
-
-    This performs approximately `4 * subspace_ndim^2 + 1` evaluations per step;
-    if `diag` is True, performs `subspace_ndim * 2 + 1` evaluations per step.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        subspace_ndim (float, optional): number of random subspace dimensions.
-        lr (float, optional): learning rate.
-        eps (float, optional): epsilon for finite difference.
-            Note that with float32 this needs to be quite high to avoid numerical instability. Defaults to 1e-2.
-        diag (bool, optional): whether to only approximate diagonal elements of the hessian.
-        solver (LinearSystemSolvers, optional):
-            solver for Hx = g. Defaults to "cholesky_lu" (cholesky or LU if it fails).
-        fallback (FallbackLinearSystemSolvers, optional):
-            what to do if solver fails. Defaults to "safe_diag"
-            (takes nonzero diagonal elements, or fallbacks to gradient descent if all elements are 0).
-        validate (bool, optional):
-            validate if the step didn't increase the loss by `loss * tol` with an additional forward pass.
-            If not, undo the step and perform a gradient descent step.
-        tol (float, optional):
-            only has effect if `validate` is enabled.
-            If loss increased by `loss * tol`, perform gradient descent step.
-            Set this to 0 to guarantee that loss always decreases. Defaults to 1.
-        gd_lr (float, optional):
-            only has effect if `validate` is enabled.
-            Gradient descent step learning rate. Defaults to 1e-2.
-        line_search (OptimizerModule | None, optional): line search module, can be None. Defaults to BacktrackingLS().
-        randomize_every (float, optional): generates new random projections every n steps. Defaults to 1.
-    """
-    def __init__(
-        self,
-        params,
-        subspace_ndim: int = 3,
-        lr: float = 1,
-        eps: float = 1e-2,
-        diag=False,
-        solver: LinearSystemSolvers = "cholesky_lu",
-        fallback: FallbackLinearSystemSolvers = "safe_diag",
-        max_norm: float | None = None,
-        validate=False,
-        tol: float = 2,
-        gd_lr = 1e-2,
-        line_search: LineSearches | None = 'brent',
-        randomize_every: int = 1,
-    ):
-        if subspace_ndim == 1: projections = [ProjRandom(1)]
-        else:
-            projections: list[Any] = [Proj2Masks(subspace_ndim//2)]
-            if subspace_ndim % 2 == 1: projections.append(ProjRandom(1))
-
-        modules: list[Any] = [
-            Subspace(
-                modules = _NewtonFDM(
-                    eps = eps,
-                    diag = diag,
-                    solver=solver,
-                    fallback=fallback,
-                    validate=validate,
-                    tol=tol,
-                    gd_lr=gd_lr
-                ),
-                projections = projections,
-                update_every=randomize_every),
-        ]
-        if max_norm is not None:
-            modules.append(ClipNorm(max_norm))
-
-        modules.append(LR(lr))
-
-        if line_search is not None:
-            modules.append(get_line_search(line_search))
-
-        super().__init__(params, modules)
-