torchzero 0.3.11__py3-none-any.whl → 0.3.13__py3-none-any.whl

torchzero/modules/trust_region/trust_region.py

@@ -0,0 +1,350 @@
+import math
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping, Sequence
+from functools import partial
+from typing import Any, Literal, Protocol, cast, final, overload
+
+import torch
+
+from ...core import Chainable, Module, Var, apply_transform
+from ...utils import TensorList, safe_dict_update_, tofloat, vec_to_tensors, generic_finfo, generic_vector_norm
+from ...utils.linalg.linear_operator import LinearOperator
+
+
+def _flatten_tensors(tensors: list[torch.Tensor]):
+    return torch.cat([t.ravel() for t in tensors])
+
+
+
+class _RadiusStrategy(Protocol):
+    def __call__(
+        self,
+        params: Sequence[torch.Tensor],
+        closure: Callable,
+        f: float,
+        g: torch.Tensor,
+        H: LinearOperator,
+        d: torch.Tensor,
+        trust_radius: float,
+        eta: float, # 0.0
+        nplus: float, # 3.5
+        nminus: float, # 0.25
+        rho_good: float, # 0.99
+        rho_bad: float, # 1e-4
+        boundary_tol: float | None,
+        init: float,
+        state: Mapping[str, Any],
+        settings: Mapping[str, Any],
+        radius_fn: Callable | None = torch.linalg.vector_norm,
+    ) -> tuple[float, bool]:
+        """returns (new trust_region value, success).
+
+        Args:
+            params (Sequence[torch.Tensor]): params tensor list
+            closure (Callable): closure
+            d (torch.Tensor):
+                current update vector with current trust_region, which is SUBTRACTED from parameters.
+                May be exact solution to (B+yI)x=g, approximate, or a solution to a different subproblem
+                (e.g. cubic regularization).
+            f (float | torch.Tensor): loss at x0
+            g (torch.Tensor): gradient vector
+            H (LinearOperator | None): hessian approximation
+            trust_radius (float): current trust region value
+            eta (float, optional):
+                if ratio of actual to predicted rediction is larger than this, step is accepted.
+                When :code:`hess_module` is GaussNewton, this can be set to 0. Defaults to 0.15.
+            nplus (float, optional): increase factor on successful steps. Defaults to 1.5.
+            nminus (float, optional): decrease factor on unsuccessful steps. Defaults to 0.75.
+            rho_good (float, optional):
+                if ratio of actual to predicted rediction is larger than this, trust region size is multiplied by `nplus`.
+            rho_bad (float, optional):
+                if ratio of actual to predicted rediction is less than this, trust region size is multiplied by `nminus`.
+            boundary_tol (float | None, optional):
+                The trust region only increases when suggested step's norm is at least `(1-boundary_tol)*trust_region`.
+                This prevents increasing trust region when solution is not on the boundary. Defaults to 1e-2.
+            init (float, optional): Initial trust region value. Defaults to 1.
+            state (dict, optional): global state of the module for storing persistent info.
+            settings (dict, optional): all settings in case this strategy has other settings.
+            radius_fn (Callable | None, optional):
+                function that accepts ``(d: torch.Tensor)`` and returns the actual region of ``d``
+                (e.g. L2) norm for L2 trust region.
+        """
+        ... # pylint:disable=unnecessary-ellipsis
+
+def _get_rho(params: Sequence[torch.Tensor], closure:Callable,
+             f: float, g: torch.Tensor, H: LinearOperator, d:torch.Tensor, ):
+    """rho is reduction/pred_reduction"""
+
+    # evaluate actual loss reduction
+    update_unflattned = vec_to_tensors(d, params)
+    params = TensorList(params)
+    x0 = params.clone() # same as in line searches, large directions are undone very imprecisely
+
+    params -= update_unflattned
+    f_star = closure(False)
+    params.set_(x0)
+
+    reduction = f - f_star
+
+    # expected reduction is g.T @ p + 0.5 * p.T @ B @ p
+    Hu = H.matvec(d)
+    pred_reduction = g.dot(d) - 0.5 * d.dot(Hu)
+
+    rho = reduction / (pred_reduction.clip(min=torch.finfo(g.dtype).tiny * 2))
+    return rho, f_star, reduction, pred_reduction
+
+def _get_rho_tensorlist(params: Sequence[torch.Tensor], closure:Callable,
+             f: float, g: TensorList, Hvp: Callable[[TensorList], TensorList], d:TensorList):
+    """rho is reduction/pred_reduction"""
+    params = TensorList(params)
+    x0 = params.clone() # same as in line searches, large directions are undone very imprecisely
+
+    # evaluate before modifying params to not break autograd
+    Hu = Hvp(d)
+
+    # actual f
+    params -= d
+    f_star = closure(False)
+    params.copy_(x0)
+
+    reduction = f - f_star
+
+    # expected f is g.T @ p + 0.5 * p.T @ B @ p
+    pred_reduction = g.dot(d) - 0.5 * d.dot(Hu)
+
+    rho = reduction / (pred_reduction.clip(min=torch.finfo(g[0].dtype).tiny * 2))
+    return rho, f_star, reduction, pred_reduction
+
+@torch.no_grad
+def default_radius(
+    params: Sequence[torch.Tensor],
+    closure: Callable,
+    f: float,
+    g: torch.Tensor | TensorList,
+    H: LinearOperator | Callable,
+    d: torch.Tensor | TensorList,
+    trust_radius: float,
+    eta: float, # 0.0
+    nplus: float, # 3.5
+    nminus: float, # 0.25
+    rho_good: float, # 0.99
+    rho_bad: float, # 1e-4
+    boundary_tol: float | None,
+    init: float,
+    state: Mapping[str, Any],
+    settings: Mapping[str, Any],
+    radius_fn: Callable | None = generic_vector_norm,
+    check_overflow: bool = True,
+    # dynamic_nminus: bool=False,
+) -> tuple[float, bool]:
+
+    # when rho_bad < rho < eta, no update is made but trust region is not updated.
+    if eta > rho_bad:
+        warnings.warn(f"trust region eta={eta} is larger than rho_bad={rho_bad}, "
+                      "this can lead to trust region getting stuck.")
+
+    if isinstance(g, torch.Tensor):
+        rho, f_star, _, _ = _get_rho(params=params, closure=closure, f=f, g=g, H=H, d=d) # pyright:ignore[reportArgumentType]
+    else:
+        rho, f_star, _, _ = _get_rho_tensorlist(params=params, closure=closure, f=f, g=g, Hvp=H, d=d) # pyright:ignore[reportArgumentType]
+
+    is_finite = math.isfinite(f_star)
+
+    # find boundary of current step
+    if radius_fn is None: d_radius = trust_radius
+    else: d_radius = radius_fn(d)
+
+    # failed step
+    if rho < rho_bad or not is_finite:
+        # if dynamic_nminus and rho > 0: nminus = nminus * max(rho, 1e-4)
+        trust_radius = d_radius*nminus
+
+    # very good step
+    elif rho > rho_good and is_finite:
+        if (boundary_tol is None) or (trust_radius-d_radius)/trust_radius < boundary_tol:
+            trust_radius = max(trust_radius, d_radius*nplus)
+
+    # prevent very small or large values
+    if check_overflow:
+        finfo = generic_finfo(g)
+        if trust_radius < finfo.tiny*2 or trust_radius > finfo.max/2:
+            trust_radius = init
+
+    # return new trust region and success boolean
+    return tofloat(trust_radius), rho > eta and is_finite
+
+
+def fixed_radius(
+    params: Sequence[torch.Tensor],
+    closure: Callable,
+    f: float,
+    g: torch.Tensor,
+    H: LinearOperator,
+    d: torch.Tensor,
+    trust_radius: float,
+    eta: float, # 0.0
+    nplus: float, # 3.5
+    nminus: float, # 0.25
+    rho_good: float, # 0.99
+    rho_bad: float, # 1e-4
+    boundary_tol: float | None,
+    init: float,
+    state: Mapping[str, Any],
+    settings: Mapping[str, Any],
+    radius_fn: Callable | None = torch.linalg.vector_norm,
+) -> tuple[float, bool]:
+    return init, True
+
+_RADIUS_KEYS = Literal['default', 'fixed']
+_RADIUS_STRATEGIES: dict[_RADIUS_KEYS, _RadiusStrategy] = {
+    "default": default_radius,
+    "fixed": fixed_radius,
+    # "dynamic": partial(default_radius, dynamic_nminus=True)
+}
+
+class TrustRegionBase(Module, ABC):
+    def __init__(
+        self,
+        defaults: dict | None,
+        hess_module: Chainable,
+        # suggested default values:
+        # Gould, Nicholas IM, et al. "Sensitivity of trust-region algorithms to their parameters." 4OR 3.3 (2005): 227-241.
+        # which I found from https://github.com/patrick-kidger/optimistix/blob/c1dad7e75fc35bd5a4977ac3a872991e51e83d2c/optimistix/_solver/trust_region.py#L113-200
+        eta: float, # 0.0
+        nplus: float, # 3.5
+        nminus: float, # 0.25
+        rho_good: float, # 0.99
+        rho_bad: float, # 1e-4
+        boundary_tol: float | None, # None or 1e-1
+        init: float, # 1
+        max_attempts: int, # 10
+        radius_strategy: _RadiusStrategy | _RADIUS_KEYS, # "default"
+        radius_fn: Callable | None, # torch.linalg.vector_norm
+        update_freq: int = 1,
+        inner: Chainable | None = None,
+    ):
+        if isinstance(radius_strategy, str): radius_strategy = _RADIUS_STRATEGIES[radius_strategy]
+        if defaults is None: defaults = {}
+
+        safe_dict_update_(
+            defaults,
+            dict(eta=eta, nplus=nplus, nminus=nminus, rho_good=rho_good, rho_bad=rho_bad, init=init,
+                 update_freq=update_freq, max_attempts=max_attempts, radius_strategy=radius_strategy,
+                 boundary_tol=boundary_tol)
+        )
+
+        super().__init__(defaults)
+
+        self._radius_fn = radius_fn
+        self.set_child('hess_module', hess_module)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @abstractmethod
+    def trust_solve(
+        self,
+        f: float,
+        g: torch.Tensor,
+        H: LinearOperator,
+        radius: float,
+        params: list[torch.Tensor],
+        closure: Callable,
+        settings: Mapping[str, Any],
+    ) -> torch.Tensor:
+        """Solve Hx=g with a trust region penalty/bound defined by `radius`"""
+        ... # pylint:disable=unnecessary-ellipsis
+
+    def trust_region_update(self, var: Var, H: LinearOperator | None) -> None:
+        """updates the state of this module after H or B have been updated, if necessary"""
+
+    def trust_region_apply(self, var: Var, tensors:list[torch.Tensor], H: LinearOperator | None) -> Var:
+        """Solves the trust region subproblem and outputs ``Var`` with the solution direction."""
+        assert H is not None
+
+        params = TensorList(var.params)
+        settings = self.settings[params[0]]
+        g = _flatten_tensors(tensors)
+
+        max_attempts = settings['max_attempts']
+
+        # loss at x_0
+        loss = var.loss
+        closure = var.closure
+        if closure is None: raise RuntimeError("Trust region requires closure")
+        if loss is None: loss = var.get_loss(False)
+        loss = tofloat(loss)
+
+        # trust region step and update
+        success = False
+        d = None
+        while not success:
+            max_attempts -= 1
+            if max_attempts < 0: break
+
+            trust_radius = self.global_state.get('trust_radius', settings['init'])
+
+            # solve Hx=g
+            d = self.trust_solve(f=loss, g=g, H=H, radius=trust_radius, params=params, closure=closure, settings=settings)
+
+            # update trust radius
+            radius_strategy: _RadiusStrategy = settings['radius_strategy']
+            self.global_state["trust_radius"], success = radius_strategy(
+                params=params,
+                closure=closure,
+                d=d,
+                f=loss,
+                g=g,
+                H=H,
+                trust_radius=trust_radius,
+
+                eta=settings["eta"],
+                nplus=settings["nplus"],
+                nminus=settings["nminus"],
+                rho_good=settings["rho_good"],
+                rho_bad=settings["rho_bad"],
+                boundary_tol=settings["boundary_tol"],
+                init=settings["init"],
+
+                state=self.global_state,
+                settings=settings,
+                radius_fn=self._radius_fn,
+            )
+
+        assert d is not None
+        if success: var.update = vec_to_tensors(d, params)
+        else: var.update = params.zeros_like()
+
+        return var
+
+
+    @final
+    @torch.no_grad
+    def update(self, var):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        if step % self.defaults["update_freq"] == 0:
+
+            hessian_module = self.children['hess_module']
+            hessian_module.update(var)
+            H = hessian_module.get_H(var)
+            self.global_state["H"] = H
+
+            self.trust_region_update(var, H=H)
+
+
+    @final
+    @torch.no_grad
+    def apply(self, var):
+        H = self.global_state.get('H', None)
+
+        # -------------------------------- inner step -------------------------------- #
+        update = var.get_update()
+        if 'inner' in self.children:
+            update = apply_transform(self.children['inner'], update, params=var.params, grads=var.grad, var=var)
+
+        # ----------------------------------- apply ---------------------------------- #
+        return self.trust_region_apply(var=var, tensors=update, H=H)
+

torchzero/modules/variance_reduction/__init__.py

@@ -0,0 +1 @@
+from .svrg import SVRG

torchzero/modules/variance_reduction/svrg.py

@@ -0,0 +1,208 @@
+import warnings
+from functools import partial
+
+import torch
+
+from ...core.module import Module
+from ...utils import tofloat
+
+
+def _reset_except_self(optimizer, var, self: Module):
+    for m in optimizer.unrolled_modules:
+        if m is not self:
+            m.reset()
+
+class SVRG(Module):
+    """Stochastic variance reduced gradient method (SVRG).
+
+    To use, put SVRG as the first module, it can be used with any other modules.
+    To reduce variance of a gradient estimator, put the gradient estimator before SVRG.
+
+    First it uses first ``accum_steps`` batches to compute full gradient at initial
+    parameters using gradient accumulation, the model will not be updated during this.
+
+    Then it performs ``svrg_steps`` SVRG steps, each requires two forward and backward passes.
+
+    After ``svrg_steps``, it goes back to full gradient computation step step.
+
+    As an alternative to gradient accumulation you can pass "full_closure" argument to the ``step`` method,
+    which should compute full gradients, set them to ``.grad`` attributes of the parameters,
+    and return full loss.
+
+    Args:
+        svrg_steps (int): number of steps before calculating full gradient. This can be set to length of the dataloader.
+        accum_steps (int | None, optional):
+            number of steps to accumulate the gradient for. Not used if "full_closure" is passed to the ``step`` method. If None, uses value of ``svrg_steps``. Defaults to None.
+        reset_before_accum (bool, optional):
+            whether to reset all other modules when re-calculating full gradient. Defaults to True.
+        svrg_loss (bool, optional):
+            whether to replace loss with SVRG loss (calculated by same formula as SVRG gradient). Defaults to True.
+        alpha (float, optional):
+            multiplier to ``g_full(x_0) - g_batch(x_0)`` term, can be annealed linearly from 1 to 0 as suggested in https://arxiv.org/pdf/2311.05589#page=6
+
+    ## Examples:
+    SVRG-LBFGS
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.SVRG(len(dataloader)),
+        tz.m.LBFGS(),
+        tz.m.Backtracking(),
+    )
+    ```
+
+    For extra variance reduction one can use Online versions of algorithms, although it won't always help.
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.SVRG(len(dataloader)),
+        tz.m.Online(tz.m.LBFGS()),
+        tz.m.Backtracking(),
+    )
+
+    Variance reduction can also be applied to gradient estimators.
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.SPSA(),
+        tz.m.SVRG(100),
+        tz.m.LR(1e-2),
+    )
+    ```
+    ## Notes
+
+    The SVRG gradient is computed as ``g_b(x) - alpha * g_b(x_0) - g_f(x0.)``, where:
+    - ``x`` is current parameters
+    - ``x_0`` is initial parameters, where full gradient was computed
+    - ``g_b`` refers to mini-batch gradient at ``x`` or ``x_0``
+    - ``g_f`` refers to full gradient at ``x_0``.
+
+    The SVRG loss is computed using the same formula.
+    """
+    def __init__(self, svrg_steps: int, accum_steps: int | None = None, reset_before_accum:bool=True, svrg_loss:bool=True, alpha:float=1):
+        defaults = dict(svrg_steps = svrg_steps, accum_steps=accum_steps, reset_before_accum=reset_before_accum, svrg_loss=svrg_loss, alpha=alpha)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        params = var.params
+        closure = var.closure
+        assert closure is not None
+
+        if "full_grad" not in self.global_state:
+
+            # -------------------------- calculate full gradient ------------------------- #
+            if "full_closure" in var.storage:
+                full_closure = var.storage['full_closure']
+                with torch.enable_grad():
+                    full_loss = full_closure()
+                    if all(p.grad is None for p in params):
+                        warnings.warn("all gradients are None after evaluating full_closure.")
+
+                    full_grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+                    self.global_state["full_loss"] = full_loss
+                    self.global_state["full_grad"] = full_grad
+                    self.global_state['x_0'] = [p.clone() for p in params]
+
+                # current batch will be used for svrg update
+
+            else:
+                # accumulate gradients over n steps
+                accum_steps = self.defaults['accum_steps']
+                if accum_steps is None: accum_steps = self.defaults['svrg_steps']
+
+                current_accum_step = self.global_state.get('current_accum_step', 0) + 1
+                self.global_state['current_accum_step'] = current_accum_step
+
+                # accumulate grads
+                accumulator = self.get_state(params, 'accumulator')
+                grad = var.get_grad()
+                torch._foreach_add_(accumulator, grad)
+
+                # accumulate loss
+                loss_accumulator = self.global_state.get('loss_accumulator', 0)
+                loss_accumulator += tofloat(var.loss)
+                self.global_state['loss_accumulator'] = loss_accumulator
+
+                # on nth step, use the accumulated gradient
+                if current_accum_step >= accum_steps:
+                    torch._foreach_div_(accumulator, accum_steps)
+                    self.global_state["full_grad"] = accumulator
+                    self.global_state["full_loss"] = loss_accumulator / accum_steps
+
+                    self.global_state['x_0'] = [p.clone() for p in params]
+                    self.clear_state_keys('accumulator')
+                    del self.global_state['current_accum_step']
+
+                # otherwise skip update until enough grads are accumulated
+                else:
+                    var.update = None
+                    var.stop = True
+                    var.skip_update = True
+                    return var
+
+
+        svrg_steps = self.defaults['svrg_steps']
+        current_svrg_step = self.global_state.get('current_svrg_step', 0) + 1
+        self.global_state['current_svrg_step'] = current_svrg_step
+
+        # --------------------------- SVRG gradient closure -------------------------- #
+        x0 = self.global_state['x_0']
+        gf_x0 = self.global_state["full_grad"]
+        ff_x0 = self.global_state['full_loss']
+        use_svrg_loss = self.defaults['svrg_loss']
+        alpha = self.get_settings(params, 'alpha')
+        alpha_0 = alpha[0]
+        if all(a == 1 for a in alpha): alpha = None
+
+        def svrg_closure(backward=True):
+            # g_b(x) - α * (g_f(x_0) - g_b(x_0)) and same for loss
+            with torch.no_grad():
+                x = [p.clone() for p in params]
+
+                if backward:
+                    # f and g at x
+                    with torch.enable_grad(): fb_x = closure()
+                    gb_x = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+
+                    # f and g at x_0
+                    torch._foreach_copy_(params, x0)
+                    with torch.enable_grad(): fb_x0 = closure()
+                    gb_x0 = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+                    torch._foreach_copy_(params, x)
+
+                    # g_svrg = gb_x - alpha * (gf_x0 - gb_x0)
+                    correction = torch._foreach_sub(gb_x0, gf_x0)
+                    if alpha is not None: torch._foreach_mul_(correction, alpha)
+                    g_svrg = torch._foreach_sub(gb_x, correction)
+
+                    f_svrg = fb_x - alpha_0 * (fb_x0 - ff_x0)
+                    for p, g in zip(params, g_svrg):
+                        p.grad = g
+
+                    if use_svrg_loss: return f_svrg
+                    return fb_x
+
+            # no backward
+            if use_svrg_loss:
+                fb_x = closure(False)
+                torch._foreach_copy_(params, x0)
+                fb_x0 = closure(False)
+                torch._foreach_copy_(params, x)
+                f_svrg = fb_x - alpha_0 * (fb_x0 - ff_x0)
+                return f_svrg
+
+            return closure(False)
+
+        var.closure = svrg_closure
+
+        # --- after svrg_steps steps reset so that new full gradient is calculated on next step --- #
+        if current_svrg_step >= svrg_steps:
+            del self.global_state['current_svrg_step']
+            del self.global_state['full_grad']
+            del self.global_state['full_loss']
+            del self.global_state['x_0']
+            if self.defaults['reset_before_accum']:
+                var.post_step_hooks.append(partial(_reset_except_self, self=self))
+
+        return var