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

torchzero/optim/experimental/experimental.py

@@ -1,343 +0,0 @@
-from typing import Literal
-
-from ...modules import (
-    LR,
-    SGD,
-    Abs,
-    Adam,
-    Add,
-    AddMagnitude,
-    Cautious,
-    Div,
-    Divide,
-    Grad,
-    HeavyBall,
-    Interpolate,
-    Lerp,
-    Multistep,
-    NanToNum,
-    NesterovMomentum,
-    Normalize,
-    Random,
-    RDiv,
-    Reciprocal,
-    UseGradSign,
-    WeightDecay,
-)
-from ...modules import RandomCoordinateMomentum as _RandomCoordinateMomentum
-from ...modules.experimental import GradMin as _GradMin
-from ...modules.experimental import (
-    HVPDiagNewton as _HVPDiagNewton,
-)
-from ...modules.experimental import MinibatchRprop as _MinibatchRprop
-from ...modules.experimental import ReduceOutwardLR
-from ...random import Distributions
-from ..modular import Modular
-
-
-class HVPDiagNewton(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - this should approximate newton method with 2 backward passes, but only if hessian is purely diagonal"""
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-1,
-        eps: float = 1e-2,
-    ):
-        modules = [_HVPDiagNewton(eps = eps), LR(lr)]
-        super().__init__(params, modules)
-
-
-class ReciprocalSGD(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - this basically uses normalized *1 / (gradient + eps)*."""
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-2,
-        eps: float = 1e-2,
-        momentum: float = 0,
-        dampening: float = 0,
-        nesterov: bool = False,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list = [
-            AddMagnitude(eps, add_to_zero=False),
-            Reciprocal(),
-            NanToNum(0,0,0),
-            Normalize(1),
-            SGD(momentum = momentum, dampening = dampening, weight_decay = 0, nesterov = nesterov),
-            LR(lr),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-
-        super().__init__(params, modules)
-
-class NoiseSign(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - uses random vector with gradient sign, and works quite well despite being completely random."""
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-2,
-        distribution: Distributions = 'normal',
-        momentum: float = 0,
-        dampening: float = 0,
-        nesterov: bool = False,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list = [
-            Random(1, distribution),
-            UseGradSign(),
-            SGD(momentum = momentum, dampening = dampening, weight_decay = 0, nesterov = nesterov),
-            LR(lr),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(2, WeightDecay(weight_decay))
-
-        super().__init__(params, modules)
-
-class MomentumNumerator(Modular):
-    """for experiments, unlikely to work well on most problems. (somewhat promising)
-
-    explanation - momentum divided by gradient."""
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-2,
-        momentum: float = 0.9,
-        nesterov: bool = True,
-        eps: float = 1e-2,
-        weight_decay: float = 0,
-        decoupled=True,    ):
-
-        modules: list = [
-            Divide(
-                numerator = SGD(momentum = momentum, nesterov=nesterov),
-                denominator=[Abs(), Add(eps)]
-            ),
-            Normalize(),
-            LR(lr),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-class MomentumDenominator(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - gradient divided by normalized momentum."""
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-2,
-        momentum: float = 0.9,
-        nesterov: bool = True,
-        eps: float = 1e-2,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list = [
-            Div([SGD(momentum=momentum, nesterov=nesterov), Abs(), Add(eps), Normalize(1)]),
-            Normalize(),
-            LR(lr),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-
-class ExaggeratedNesterov(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - exaggerates difference between heavyball and nesterov momentum."""
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-2,
-        momentum: float = 0.9,
-        dampening: float = 0,
-        strength: float = 5,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-
-        modules: list = [
-            Interpolate(HeavyBall(momentum, dampening), NesterovMomentum(momentum, dampening), strength),
-            LR(lr),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-class ExtraCautiousAdam(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - caution with true backtracking."""
-    def __init__(
-        self,
-        params,
-        lr: float = 1,
-        beta1: float = 0.9,
-        beta2: float = 0.999,
-        eps: float = 1e-8,
-        amsgrad=False,
-        normalize = False,
-        c_eps = 1e-6,
-        mode: Literal['zero', 'grad', 'backtrack'] = 'zero',
-        strength = 5,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list = [
-            Adam(beta1, beta2, eps, amsgrad=amsgrad),
-            Lerp(Cautious(normalize, c_eps, mode), strength),
-            LR(lr),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-class InwardSGD(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - reduces lrs for updates that move weights away from 0."""
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-3,
-        momentum: float = 0,
-        dampening: float = 0,
-        nesterov: bool = False,
-        mul = 0.5,
-        use_grad=False,
-        invert=False,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list = [
-            SGD(momentum = momentum, dampening = dampening, weight_decay = 0, nesterov = nesterov),
-            LR(lr),
-            ReduceOutwardLR(mul, use_grad, invert),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-class MultistepSGD(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - perform multiple steps per batch. Momentum applies to the total update over multiple step"""
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-3,
-        momentum: float = 0,
-        dampening: float = 0,
-        nesterov: bool = False,
-        num_steps=2,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        # lr, lr_module = _get_baked_in_and_module_lr(lr, kwargs) # multistep must use lr
-
-        modules: list = [
-            Multistep(LR(lr), num_steps=num_steps),
-            SGD(momentum = momentum, dampening = dampening, weight_decay = 0, nesterov = nesterov),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-
-class MinibatchRprop(Modular):
-    """
-    for experiments, unlikely to work well on most problems.
-
-    explanation: does 2 steps per batch, applies rprop rule on the second step.
-    """
-    def __init__(
-        self,
-        params,
-        lr: float = 1,
-        nplus: float = 1.2,
-        nminus: float = 0.5,
-        lb: float | None = 1e-6,
-        ub: float | None = 50,
-        backtrack=True,
-        next_mode = 'continue',
-        increase_mul = 0.5,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list = [
-            _MinibatchRprop(nplus=nplus,nminus=nminus,lb=lb,ub=ub,backtrack=backtrack,next_mode=next_mode,increase_mul=increase_mul),
-            LR(lr),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-
-class RandomCoordinateMomentum(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    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:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float): learning rate (default: 1e-3).
-        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,
-        params,
-        lr: float = 1e-3,
-        p: float = 0.1,
-        nesterov: bool = True,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list = [_RandomCoordinateMomentum(p, nesterov), LR(lr)]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-class GradMin(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - this uses gradient wrt sum of gradients + loss."""
-
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-2,
-        loss_term: float = 1,
-        square: bool = False,
-        maximize_grad: bool = False,
-        momentum: float = 0,
-        dampening: float = 0,
-        nesterov: bool = False,
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list = [
-            _GradMin(loss_term, square, maximize_grad),
-            SGD(momentum = momentum, dampening = dampening, weight_decay = 0, nesterov = nesterov),
-            LR(lr),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-

torchzero/optim/experimental/ray_search.py

@@ -1,83 +0,0 @@
-from typing import Literal, Any
-
-import torch
-
-from ...core import OptimizerModule
-from ...modules import (SGD, LineSearches, NewtonFDM,
-                        get_line_search, LR, WrapClosure)
-from ...modules.experimental.subspace import Subspace, ProjNormalize, ProjAscentRay
-from ..modular import Modular
-
-
-class NewtonFDMRaySearch(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - like a fancy line search, instead of a line searches in a cone using FDM newton."""
-    def __init__(
-        self,
-        params,
-        lr = 1e-2,
-        momentum:float = 0,
-        weight_decay:float = 0,
-        dampening: float = 0,
-        nesterov:bool = False,
-        n_rays = 3,
-        eps = 1e-2,
-        ray_width: float = 1e-1,
-        line_search: LineSearches | None = 'brent'
-    ):
-        modules: list[Any] = [
-            SGD(momentum=momentum, weight_decay=weight_decay, dampening=dampening, nesterov=nesterov),
-            LR(lr),
-            Subspace(NewtonFDM(eps = eps), ProjNormalize(ProjAscentRay(ray_width, n = n_rays))),
-        ]
-        if lr != 1:
-            modules.append(LR(lr))
-
-        if line_search is not None:
-            modules.append(get_line_search(line_search))
-
-        super().__init__(params, modules)
-
-
-class LBFGSRaySearch(Modular):
-    """for experiments, unlikely to work well on most problems.
-
-    explanation - like a fancy line search, instead of a line searches in a cone using LBFGS."""
-    def __init__(
-        self,
-        params,
-        lr = 1,
-        momentum:float = 0,
-        weight_decay:float = 0,
-        dampening: float = 0,
-        nesterov:bool = False,
-        n_rays = 24,
-        ray_width: float = 1e-1,
-        max_iter: int = 20,
-        max_eval: int | None = None,
-        tolerance_grad: float = 1e-7,
-        tolerance_change: float = 1e-9,
-        history_size: int = 100,
-        line_search_fn: str | Literal['strong_wolfe'] | None = None,
-    ):
-        lbfgs = WrapClosure(
-            torch.optim.LBFGS,
-            lr=lr,
-            max_iter=max_iter,
-            max_eval=max_eval,
-            tolerance_grad=tolerance_grad,
-            tolerance_change=tolerance_change,
-            history_size=history_size,
-            line_search_fn=line_search_fn,
-        )
-        modules: list[OptimizerModule] = [
-            SGD(momentum=momentum, weight_decay=weight_decay, dampening=dampening, nesterov=nesterov),
-            Subspace(lbfgs, ProjNormalize(ProjAscentRay(ray_width, n = n_rays))),
-
-        ]
-
-        super().__init__(params, modules)
-
-
-

torchzero/optim/first_order/__init__.py

@@ -1,18 +0,0 @@
-from .cautious import CautiousAdamW, CautiousLion, CautiousSGD
-from .optimizers import (
-    GD,
-    SGD,
-    Adagrad,
-    Adam,
-    AdamW,
-    Grams,
-    LaplacianSmoothingSGD,
-    Lion,
-    NestedNesterov,
-    NoisySGD,
-    NormSGD,
-    RMSProp,
-    Rprop,
-    SignSGD,
-)
-from .forward_gradient import ForwardGradient

torchzero/optim/first_order/cautious.py

@@ -1,158 +0,0 @@
-from typing import Literal
-
-
-from ...core import OptimizerModule
-from ...modules import Cautious, Adam, SGD, Lion, WeightDecay, LR
-from ..modular import Modular
-
-
-class CautiousAdamW(Modular):
-    """Adam, but updates for parameters where update and gradient sign is inconsistent are negated.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float): learning rate (default: 1e-3).
-        beta1 (float, optional): exponential decay rate of gradient moving average. Defaults to 0.9.
-        beta2 (float, optional): exponential decay rate of squared gradient moving average. Defaults to 0.999.
-        eps (float, optional): epsilon for numerical stability. Defaults to 1e-8.
-        amsgrad (bool, optional):
-            whether to use the AMSGrad variant of this algorithm from
-            On the Convergence of Adam and Beyond (default: False).
-        normalize (bool, optional):
-            renormalize update after masking.
-            only has effect when mode is 'zero'. Defaults to False.
-        c_eps (float, optional): epsilon for normalization after applying cautioning mask. 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
-
-            "negate" - negate them (same as using update magnitude and gradient sign).
-        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,
-        beta1: float = 0.9,
-        beta2: float = 0.999,
-        eps: float = 1e-8,
-        amsgrad=False,
-        normalize = False,
-        c_eps = 1e-6,
-        mode: Literal['zero', 'grad', 'backtrack'] = 'zero',
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list[OptimizerModule] = [
-            Adam(beta1 = beta1, beta2 = beta2, eps = eps, amsgrad = amsgrad),
-            LR(lr),
-            Cautious(normalize = normalize, eps = c_eps, mode = mode),
-        ]
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-        super().__init__(params, modules)
-
-
-class CautiousSGD(Modular):
-    """SGD with momentum, but updates for parameters where update and gradient sign is inconsistent are negated.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float): learning rate (default: 1e-3).
-        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.
-        normalize (bool, optional):
-            renormalize update after masking.
-            only has effect when mode is 'zero'. Defaults to False.
-        c_eps (float, optional): epsilon for normalization after applying cautioning mask. 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
-
-            "negate" - negate them (same as using update magnitude and gradient sign).
-        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,
-        momentum: float = 0.9,
-        dampening: float = 0,
-        nesterov: bool = True,
-        c_eps = 1e-6,
-        normalize = False,
-        mode: Literal['zero', 'grad', 'backtrack'] = 'zero',
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list[OptimizerModule] = [
-            SGD(momentum = momentum, dampening = dampening, weight_decay = 0, nesterov = nesterov),
-            LR(lr),
-            Cautious(normalize = normalize, eps = c_eps, mode = mode),
-        ]
-
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-
-        super().__init__(params, modules)
-
-
-class CautiousLion(Modular):
-    """Lion optimizer, but updates for parameters where update and gradient sign is inconsistent are negated.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float): learning rate (default: 1e-3).
-        beta1 (float, optional): dampening for momentum. Defaults to 0.9.
-        beta2 (float, optional): momentum factor. Defaults to 0.99.
-        normalize (bool, optional):
-            renormalize update after masking.
-            only has effect when mode is 'zero'. Defaults to False.
-        c_eps (float, optional): epsilon for normalization after applying cautioning mask. 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
-
-            "negate" - negate them (same as using update magnitude and gradient sign).
-        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,
-        beta1: float = 0.9,
-        beta2: float = 0.99,
-        c_eps = 1e-6,
-        normalize = False,
-        mode: Literal['zero', 'grad', 'backtrack'] = 'zero',
-        weight_decay: float = 0,
-        decoupled=True,
-    ):
-        modules: list[OptimizerModule] = [
-            Lion(beta1, beta2),
-            LR(lr),
-            Cautious(normalize = normalize, eps = c_eps, mode = mode),
-        ]
-
-        if decoupled: modules.append(WeightDecay(weight_decay))
-        else: modules.insert(0, WeightDecay(weight_decay))
-
-        super().__init__(params, modules)

torchzero/optim/first_order/forward_gradient.py

@@ -1,70 +0,0 @@
-from typing import Literal
-
-import torch
-
-from ...core import OptimizationVars, OptimizerModule
-from ...modules import ForwardGradient as _ForwardGradient, SGD, WeightDecay, LR
-from ...tensorlist import Distributions
-from ..modular import Modular
-
-
-class ForwardGradient(Modular):
-    """
-
-    Evaluates jacobian-vector product with a random vector using forward mode autodiff (torch.func.jvp), which is
-    the true directional derivative in the direction of that vector.
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        lr (float, optional): learning rate. Defaults to 1e-3.
-        n_samples (int): number of forward gradients to evaluate and average.
-        distribution (Distributions): distribution for random tangent vector.
-        mode (str):
-            "jvp" - uses forward mode AD, usually slightly slower than backward mode AD  but uses significantly less memory.
-
-            "grad" - evaluates gradient with `loss.backward()` which may be faster but uses all the memory, mainly useful for
-            benchmarking as there is probably no point in forward gradient if full gradient is available.
-
-            "fd" - uses finite difference to estimate JVP, doesn't require gradients to be known. Equivalent to randomized FDM.
-
-        fd_eps (float, optional): epsilon for finite difference, only has effect if mode is "fd". Defaults to 1e-4.
-        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.
-
-    Reference:
-        Baydin, A. G., Pearlmutter, B. A., Syme, D., Wood, F., & Torr, P. (2022).
-        Gradients without backpropagation. arXiv preprint arXiv:2202.08587.
-        https://arxiv.org/abs/2202.08587
-    """
-
-    def __init__(
-        self,
-        params,
-        lr: float = 1e-3,
-        n_samples: int = 1,
-        distribution: Distributions = "normal",
-        mode: Literal["jvp", "grad", "fd"] = "jvp",
-        fd_eps: float = 1e-4,
-        momentum: float = 0,
-        dampening: float = 0,
-        nesterov: bool = False,
-        weight_decay: float = 0,
-        decoupled=False,
-    ):
-        modules: list = [
-            _ForwardGradient(
-                n_samples=n_samples,
-                distribution=distribution,
-                mode=mode,
-                fd_eps=fd_eps,
-            ),
-            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)