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

torchzero/modules/gradient_approximation/base_approximator.py

@@ -1,105 +0,0 @@
-from abc import ABC, abstractmethod
-from functools import partial
-from typing import Any, Literal
-
-import torch
-
-from ...core import (
-    OptimizationVars,
-    OptimizerModule,
-    _ClosureType,
-    _maybe_pass_backward,
-    _ScalarLoss,
-)
-from ...tensorlist import TensorList
-
-
-class GradientApproximatorBase(OptimizerModule, ABC):
-    """Base gradient approximator class. This is an abstract class, please don't use it as the optimizer.
-
-    When inheriting from this class the easiest way is to override `_make_ascent`, which should
-    return the ascent direction (like approximated gradient).
-
-    Args:
-        defaults (dict[str, Any]): defaults
-        requires_fx0 (bool):
-            if True, makes sure to calculate fx0 beforehand.
-            This means `_make_ascent` will always receive a pre-calculated `fx0` that won't be None.
-
-        target (str, optional):
-            determines what this module sets.
-
-            "ascent" - it creates a new ascent direction but doesn't treat is as gradient.
-
-            "grad" - it creates the gradient and sets it to `.grad` attributes (default).
-
-            "closure" - it makes a new closure that sets the estimated gradient to the `.grad` attributes.
-    """
-    def __init__(self, defaults: dict[str, Any], requires_fx0: bool, target: Literal['ascent', 'grad', 'closure']):
-        super().__init__(defaults, target)
-        self.requires_fx0 = requires_fx0
-
-    def _step_make_closure_(self, vars: OptimizationVars, params: TensorList):
-        if vars.closure is None: raise ValueError("gradient approximation requires closure")
-        closure = vars.closure
-
-        if self.requires_fx0: fx0 = vars.evaluate_fx0_(False)
-        else: fx0 = vars.fx0
-
-        def new_closure(backward=True) -> _ScalarLoss:
-            if backward:
-                g, ret_fx0, ret_fx0_approx = self._make_ascent(closure, params, fx0)
-                params.set_grad_(g)
-
-                if ret_fx0 is None: return ret_fx0_approx # type:ignore
-                return ret_fx0
-
-            return closure(False)
-
-        vars.closure = new_closure
-
-    def _step_make_target_(self, vars: OptimizationVars, params: TensorList):
-        if vars.closure is None: raise ValueError("gradient approximation requires closure")
-
-        if self.requires_fx0: fx0 = vars.evaluate_fx0_(False)
-        else: fx0 = vars.fx0
-
-        g, vars.fx0, vars.fx0_approx = self._make_ascent(vars.closure, params, fx0)
-        if self._default_step_target == 'ascent': vars.ascent = g
-        elif self._default_step_target == 'grad': vars.set_grad_(g, params)
-        else: raise ValueError(f"Unknown target {self._default_step_target}")
-
-    @torch.no_grad
-    def step(self, vars: OptimizationVars):
-        params = self.get_params()
-        if self._default_step_target == 'closure':
-            self._step_make_closure_(vars, params)
-
-        else:
-            self._step_make_target_(vars, params)
-
-        return self._update_params_or_step_with_next(vars, params)
-
-    @abstractmethod
-    @torch.no_grad
-    def _make_ascent(
-        self,
-        # vars: OptimizationVars,
-        closure: _ClosureType,
-        params: TensorList,
-        fx0: Any,
-    ) -> tuple[TensorList, _ScalarLoss | None, _ScalarLoss | None]:
-        """This should return a tuple of 3 elements:
-
-        .. code:: py
-
-            (ascent, fx0, fx0_approx)
-
-        Args:
-            closure (_ClosureType): closure
-            params (TensorList): parameters
-            fx0 (Any): fx0, can be None unless :target:`requires_fx0` is True on this module.
-
-        Returns:
-            (ascent, fx0, fx0_approx)
-        """

torchzero/modules/gradient_approximation/fdm.py

@@ -1,125 +0,0 @@
-from typing import Literal, Any
-from warnings import warn
-import torch
-
-from ...utils.python_tools import _ScalarLoss
-from ...tensorlist import TensorList
-from ...core import _ClosureType, OptimizerModule, OptimizationVars
-from ._fd_formulas import _FD_Formulas
-from .base_approximator import GradientApproximatorBase
-
-def _two_point_fd_(closure: _ClosureType, idx: int, pvec: torch.Tensor, gvec: torch.Tensor, eps: _ScalarLoss, fx0: _ScalarLoss, ):
-    """Two point finite difference (same signature for all other finite differences functions).
-
-    Args:
-        closure (Callable): A closure that reevaluates the model and returns the loss.
-        idx (int): Flat index of the current parameter.
-        pvec (Tensor): Flattened view of the current parameter tensor.
-        gvec (Tensor): Flattened view of the current parameter tensor gradient.
-        eps (float): Finite difference epsilon.
-        fx0 (ScalarType): Loss at fx0, to avoid reevaluating it each time. On some functions can be None when it isn't needed.
-
-    Returns:
-        This modifies `gvec` in place.
-        This returns loss, not necessarily at fx0 (for example central difference never evaluate at fx0).
-        So this should be assigned to fx0_approx.
-    """
-    pvec[idx] += eps
-    fx1 = closure(False)
-    gvec[idx] = (fx1 - fx0) / eps
-    pvec[idx] -= eps
-    return fx0
-
-def _two_point_bd_(closure: _ClosureType, idx: int, pvec: torch.Tensor, gvec: torch.Tensor, eps: _ScalarLoss, fx0: _ScalarLoss, ):
-    pvec[idx] += eps
-    fx1 = closure(False)
-    gvec[idx] = (fx0 - fx1) / eps
-    pvec[idx] -= eps
-    return fx0
-
-def _two_point_cd_(closure: _ClosureType, idx: int, pvec: torch.Tensor, gvec: torch.Tensor, eps: _ScalarLoss, fx0 = None, ):
-    pvec[idx] += eps
-    fxplus = closure(False)
-    pvec[idx] -= eps * 2
-    fxminus = closure(False)
-    gvec[idx] = (fxplus - fxminus) / (2 * eps)
-    pvec[idx] += eps
-    return fxplus
-
-def _three_point_fd_(closure: _ClosureType, idx: int, pvec: torch.Tensor, gvec: torch.Tensor, eps: _ScalarLoss, fx0: _ScalarLoss, ):
-    pvec[idx] += eps
-    fx1 = closure(False)
-    pvec[idx] += eps
-    fx2 = closure(False)
-    gvec[idx] = (-3*fx0 + 4*fx1 - fx2) / (2 * eps)
-    pvec[idx] -= 2 * eps
-    return fx0
-
-def _three_point_bd_(closure: _ClosureType, idx: int, pvec: torch.Tensor, gvec: torch.Tensor, eps: _ScalarLoss, fx0: _ScalarLoss, ):
-    pvec[idx] -= eps
-    fx1 = closure(False)
-    pvec[idx] -= eps
-    fx2 = closure(False)
-    gvec[idx] = (fx2 - 4*fx1 + 3*fx0) / (2 * eps)
-    pvec[idx] += 2 * eps
-    return fx0
-
-
-class FDM(GradientApproximatorBase):
-    """Gradient approximation via finite difference.
-
-    This performs :math:`num_parameters + 1` or :math:`num_parameters * 2` evaluations per step, depending on formula.
-
-    Args:
-        eps (float, optional): finite difference epsilon. Defaults to 1e-5.
-        formula (_FD_Formulas, optional): finite difference formula. Defaults to 'forward'.
-        n_points (T.Literal[2, 3], optional): number of points, 2 or 3. Defaults to 2.
-        target (str, optional):
-            determines what this module sets.
-
-            "ascent" - it creates a new ascent direction but doesn't treat is as gradient.
-
-            "grad" - it creates the gradient and sets it to `.grad` attributes (default).
-
-            "closure" - it makes a new closure that sets the estimated gradient to the `.grad` attributes.
-    """
-    def __init__(
-        self,
-        eps: float = 1e-5,
-        formula: _FD_Formulas = "forward",
-        n_points: Literal[2, 3] = 2,
-        target: Literal["ascent", "grad", "closure"] = "grad",
-    ):
-        defaults = dict(eps = eps)
-
-        if formula == 'central':
-            self._finite_difference_ = _two_point_cd_ # this is both 2 and 3 point formula
-            requires_fx0 = False
-
-        elif formula == 'forward':
-            if n_points == 2: self._finite_difference_ = _two_point_fd_
-            else: self._finite_difference_ = _three_point_fd_
-            requires_fx0 = True
-
-        elif formula == 'backward':
-            if n_points == 2: self._finite_difference_ = _two_point_bd_
-            else: self._finite_difference_ = _three_point_bd_
-            requires_fx0 = True
-
-        else: raise ValueError(f'{formula} is not valid.')
-
-        super().__init__(defaults, requires_fx0=requires_fx0, target = target)
-
-    @torch.no_grad
-    def _make_ascent(self, closure, params, fx0):
-        grads = params.zeros_like()
-        epsilons = self.get_group_key('eps')
-
-        fx0_approx = None
-        for p, g, eps in zip(params, grads, epsilons):
-            flat_param = p.view(-1)
-            flat_grad = g.view(-1)
-            for idx in range(flat_param.numel()):
-                fx0_approx = self._finite_difference_(closure, idx, flat_param, flat_grad, eps, fx0)
-
-        return grads, fx0, fx0_approx

torchzero/modules/gradient_approximation/forward_gradient.py

@@ -1,163 +0,0 @@
-from collections.abc import Iterable
-from typing import Literal
-
-import torch
-import torch.autograd.forward_ad as fwAD
-
-from ...core import OptimizerModule, _ClosureType
-from ...tensorlist import TensorList
-from ...random import Distributions
-from ...utils.torch_tools import swap_tensors_no_use_count_check
-from .base_approximator import GradientApproximatorBase
-
-def get_forward_gradient(
-    params: Iterable[torch.Tensor],
-    closure: _ClosureType,
-    n_samples: int,
-    distribution: Distributions,
-    mode: Literal["jvp", "grad", "fd"],
-    fd_eps: float = 1e-4,
-):
-    """Evaluates forward gradient of a closure w.r.t iterable of parameters with a random tangent vector.
-
-    Args:
-        params (Iterable[torch.Tensor]): iterable of parameters of the model.
-        closure (_ClosureType):
-            A closure that reevaluates the model and returns the loss.
-            Closure must accept `backward = True` boolean argument. Forward gradient will always call it as
-            `closure(False)`, unless `mode = "grad"` which requires a backward pass.
-        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 in two forward passes,
-            doesn't require the objective to be autodiffable. Equivalent to randomized FDM.
-
-        fd_eps (float, optional): epsilon for finite difference, only has effect if mode is "fd". Defaults to 1e-4.
-
-    Returns:
-        TensorList: list of estimated gradients of the same structure and shape as `params`.
-
-    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
-    """
-    if not isinstance(params, TensorList): params = TensorList(params)
-    params = params.with_requires_grad()
-
-    orig_params = None
-    grad = None
-    loss = None
-    for _ in range(n_samples):
-
-        # generate random vector
-        tangents = params.sample_like(fd_eps if mode == 'fd' else 1, distribution)
-
-        if mode == 'jvp':
-            if orig_params is None:
-                orig_params = params.clone().requires_grad_()
-
-            # evaluate jvp with it
-            with fwAD.dual_level():
-
-                # swap to duals
-                for param, clone, tangent in zip(params, orig_params, tangents):
-                    dual = fwAD.make_dual(clone, tangent)
-                    torch.utils.swap_tensors(param, dual)
-
-                loss = closure(False)
-                jvp = fwAD.unpack_dual(loss).tangent
-
-        elif mode == 'grad':
-            with torch.enable_grad(): loss = closure()
-            jvp = tangents.mul(params.ensure_grad_().grad).sum()
-
-        elif mode == 'fd':
-            loss = closure(False)
-            params += tangents
-            loss2 = closure(False)
-            params -= tangents
-            jvp = (loss2 - loss) / fd_eps**2
-
-        else:
-            raise ValueError(mode)
-
-        # update grad estimate
-        if grad is None: grad = tangents * jvp
-        else: grad += tangents * jvp
-
-    # swap back to original params
-    if orig_params is not None:
-        for param, orig in zip(params, orig_params):
-            swap_tensors_no_use_count_check(param, orig)
-
-    assert grad is not None
-    assert loss is not None
-    if n_samples > 1:
-        grad /= n_samples
-
-    return grad, loss
-
-class ForwardGradient(GradientApproximatorBase):
-    """Evaluates jacobian-vector product with a random vector using forward mode autodiff (torch.autograd.forward_ad), which is
-    the true directional derivative in the direction of that vector.
-
-    Args:
-        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,
-            because it doesn't have to store intermediate activations.
-
-            "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 in two forward passes,
-            doesn't require the objective to be autodiffable. Equivalent to randomized FDM.
-
-        fd_eps (float, optional): epsilon for finite difference, only has effect if mode is "fd". Defaults to 1e-4.
-        target (str, optional):
-            determines what this module sets.
-
-            "ascent" - it creates a new ascent direction but doesn't treat is as gradient.
-
-            "grad" - it creates the gradient and sets it to `.grad` attributes (default).
-
-            "closure" - it makes a new closure that sets the estimated gradient to the `.grad` attributes.
-
-    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,
-        n_samples: int = 1,
-        distribution: Distributions = "normal",
-        mode: Literal["jvp", "grad", "fd"] = "jvp",
-        fd_eps: float = 1e-4,
-        target: Literal['ascent', 'grad', 'closure'] = 'grad',
-    ):
-        super().__init__({}, requires_fx0=False, target = target)
-        self.distribution: Distributions = distribution
-        self.n_samples = n_samples
-        self.mode: Literal["jvp", "grad", "fd"] = mode
-        self.fd_eps = fd_eps
-
-
-    def _make_ascent(self, closure, params, fx0):
-        g, fx0 = get_forward_gradient(
-            params=params,
-            closure=closure,
-            n_samples=self.n_samples,
-            distribution=self.distribution,
-            mode=self.mode,
-            fd_eps=self.fd_eps,
-        )
-
-        return g, fx0, None

torchzero/modules/gradient_approximation/newton_fdm.py

@@ -1,198 +0,0 @@
-import typing as T
-
-import torch
-
-from ...utils.python_tools import _ScalarLoss
-from ...tensorlist import TensorList
-from ...core import _ClosureType, OptimizerModule
-from ..second_order.newton import (LINEAR_SYSTEM_SOLVERS,
-                                   FallbackLinearSystemSolvers,
-                                   LinearSystemSolvers, _fallback_gd)
-from ._fd_formulas import _FD_Formulas
-
-
-def _three_point_2cd_(
-    closure: _ClosureType,
-    idx1: int,
-    idx2: int,
-    p1: torch.Tensor,
-    p2: torch.Tensor,
-    g1: torch.Tensor,
-    hessian: torch.Tensor,
-    eps1: _ScalarLoss,
-    eps2: _ScalarLoss,
-    i1: int,
-    i2: int,
-    fx0: _ScalarLoss,
-):
-    """Second order three point finite differences (same signature for all other 2nd order finite differences functions).
-
-    Args:
-        closure (ClosureType): _description_
-        idx1 (int): _description_
-        idx2 (int): _description_
-        p1 (torch.Tensor): _description_
-        p2 (torch.Tensor): _description_
-        g1 (torch.Tensor): _description_
-        g2 (torch.Tensor): _description_
-        hessian (torch.Tensor): _description_
-        eps1 (ScalarType): _description_
-        eps2 (ScalarType): _description_
-        i1 (int): _description_
-        i23 (int): _description_
-        fx0 (ScalarType): _description_
-
-    """
-    # same param
-    if i1 == i2 and idx1 == idx2:
-        p1[idx1] += eps1
-        fxplus = closure(False)
-
-        p1[idx1] -= 2*eps1
-        fxminus = closure(False)
-
-        p1[idx1] += eps1
-
-        g1[idx1] = (fxplus - fxminus) / (2 * eps1)
-        hessian[i1, i2] = (fxplus - 2*fx0 + fxminus) / eps1**2
-
-    else:
-        p1[idx1] += eps1
-        p2[idx2] += eps2
-        fxpp = closure(False)
-        p1[idx1] -= eps1*2
-        fxnp = closure(False)
-        p2[idx2] -= eps2*2
-        fxnn = closure(False)
-        p1[idx1] += eps1*2
-        fxpn = closure(False)
-
-        p1[idx1] -= eps1
-        p2[idx2] += eps2
-
-        hessian[i1, i2] = (fxpp - fxpn - fxnp + fxnn) / (4 * eps1 * eps2)
-
-
-class NewtonFDM(OptimizerModule):
-    """Newton method with gradient and hessian approximated via finite difference.
-
-    Args:
-        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.
-            If true, ignores `solver` and `fallback`. 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.
-
-    """
-    def __init__(
-        self,
-        eps: float = 1e-2,
-        diag=False,
-        solver: LinearSystemSolvers = "cholesky_lu",
-        fallback: FallbackLinearSystemSolvers = "safe_diag",
-        validate=False,
-        tol: float = 1,
-        gd_lr = 1e-2,
-    ):
-        defaults = dict(eps = eps)
-        super().__init__(defaults)
-        self.diag = diag
-        self.solver = LINEAR_SYSTEM_SOLVERS[solver]
-        self.fallback = LINEAR_SYSTEM_SOLVERS[fallback]
-
-        self.validate = validate
-        self.gd_lr = gd_lr
-        self.tol = tol
-
-    @torch.no_grad
-    def step(self, vars):
-        """Returns a new ascent direction."""
-        if vars.closure is None: raise ValueError('NewtonFDM requires a closure.')
-        if vars.ascent is not None: raise ValueError('NewtonFDM got ascent direction')
-
-        params = self.get_params()
-        epsilons = self.get_group_key('eps')
-
-        # evaluate fx0.
-        if vars.fx0 is None: vars.fx0 = vars.closure(False)
-
-        # evaluate gradients and hessian via finite differences.
-        grads = params.zeros_like()
-        numel = params.total_numel()
-        hessian = torch.zeros((numel, numel), dtype = params[0].dtype, device = params[0].device)
-
-        cur1 = 0
-        for p1, g1, eps1 in zip(params, grads, epsilons):
-            flat_param1 = p1.view(-1)
-            flat_grad1 = g1.view(-1)
-            for idx1 in range(flat_param1.numel()):
-
-                cur2 = 0
-                for p2, eps2 in zip(params, epsilons):
-
-                    flat_param2 = p2.view(-1)
-                    for idx2 in range(flat_param2.numel()):
-                        if self.diag and (idx1 != idx2 or cur1 != cur2):
-                            cur2 += 1
-                            continue
-                        _three_point_2cd_(
-                            closure = vars.closure,
-                            idx1 = idx1,
-                            idx2 = idx2,
-                            p1 = flat_param1,
-                            p2 = flat_param2,
-                            g1 = flat_grad1,
-                            hessian = hessian,
-                            eps1 = eps1,
-                            eps2 = eps2,
-                            fx0 = vars.fx0,
-                            i1 = cur1,
-                            i2 = cur2,
-                        )
-                        cur2 += 1
-                cur1 += 1
-
-        gvec = grads.to_vec()
-        if self.diag:
-            hdiag = hessian.diag()
-            hdiag[hdiag == 0] = 1
-            newton_step = gvec / hdiag
-        else:
-            newton_step, success = self.solver(hessian, gvec)
-            if not success:
-                newton_step, success = self.fallback(hessian, gvec)
-                if not success:
-                    newton_step, success = _fallback_gd(hessian, gvec)
-
-        # update params or pass the gradients to the child.
-        vars.ascent = grads.from_vec(newton_step)
-
-
-        # validate if newton step decreased loss
-        if self.validate:
-
-            params.sub_(vars.ascent)
-            fx1 = vars.closure(False)
-            params.add_(vars.ascent)
-
-            # if loss increases, set ascent direction to gvec times lr
-            if fx1 - vars.fx0 > vars.fx0 * self.tol:
-                vars.ascent = grads.from_vec(gvec) * self.gd_lr
-
-        return self._update_params_or_step_with_next(vars, params)