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

torchzero/modules/grad_approximation/grad_approximator.py

@@ -0,0 +1,66 @@
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
+from typing import Any, Literal
+
+import torch
+
+from ...core import Module, Vars
+
+GradTarget = Literal['update', 'grad', 'closure']
+_Scalar = torch.Tensor | float
+
+class GradApproximator(Module, ABC):
+    """Base class for gradient approximations.
+    This is an abstract class, to use it, subclass it and override `approximate`.
+
+    Args:
+        defaults (dict[str, Any] | None, optional): dict with defaults. Defaults to None.
+        target (str, optional):
+            whether to set `vars.grad`, `vars.update` or 'vars.closure`. Defaults to 'closure'.
+    """
+    def __init__(self, defaults: dict[str, Any] | None = None, target: GradTarget = 'closure'):
+        super().__init__(defaults)
+        self._target: GradTarget = target
+
+    @abstractmethod
+    def approximate(self, closure: Callable, params: list[torch.Tensor], loss: _Scalar | None, vars: Vars) -> tuple[Iterable[torch.Tensor], _Scalar | None, _Scalar | None]:
+        """Returns a tuple: (grad, loss, loss_approx), make sure this resets parameters to their original values!"""
+
+    def pre_step(self, vars: Vars) -> Vars | None:
+        """This runs once before each step, whereas `approximate` may run multiple times per step if further modules
+        evaluate gradients at multiple points. This is useful for example to pre-generate new random perturbations."""
+        return vars
+
+    @torch.no_grad
+    def step(self, vars):
+        ret = self.pre_step(vars)
+        if isinstance(ret, Vars): vars = ret
+
+        if vars.closure is None: raise RuntimeError("Gradient approximation requires closure")
+        params, closure, loss = vars.params, vars.closure, vars.loss
+
+        if self._target == 'closure':
+
+            def approx_closure(backward=True):
+                if backward:
+                    # set loss to None because closure might be evaluated at different points
+                    grad, l, l_approx = self.approximate(closure=closure, params=params, loss=None, vars=vars)
+                    for p, g in zip(params, grad): p.grad = g
+                    return l if l is not None else l_approx
+                return closure(False)
+
+            vars.closure = approx_closure
+            return vars
+
+        # if vars.grad is not None:
+        #     warnings.warn('Using grad approximator when `vars.grad` is already set.')
+        grad,loss,loss_approx = self.approximate(closure=closure, params=params, loss=loss, vars=vars)
+        if loss_approx is not None: vars.loss_approx = loss_approx
+        if loss is not None: vars.loss = vars.loss_approx = loss
+        if self._target == 'grad': vars.grad = list(grad)
+        elif self._target == 'update': vars.update = list(grad)
+        else: raise ValueError(self._target)
+        return vars
+
+_FD_Formula = Literal['forward2', 'backward2', 'forward3', 'backward3', 'central2', 'central4']

torchzero/modules/grad_approximation/rfdm.py

@@ -0,0 +1,259 @@
+from collections.abc import Callable
+from typing import Any
+from functools import partial
+import torch
+
+from ...utils import TensorList, Distributions, NumberList, generic_eq
+from .grad_approximator import GradApproximator, GradTarget, _FD_Formula
+
+
+def _rforward2(closure: Callable[..., float], params:TensorList, p_fn:Callable[[], TensorList], h, v_0: float | None):
+    """p_fn is a function that returns the perturbation.
+    It may return pre-generated one or generate one deterministically from a seed as in MeZO.
+    Returned perturbation must be multiplied by `h`."""
+    if v_0 is None: v_0 = closure(False)
+    params += p_fn()
+    v_plus = closure(False)
+    params -= p_fn()
+    h = h**2 # because perturbation already multiplied by h
+    return v_0, v_0, (v_plus - v_0) / h # (loss, loss_approx, grad)
+
+def _rbackward2(closure: Callable[..., float], params:TensorList, p_fn:Callable[[], TensorList], h, v_0: float | None):
+    if v_0 is None: v_0 = closure(False)
+    params -= p_fn()
+    v_minus = closure(False)
+    params += p_fn()
+    h = h**2 # because perturbation already multiplied by h
+    return v_0, v_0, (v_0 - v_minus) / h
+
+def _rcentral2(closure: Callable[..., float], params:TensorList, p_fn:Callable[[], TensorList], h, v_0: Any):
+    params += p_fn()
+    v_plus = closure(False)
+
+    params -= p_fn() * 2
+    v_minus = closure(False)
+
+    params += p_fn()
+    h = h**2 # because perturbation already multiplied by h
+    return v_0, v_plus, (v_plus - v_minus) / (2 * h)
+
+def _rforward3(closure: Callable[..., float], params:TensorList, p_fn:Callable[[], TensorList], h, v_0: float | None):
+    if v_0 is None: v_0 = closure(False)
+    params += p_fn()
+    v_plus1 = closure(False)
+
+    params += p_fn()
+    v_plus2 = closure(False)
+
+    params -= p_fn() * 2
+    h = h**2 # because perturbation already multiplied by h
+    return v_0, v_0, (-3*v_0 + 4*v_plus1 - v_plus2) / (2 * h)
+
+def _rbackward3(closure: Callable[..., float], params:TensorList, p_fn:Callable[[], TensorList], h, v_0: float | None):
+    if v_0 is None: v_0 = closure(False)
+
+    params -= p_fn()
+    v_minus1 = closure(False)
+
+    params -= p_fn()
+    v_minus2 = closure(False)
+
+    params += p_fn() * 2
+    h = h**2 # because perturbation already multiplied by h
+    return v_0, v_0, (v_minus2 - 4*v_minus1 + 3*v_0) / (2 * h)
+
+def _rcentral4(closure: Callable[..., float], params:TensorList, p_fn:Callable[[], TensorList], h, v_0: float | None):
+    params += p_fn()
+    v_plus1 = closure(False)
+
+    params += p_fn()
+    v_plus2 = closure(False)
+
+    params -= p_fn() * 3
+    v_minus1 = closure(False)
+
+    params -= p_fn()
+    v_minus2 = closure(False)
+
+    params += p_fn() * 2
+    h = h**2 # because perturbation already multiplied by h
+    return v_0, v_plus1, (v_minus2 - 8*v_minus1 + 8*v_plus1 - v_plus2) / (12 * h)
+
+_RFD_FUNCS = {
+    "forward2": _rforward2,
+    "backward2": _rbackward2,
+    "central2": _rcentral2,
+    "forward3": _rforward3,
+    "backward3": _rbackward3,
+    "central4": _rcentral4,
+}
+
+
+class RandomizedFDM(GradApproximator):
+    PRE_MULTIPLY_BY_H = True
+    def __init__(
+        self,
+        h: float = 1e-3,
+        n_samples: int = 1,
+        formula: _FD_Formula = "central2",
+        distribution: Distributions = "rademacher",
+        beta: float = 0,
+        pre_generate = True,
+        target: GradTarget = "closure",
+        seed: int | None | torch.Generator = None,
+    ):
+        defaults = dict(h=h, formula=formula, n_samples=n_samples, distribution=distribution, beta=beta, pre_generate=pre_generate, seed=seed)
+        super().__init__(defaults, target=target)
+
+    def reset(self):
+        self.state.clear()
+        generator = self.global_state.get('generator', None) # avoid resetting generator
+        self.global_state.clear()
+        if generator is not None: self.global_state['generator'] = generator
+
+    def _get_generator(self, seed: int | None | torch.Generator, params: list[torch.Tensor]):
+        if 'generator' not in self.global_state:
+            if isinstance(seed, torch.Generator): self.global_state['generator'] = seed
+            elif seed is not None: self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
+            else: self.global_state['generator'] = None
+        return self.global_state['generator']
+
+    def pre_step(self, vars):
+        h, beta = self.get_settings('h', 'beta', params=vars.params)
+        settings = self.settings[vars.params[0]]
+        n_samples = settings['n_samples']
+        distribution = settings['distribution']
+        pre_generate = settings['pre_generate']
+
+        if pre_generate:
+            params = TensorList(vars.params)
+            generator = self._get_generator(settings['seed'], vars.params)
+            perturbations = [params.sample_like(distribution=distribution, generator=generator) for _ in range(n_samples)]
+
+            if self.PRE_MULTIPLY_BY_H:
+                torch._foreach_mul_([p for l in perturbations for p in l], [v for vv in h for v in [vv]*n_samples])
+
+            if all(i==0 for i in beta):
+                # just use pre-generated perturbations
+                for param, prt in zip(params, zip(*perturbations)):
+                    self.state[param]['perturbations'] = prt
+
+            else:
+                # lerp old and new perturbations. This makes the subspace change gradually
+                # which in theory might improve algorithms with history
+                for i,p in enumerate(params):
+                    state = self.state[p]
+                    if 'perturbations' not in state: state['perturbations'] = [p[i] for p in perturbations]
+
+                cur = [self.state[p]['perturbations'][:n_samples] for p in params]
+                cur_flat = [p for l in cur for p in l]
+                new_flat = [p for l in zip(*perturbations) for p in l]
+                betas = [1-v for b in beta for v in [b]*n_samples]
+                torch._foreach_lerp_(cur_flat, new_flat, betas)
+
+    @torch.no_grad
+    def approximate(self, closure, params, loss, vars):
+        params = TensorList(params)
+        loss_approx = None
+
+        h = self.get_settings('h', params=vars.params, cls=NumberList)
+        settings = self.settings[params[0]]
+        n_samples = settings['n_samples']
+        fd_fn = _RFD_FUNCS[settings['formula']]
+        default = [None]*n_samples
+        perturbations = list(zip(*(self.state[p].get('perturbations', default) for p in params)))
+        distribution = settings['distribution']
+        generator = self._get_generator(settings['seed'], params)
+
+        grad = None
+        for i in range(n_samples):
+            prt = perturbations[i]
+            if prt[0] is None: prt = params.sample_like(distribution=distribution, generator=generator).mul_(h)
+            else: prt = TensorList(prt)
+
+            loss, loss_approx, d = fd_fn(closure=closure, params=params, p_fn=lambda: prt, h=h, v_0=loss)
+            if grad is None: grad = prt * d
+            else: grad += prt * d
+
+        assert grad is not None
+        if n_samples > 1: grad.div_(n_samples)
+        return grad, loss, loss_approx
+
+SPSA = RandomizedFDM
+
+class RDSA(RandomizedFDM):
+    def __init__(
+        self,
+        h: float = 1e-3,
+        n_samples: int = 1,
+        formula: _FD_Formula = "central2",
+        distribution: Distributions = "gaussian",
+        beta: float = 0,
+        pre_generate = True,
+        target: GradTarget = "closure",
+        seed: int | None | torch.Generator = None,
+    ):
+        super().__init__(h=h, n_samples=n_samples,formula=formula,distribution=distribution,beta=beta,pre_generate=pre_generate,target=target,seed=seed)
+
+class GaussianSmoothing(RandomizedFDM):
+    def __init__(
+        self,
+        h: float = 1e-2,
+        n_samples: int = 100,
+        formula: _FD_Formula = "central2",
+        distribution: Distributions = "gaussian",
+        beta: float = 0,
+        pre_generate = True,
+        target: GradTarget = "closure",
+        seed: int | None | torch.Generator = None,
+    ):
+        super().__init__(h=h, n_samples=n_samples,formula=formula,distribution=distribution,beta=beta,pre_generate=pre_generate,target=target,seed=seed)
+
+class MeZO(GradApproximator):
+    def __init__(self, h: float=1e-3, n_samples: int = 1, formula: _FD_Formula = 'central2',
+                 distribution: Distributions = 'rademacher', target: GradTarget = 'closure'):
+        defaults = dict(h=h, formula=formula, n_samples=n_samples, distribution=distribution)
+        super().__init__(defaults, target=target)
+
+    def _seeded_perturbation(self, params: list[torch.Tensor], distribution, seed, h):
+        return TensorList(params).sample_like(
+            distribution=distribution, generator=torch.Generator(params[0].device).manual_seed(seed)
+        ).mul_(h)
+
+    def pre_step(self, vars):
+        h = self.get_settings('h', params=vars.params)
+        settings = self.settings[vars.params[0]]
+        n_samples = settings['n_samples']
+        distribution = settings['distribution']
+
+        step = vars.current_step
+
+        # create functions that generate a deterministic perturbation from seed based on current step
+        prt_fns = []
+        for i in range(n_samples):
+
+            prt_fn = partial(self._seeded_perturbation, params=vars.params, distribution=distribution, seed=1_000_000*step + i, h=h)
+            prt_fns.append(prt_fn)
+
+        self.global_state['prt_fns'] = prt_fns
+
+    @torch.no_grad
+    def approximate(self, closure, params, loss, vars):
+        params = TensorList(params)
+        loss_approx = None
+
+        h = self.get_settings('h', params=vars.params, cls=NumberList)
+        settings = self.settings[params[0]]
+        n_samples = settings['n_samples']
+        fd_fn = _RFD_FUNCS[settings['formula']]
+        prt_fns = self.global_state['prt_fns']
+
+        grad = None
+        for i in range(n_samples):
+            loss, loss_approx, d = fd_fn(closure=closure, params=params, p_fn=prt_fns[i], h=h, v_0=loss)
+            if grad is None: grad = prt_fns[i]().mul_(d)
+            else: grad += prt_fns[i]().mul_(d)
+
+        assert grad is not None
+        if n_samples > 1: grad.div_(n_samples)
+        return grad, loss, loss_approx

torchzero/modules/line_search/__init__.py

@@ -1,30 +1,5 @@
-r"""
-Line searches.
-"""
-
-from typing import Literal
-
-from ...core import OptimizerModule
-from ..regularization import Normalize
-from .grid_ls import (ArangeLS, BacktrackingLS, GridLS, LinspaceLS,
-                      MultiplicativeLS)
-# from .quad_interp import QuadraticInterpolation2Point
-from .directional_newton import DirectionalNewton3Points, DirectionalNewton
-from .scipy_minimize_scalar import ScipyMinimizeScalarLS
-from .armijo import ArmijoLS
-
-LineSearches = Literal['backtracking', 'brent', 'brent-exact', 'brent-norm', 'multiplicative', 'newton', 'newton3', 'armijo'] | OptimizerModule
-
-def get_line_search(name:str | OptimizerModule) -> OptimizerModule | list[OptimizerModule]:
-    if isinstance(name, str):
-        name = name.strip().lower()
-        if name == 'backtracking': return BacktrackingLS()
-        if name == 'multiplicative': return MultiplicativeLS()
-        if name == 'brent': return ScipyMinimizeScalarLS(maxiter=8)
-        if name == 'brent-exact': return ScipyMinimizeScalarLS()
-        if name == 'brent-norm': return [Normalize(), ScipyMinimizeScalarLS(maxiter=16)]
-        if name == 'newton': return DirectionalNewton(1)
-        if name == 'newton3': return DirectionalNewton3Points(1)
-        if name == 'armijo': return ArmijoLS(1)
-        raise ValueError(f"Unknown line search method: {name}")
-    return name
+from .line_search import LineSearch, GridLineSearch
+from .backtracking import backtracking_line_search, Backtracking, AdaptiveBacktracking
+from .strong_wolfe import StrongWolfe
+from .scipy import ScipyMinimizeScalar
+from .trust_region import TrustRegion

torchzero/modules/line_search/backtracking.py

@@ -0,0 +1,186 @@
+import math
+from collections.abc import Callable
+from operator import itemgetter
+
+import torch
+
+from .line_search import LineSearch
+
+
+def backtracking_line_search(
+    f: Callable[[float], float],
+    g_0: float | torch.Tensor,
+    init: float = 1.0,
+    beta: float = 0.5,
+    c: float = 1e-4,
+    maxiter: int = 10,
+    a_min: float | None = None,
+    try_negative: bool = False,
+) -> float | None:
+    """
+
+    Args:
+        objective_fn: evaluates step size along some descent direction.
+        dir_derivative: directional derivative along the descent direction.
+        alpha_init: initial step size.
+        beta: The factor by which to decrease alpha in each iteration
+        c: The constant for the Armijo sufficient decrease condition
+        max_iter: Maximum number of backtracking iterations (default: 10).
+        min_alpha: Minimum allowable step size to prevent near-zero values (default: 1e-16).
+
+    Returns:
+        step size
+    """
+
+    a = init
+    f_x = f(0)
+
+    for iteration in range(maxiter):
+        f_a = f(a)
+
+        if f_a <= f_x + c * a * min(g_0, 0): # pyright: ignore[reportArgumentType]
+            # found an acceptable alpha
+            return a
+
+        # decrease alpha
+        a *= beta
+
+        # alpha too small
+        if a_min is not None and a < a_min:
+            return a_min
+
+    # fail
+    if try_negative:
+        def inv_objective(alpha): return f(-alpha)
+
+        v = backtracking_line_search(
+            inv_objective,
+            g_0=-g_0,
+            beta=beta,
+            c=c,
+            maxiter=maxiter,
+            a_min=a_min,
+            try_negative=False,
+        )
+        if v is not None: return -v
+
+    return None
+
+class Backtracking(LineSearch):
+    def __init__(
+        self,
+        init: float = 1.0,
+        beta: float = 0.5,
+        c: float = 1e-4,
+        maxiter: int = 10,
+        min_alpha: float | None = None,
+        adaptive=True,
+        try_negative: bool = False,
+    ):
+        defaults=dict(init=init,beta=beta,c=c,maxiter=maxiter,min_alpha=min_alpha,adaptive=adaptive, try_negative=try_negative)
+        super().__init__(defaults=defaults)
+        self.global_state['beta_scale'] = 1.0
+
+    def reset(self):
+        super().reset()
+        self.global_state['beta_scale'] = 1.0
+
+    @torch.no_grad
+    def search(self, update, vars):
+        init, beta, c, maxiter, min_alpha, adaptive, try_negative = itemgetter(
+            'init', 'beta', 'c', 'maxiter', 'min_alpha', 'adaptive', 'try_negative')(self.settings[vars.params[0]])
+
+        objective = self.make_objective(vars=vars)
+
+        # # directional derivative
+        d = -sum(t.sum() for t in torch._foreach_mul(vars.get_grad(), vars.get_update()))
+
+        # scale beta (beta is multiplicative and i think may be better than scaling initial step size)
+        if adaptive: beta = beta * self.global_state['beta_scale']
+
+        step_size = backtracking_line_search(objective, d, init=init,beta=beta,
+                                        c=c,maxiter=maxiter,a_min=min_alpha, try_negative=try_negative)
+
+        # found an alpha that reduces loss
+        if step_size is not None:
+            self.global_state['beta_scale'] = min(1.0, self.global_state['beta_scale'] * math.sqrt(1.5))
+            return step_size
+
+        # on fail reduce beta scale value
+        self.global_state['beta_scale'] /= 1.5
+        return 0
+
+def _lerp(start,end,weight):
+    return start + weight * (end - start)
+
+class AdaptiveBacktracking(LineSearch):
+    def __init__(
+        self,
+        init: float = 1.0,
+        beta: float = 0.5,
+        c: float = 1e-4,
+        maxiter: int = 20,
+        min_alpha: float | None = None,
+        target_iters = 1,
+        nplus = 2.0,
+        scale_beta = 0.0,
+        try_negative: bool = False,
+    ):
+        defaults=dict(init=init,beta=beta,c=c,maxiter=maxiter,min_alpha=min_alpha,target_iters=target_iters,nplus=nplus,scale_beta=scale_beta, try_negative=try_negative)
+        super().__init__(defaults=defaults)
+
+        self.global_state['beta_scale'] = 1.0
+        self.global_state['initial_scale'] = 1.0
+
+    def reset(self):
+        super().reset()
+        self.global_state['beta_scale'] = 1.0
+        self.global_state['initial_scale'] = 1.0
+
+    @torch.no_grad
+    def search(self, update, vars):
+        init, beta, c, maxiter, min_alpha, target_iters, nplus, scale_beta, try_negative=itemgetter(
+            'init','beta','c','maxiter','min_alpha','target_iters','nplus','scale_beta', 'try_negative')(self.settings[vars.params[0]])
+
+        objective = self.make_objective(vars=vars)
+
+        # directional derivative (0 if c = 0 because it is not needed)
+        if c == 0: d = 0
+        else: d = -sum(t.sum() for t in torch._foreach_mul(vars.get_grad(), update))
+
+        # scale beta
+        beta = beta * self.global_state['beta_scale']
+
+        # scale step size so that decrease is expected at target_iters
+        init = init * self.global_state['initial_scale']
+
+        step_size = backtracking_line_search(objective, d, init=init, beta=beta,
+                                        c=c,maxiter=maxiter,a_min=min_alpha, try_negative=try_negative)
+
+        # found an alpha that reduces loss
+        if step_size is not None:
+
+            # update initial_scale
+            # initial step size satisfied conditions, increase initial_scale by nplus
+            if step_size == init and target_iters > 0:
+                self.global_state['initial_scale'] *= nplus ** target_iters
+                self.global_state['initial_scale'] = min(self.global_state['initial_scale'], 1e32) # avoid overflow error
+
+            else:
+                # otherwise make initial_scale such that target_iters iterations will satisfy armijo
+                init_target = step_size
+                for _ in range(target_iters):
+                    init_target = step_size / beta
+
+                self.global_state['initial_scale'] = _lerp(
+                    self.global_state['initial_scale'], init_target / init, 1-scale_beta
+                )
+
+            # revert beta_scale
+            self.global_state['beta_scale'] = min(1.0, self.global_state['beta_scale'] * math.sqrt(1.5))
+
+            return step_size
+
+        # on fail reduce beta scale value
+        self.global_state['beta_scale'] /= 1.5
+        return 0