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

torchzero/modules/line_search/line_search.py

@@ -10,6 +10,7 @@ import torch
 
 from ...core import Module, Target, Var
 from ...utils import tofloat, set_storage_
+from ..functional import clip_by_finfo
 
 
 class MaxLineSearchItersReached(Exception): pass
@@ -103,23 +104,18 @@ class LineSearchBase(Module, ABC):
     ):
         if not math.isfinite(step_size): return
 
-         # fixes overflow when backtracking keeps increasing alpha after converging
-        step_size = max(min(tofloat(step_size), 1e36), -1e36)
+         # avoid overflow error
+        step_size = clip_by_finfo(tofloat(step_size), torch.finfo(update[0].dtype))
 
         # skip is parameters are already at suggested step size
         if self._current_step_size == step_size: return
 
-        # this was basically causing floating point imprecision to build up
-        #if False:
-        # if abs(alpha) < abs(step_size) and step_size != 0:
-        #     torch._foreach_add_(params, update, alpha=alpha)
-
-        # else:
         assert self._initial_params is not None
         if step_size == 0:
             new_params = [p.clone() for p in self._initial_params]
         else:
             new_params = torch._foreach_sub(self._initial_params, update, alpha=step_size)
+
         for c, n in zip(params, new_params):
             set_storage_(c, n)
 
@@ -131,10 +127,7 @@ class LineSearchBase(Module, ABC):
         params: list[torch.Tensor],
         update: list[torch.Tensor],
     ):
-        # if not np.isfinite(step_size): step_size = [0 for _ in step_size]
-        # alpha = [self._current_step_size - s for s in step_size]
-        # if any(a!=0 for a in alpha):
-        #     torch._foreach_add_(params, torch._foreach_mul(update, alpha))
+
         assert self._initial_params is not None
         if not np.isfinite(step_size).all(): step_size = [0 for _ in step_size]
 
@@ -248,16 +241,14 @@ class LineSearchBase(Module, ABC):
         except MaxLineSearchItersReached:
             step_size = self._best_step_size
 
+        step_size = clip_by_finfo(step_size, torch.finfo(update[0].dtype))
+
         # set loss_approx
         if var.loss_approx is None: var.loss_approx = self._lowest_loss
 
-        # this is last module - set step size to found step_size times lr
-        if var.is_last:
-            if var.last_module_lrs is None:
-                self.set_step_size_(step_size, params=params, update=update)
-
-            else:
-                self._set_per_parameter_step_size_([step_size*lr for lr in var.last_module_lrs], params=params, update=update)
+        # if this is last module, directly update parameters to avoid redundant operations
+        if var.modular is not None and self is var.modular.modules[-1]:
+            self.set_step_size_(step_size, params=params, update=update)
 
             var.stop = True; var.skip_update = True
             return var
@@ -277,7 +268,7 @@ class GridLineSearch(LineSearchBase):
 
     @torch.no_grad
     def search(self, update, var):
-        start,end,num=itemgetter('start','end','num')(self.defaults)
+        start, end, num = itemgetter('start', 'end', 'num')(self.defaults)
 
         for lr in torch.linspace(start,end,num):
             self.evaluate_f(lr.item(), var=var, backward=False)

torchzero/modules/line_search/scipy.py

@@ -1,3 +1,4 @@
+import math
 from collections.abc import Mapping
 from operator import itemgetter
 
@@ -17,6 +18,7 @@ class ScipyMinimizeScalar(LineSearchBase):
         bounds (Sequence | None, optional):
             For method ‘bounded’, bounds is mandatory and must have two finite items corresponding to the optimization bounds. Defaults to None.
         tol (float | None, optional): Tolerance for termination. Defaults to None.
+        prev_init (bool, optional): uses previous step size as initial guess for the line search.
         options (dict | None, optional): A dictionary of solver options. Defaults to None.
 
     For more details on methods and arguments refer to https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize_scalar.html
@@ -29,9 +31,10 @@ class ScipyMinimizeScalar(LineSearchBase):
         bracket=None,
         bounds=None,
         tol: float | None = None,
+        prev_init: bool = False,
         options=None,
     ):
-        defaults = dict(method=method,bracket=bracket,bounds=bounds,tol=tol,options=options,maxiter=maxiter)
+        defaults = dict(method=method,bracket=bracket,bounds=bounds,tol=tol,options=options,maxiter=maxiter, prev_init=prev_init)
         super().__init__(defaults)
 
         import scipy.optimize
@@ -48,5 +51,14 @@ class ScipyMinimizeScalar(LineSearchBase):
             options = dict(options) if isinstance(options, Mapping) else {}
             options['maxiter'] = maxiter
 
-        res = self.scopt.minimize_scalar(objective, method=method, bracket=bracket, bounds=bounds, tol=tol, options=options)
-        return res.x
+        if self.defaults["prev_init"] and "x_prev" in self.global_state:
+            if bracket is None: bracket = (0, 1)
+            bracket = (*bracket[:-1], self.global_state["x_prev"])
+
+        x = self.scopt.minimize_scalar(objective, method=method, bracket=bracket, bounds=bounds, tol=tol, options=options).x # pyright:ignore[reportAttributeAccessIssue]
+
+        max = torch.finfo(var.params[0].dtype).max / 2
+        if (not math.isfinite(x)) or abs(x) >= max: x = 0
+
+        self.global_state['x_prev'] = x
+        return x

torchzero/modules/line_search/strong_wolfe.py

@@ -7,7 +7,7 @@ import numpy as np
 import torch
 from torch.optim.lbfgs import _cubic_interpolate
 
-from ...utils import as_tensorlist, totensor
+from ...utils import as_tensorlist, totensor, tofloat
 from ._polyinterp import polyinterp, polyinterp2
 from .line_search import LineSearchBase, TerminationCondition, termination_condition
 from ..step_size.adaptive import _bb_geom
@@ -92,7 +92,7 @@ class _StrongWolfe:
             return _apply_bounds(a_lo + 0.5 * (a_hi - a_lo), bounds)
 
         if self.interpolation in ('polynomial', 'polynomial2'):
-            finite_history = [(a, f, g) for a, (f,g) in self.history.items() if math.isfinite(a) and math.isfinite(f) and math.isfinite(g)]
+            finite_history = [(tofloat(a), tofloat(f), tofloat(g)) for a, (f,g) in self.history.items() if math.isfinite(a) and math.isfinite(f) and math.isfinite(g)]
             if bounds is None: bounds = (None, None)
             polyinterp_fn = polyinterp if self.interpolation == 'polynomial' else polyinterp2
             try:
@@ -330,7 +330,6 @@ class StrongWolfe(LineSearchBase):
         if adaptive:
             a_init *= self.global_state.get('initial_scale', 1)
 
-
         strong_wolfe = _StrongWolfe(
             f=objective,
             f_0=f_0,
@@ -360,7 +359,6 @@ class StrongWolfe(LineSearchBase):
                     if inverted: a = -a
 
         if a is not None and a != 0 and math.isfinite(a):
-            #self.global_state['initial_scale'] = min(1.0, self.global_state.get('initial_scale', 1) * math.sqrt(2))
             self.global_state['initial_scale'] = 1
             self.global_state['a_prev'] = a
             self.global_state['f_prev'] = f_0
@@ -372,6 +370,6 @@ class StrongWolfe(LineSearchBase):
             self.global_state['initial_scale'] = self.global_state.get('initial_scale', 1) * 0.5
             finfo = torch.finfo(dir[0].dtype)
             if self.global_state['initial_scale'] < finfo.tiny * 2:
-                self.global_state['initial_scale'] = finfo.max / 2
+                self.global_state['initial_scale'] = init_value * 2
 
         return 0

torchzero/modules/misc/misc.py

@@ -306,8 +306,8 @@ class RandomHvp(Module):
             for i in range(n_samples):
                 u = params.sample_like(distribution=distribution, variance=1)
 
-                Hvp, rgrad = self.Hvp(u, at_x0=True, var=var, rgrad=rgrad, hvp_method=hvp_method,
-                                    h=h, normalize=True, retain_grad=i < n_samples-1)
+                Hvp, rgrad = var.hessian_vector_product(u, at_x0=True, rgrad=rgrad, hvp_method=hvp_method,
+                                    h=h, normalize=True, retain_graph=i < n_samples-1)
 
                 if D is None: D = Hvp
                 else: torch._foreach_add_(D, Hvp)

torchzero/modules/misc/multistep.py

@@ -15,7 +15,7 @@ def _sequential_step(self: Module, var: Var, sequential: bool):
     if var.closure is None and len(modules) > 1: raise ValueError('Multistep and Sequential require closure')
 
     # store original params unless this is last module and can update params directly
-    params_before_steps = None if (var.is_last and var.last_module_lrs is None) else [p.clone() for p in params]
+    params_before_steps = [p.clone() for p in params]
 
     # first step - pass var as usual
     var = modules[0].step(var)
@@ -27,8 +27,8 @@ def _sequential_step(self: Module, var: Var, sequential: bool):
 
             # update params
             if (not new_var.skip_update):
-                if new_var.last_module_lrs is not None:
-                    torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
+                # if new_var.last_module_lrs is not None:
+                #     torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
 
                 torch._foreach_sub_(params, new_var.get_update())
 
@@ -41,16 +41,16 @@ def _sequential_step(self: Module, var: Var, sequential: bool):
 
         # final parameter update
         if (not new_var.skip_update):
-            if new_var.last_module_lrs is not None:
-                torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
+            # if new_var.last_module_lrs is not None:
+            #     torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
 
             torch._foreach_sub_(params, new_var.get_update())
 
     # if last module, update is applied so return new var
-    if params_before_steps is None:
-        new_var.stop = True
-        new_var.skip_update = True
-        return new_var
+    # if params_before_steps is None:
+    #     new_var.stop = True
+    #     new_var.skip_update = True
+    #     return new_var
 
     # otherwise use parameter difference as update
     var.update = list(torch._foreach_sub(params_before_steps, params))
@@ -106,10 +106,10 @@ class NegateOnLossIncrease(Module):
         f_1 = closure(False)
 
         if f_1 <= f_0:
-            if var.is_last and var.last_module_lrs is None:
-                var.stop = True
-                var.skip_update = True
-                return var
+            # if var.is_last and var.last_module_lrs is None:
+            #     var.stop = True
+            #     var.skip_update = True
+            #     return var
 
             torch._foreach_add_(var.params, update)
             return var

torchzero/modules/quasi_newton/__init__.py

@@ -29,3 +29,5 @@ from .quasi_newton import (
     ShorR,
     ThomasOptimalMethod,
 )
+
+from .sg2 import SG2, SPSA2

torchzero/modules/quasi_newton/quasi_newton.py

@@ -1182,16 +1182,19 @@ class ShorR(HessianUpdateStrategy):
     """Shor’s r-algorithm.
 
     Note:
-        A line search such as ``tz.m.StrongWolfe(a_init="quadratic", fallback=True)`` is required.
-        Similarly to conjugate gradient, ShorR doesn't have an automatic step size scaling,
-        so setting ``a_init`` in the line search is recommended.
+        - A line search such as ``[tz.m.StrongWolfe(a_init="quadratic", fallback=True), tz.m.Mul(1.2)]`` is required. Similarly to conjugate gradient, ShorR doesn't have an automatic step size scaling, so setting ``a_init`` in the line search is recommended.
+
+        - The line search should try to overstep by a little, therefore it can help to multiply direction given by a line search by some value slightly larger than 1 such as 1.2.
 
     References:
-        S HOR , N. Z. (1985) Minimization Methods for Non-differentiable Functions. New York: Springer.
+        Those are the original references, but neither seem to be available online:
+            - Shor, N. Z., Utilization of the Operation of Space Dilatation in the Minimization of Convex Functions, Kibernetika, No. 1, pp. 6-12, 1970.
+
+            - Skokov, V. A., Note on Minimization Methods Employing Space Stretching, Kibernetika, No. 4, pp. 115-117, 1974.
 
-        Burke, James V., Adrian S. Lewis, and Michael L. Overton. "The Speed of Shor's R-algorithm." IMA Journal of numerical analysis 28.4 (2008): 711-720. - good overview.
+        An overview is available in [Burke, James V., Adrian S. Lewis, and Michael L. Overton. "The Speed of Shor's R-algorithm." IMA Journal of numerical analysis 28.4 (2008): 711-720](https://sites.math.washington.edu/~burke/papers/reprints/60-speed-Shor-R.pdf).
 
-        Ansari, Zafar A. Limited Memory Space Dilation and Reduction Algorithms. Diss. Virginia Tech, 1998. - this is where a more efficient formula is described.
+        Reference by Skokov, V. A. describes a more efficient formula which can be found here [Ansari, Zafar A. Limited Memory Space Dilation and Reduction Algorithms. Diss. Virginia Tech, 1998.](https://camo.ici.ro/books/thesis/th.pdf)
     """
 
     def __init__(
@@ -1229,3 +1232,9 @@ class ShorR(HessianUpdateStrategy):
 
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, setting):
         return shor_r_(H=H, y=y, alpha=setting['alpha'])
+
+
+# Todd, Michael J. "The symmetric rank-one quasi-Newton method is a space-dilation subgradient algorithm." Operations research letters 5.5 (1986): 217-219.
+# TODO
+
+# Sorensen, D. C. "The q-superlinear convergence of a collinear scaling algorithm for unconstrained optimization." SIAM Journal on Numerical Analysis 17.1 (1980): 84-114.

torchzero/modules/quasi_newton/sg2.py

@@ -0,0 +1,292 @@
+import torch
+
+from ...core import Module, Chainable, apply_transform
+from ...utils import TensorList, vec_to_tensors
+from ..second_order.newton import _newton_step, _get_H
+
+def sg2_(
+    delta_g: torch.Tensor,
+    cd: torch.Tensor,
+) -> torch.Tensor:
+    """cd is c * perturbation, and must be multiplied by two if hessian estimate is two-sided
+    (or divide delta_g by two)."""
+
+    M = torch.outer(1.0 / cd, delta_g)
+    H_hat = 0.5 * (M + M.T)
+
+    return H_hat
+
+
+
+class SG2(Module):
+    """second-order stochastic gradient
+
+    SG2 with line search
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.SG2(),
+        tz.m.Backtracking()
+    )
+    ```
+
+    SG2 with trust region
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LevenbergMarquardt(tz.m.SG2()),
+    )
+    ```
+
+    """
+
+    def __init__(
+        self,
+        n_samples: int = 1,
+        h: float = 1e-2,
+        beta: float | None = None,
+        damping: float = 0,
+        eigval_fn=None,
+        one_sided: bool = False, # one-sided hessian
+        use_lstsq: bool = True,
+        seed=None,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(n_samples=n_samples, h=h, beta=beta, damping=damping, eigval_fn=eigval_fn, one_sided=one_sided, seed=seed, use_lstsq=use_lstsq)
+        super().__init__(defaults)
+
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def update(self, var):
+        k = self.global_state.get('step', 0) + 1
+        self.global_state["step"] = k
+
+        params = TensorList(var.params)
+        closure = var.closure
+        if closure is None:
+            raise RuntimeError("closure is required for SG2")
+        generator = self.get_generator(params[0].device, self.defaults["seed"])
+
+        h = self.get_settings(params, "h")
+        x_0 = params.clone()
+        n_samples = self.defaults["n_samples"]
+        H_hat = None
+
+        for i in range(n_samples):
+            # generate perturbation
+            cd = params.rademacher_like(generator=generator).mul_(h)
+
+            # one sided
+            if self.defaults["one_sided"]:
+                g_0 = TensorList(var.get_grad())
+                params.add_(cd)
+                closure()
+
+                g_p = params.grad.fill_none_(params)
+                delta_g = (g_p - g_0) * 2
+
+            # two sided
+            else:
+                params.add_(cd)
+                closure()
+                g_p = params.grad.fill_none_(params)
+
+                params.copy_(x_0)
+                params.sub_(cd)
+                closure()
+                g_n = params.grad.fill_none_(params)
+
+                delta_g = g_p - g_n
+
+            # restore params
+            params.set_(x_0)
+
+            # compute H hat
+            H_i = sg2_(
+                delta_g = delta_g.to_vec(),
+                cd = cd.to_vec(),
+            )
+
+            if H_hat is None: H_hat = H_i
+            else: H_hat += H_i
+
+        assert H_hat is not None
+        if n_samples > 1: H_hat /= n_samples
+
+        # update H
+        H = self.global_state.get("H", None)
+        if H is None: H = H_hat
+        else:
+            beta = self.defaults["beta"]
+            if beta is None: beta = k / (k+1)
+            H.lerp_(H_hat, 1-beta)
+
+        self.global_state["H"] = H
+
+
+    @torch.no_grad
+    def apply(self, var):
+        dir = _newton_step(
+            var=var,
+            H = self.global_state["H"],
+            damping = self.defaults["damping"],
+            inner = self.children.get("inner", None),
+            H_tfm=None,
+            eigval_fn=self.defaults["eigval_fn"],
+            use_lstsq=self.defaults["use_lstsq"],
+            g_proj=None,
+        )
+
+        var.update = vec_to_tensors(dir, var.params)
+        return var
+
+    def get_H(self,var=...):
+        return _get_H(self.global_state["H"], self.defaults["eigval_fn"])
+
+
+
+
+# two sided
+# we have g via x + d, x - d
+# H via g(x + d), g(x - d)
+# 1 is x, x+2d
+# 2 is x, x-2d
+# 5 evals in total
+
+# one sided
+# g via x, x + d
+# 1 is x, x + d
+# 2 is x, x - d
+# 3 evals and can use two sided for g_0
+
+class SPSA2(Module):
+    """second-order SPSA
+
+    SPSA2 with line search
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.SPSA2(),
+        tz.m.Backtracking()
+    )
+    ```
+
+    SPSA2 with trust region
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LevenbergMarquardt(tz.m.SPSA2()),
+    )
+    ```
+    """
+
+    def __init__(
+        self,
+        n_samples: int = 1,
+        h: float = 1e-2,
+        beta: float | None = None,
+        damping: float = 0,
+        eigval_fn=None,
+        use_lstsq: bool = True,
+        seed=None,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(n_samples=n_samples, h=h, beta=beta, damping=damping, eigval_fn=eigval_fn, seed=seed, use_lstsq=use_lstsq)
+        super().__init__(defaults)
+
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def update(self, var):
+        k = self.global_state.get('step', 0) + 1
+        self.global_state["step"] = k
+
+        params = TensorList(var.params)
+        closure = var.closure
+        if closure is None:
+            raise RuntimeError("closure is required for SPSA2")
+
+        generator = self.get_generator(params[0].device, self.defaults["seed"])
+
+        h = self.get_settings(params, "h")
+        x_0 = params.clone()
+        n_samples = self.defaults["n_samples"]
+        H_hat = None
+        g_0 = None
+
+        for i in range(n_samples):
+            # perturbations for g and H
+            cd_g = params.rademacher_like(generator=generator).mul_(h)
+            cd_H = params.rademacher_like(generator=generator).mul_(h)
+
+            # evaluate 4 points
+            x_p = x_0 + cd_g
+            x_n = x_0 - cd_g
+
+            params.set_(x_p)
+            f_p = closure(False)
+            params.add_(cd_H)
+            f_pp = closure(False)
+
+            params.set_(x_n)
+            f_n = closure(False)
+            params.add_(cd_H)
+            f_np = closure(False)
+
+            g_p_vec = (f_pp - f_p) / cd_H
+            g_n_vec = (f_np - f_n) / cd_H
+            delta_g = g_p_vec - g_n_vec
+
+            # restore params
+            params.set_(x_0)
+
+            # compute grad
+            g_i = (f_p - f_n) / (2 * cd_g)
+            if g_0 is None: g_0 = g_i
+            else: g_0 += g_i
+
+            # compute H hat
+            H_i = sg2_(
+                delta_g = delta_g.to_vec().div_(2.0),
+                cd = cd_g.to_vec(), # The interval is measured by the original 'cd'
+            )
+            if H_hat is None: H_hat = H_i
+            else: H_hat += H_i
+
+        assert g_0 is not None and H_hat is not None
+        if n_samples > 1:
+            g_0 /= n_samples
+            H_hat /= n_samples
+
+        # set grad to approximated grad
+        var.grad = g_0
+
+        # update H
+        H = self.global_state.get("H", None)
+        if H is None: H = H_hat
+        else:
+            beta = self.defaults["beta"]
+            if beta is None: beta = k / (k+1)
+            H.lerp_(H_hat, 1-beta)
+
+        self.global_state["H"] = H
+
+    @torch.no_grad
+    def apply(self, var):
+        dir = _newton_step(
+            var=var,
+            H = self.global_state["H"],
+            damping = self.defaults["damping"],
+            inner = self.children.get("inner", None),
+            H_tfm=None,
+            eigval_fn=self.defaults["eigval_fn"],
+            use_lstsq=self.defaults["use_lstsq"],
+            g_proj=None,
+        )
+
+        var.update = vec_to_tensors(dir, var.params)
+        return var
+
+    def get_H(self,var=...):
+        return _get_H(self.global_state["H"], self.defaults["eigval_fn"])

torchzero/modules/restarts/restars.py

@@ -60,18 +60,18 @@ class RestartStrategyBase(Module, ABC):
 
 
 class RestartOnStuck(RestartStrategyBase):
-    """Resets the state when update (difference in parameters) is close to zero for multiple steps in a row.
+    """Resets the state when update (difference in parameters) is zero for multiple steps in a row.
 
     Args:
         modules (Chainable | None):
             modules to reset. If None, resets all modules.
         tol (float, optional):
-            step is considered failed when maximum absolute parameter difference is smaller than this. Defaults to 1e-10.
+            step is considered failed when maximum absolute parameter difference is smaller than this. Defaults to None (uses twice the smallest respresentable number)
         n_tol (int, optional):
-            number of failed consequtive steps required to trigger a reset. Defaults to 4.
+            number of failed consequtive steps required to trigger a reset. Defaults to 10.
 
     """
-    def __init__(self, modules: Chainable | None, tol: float = 1e-10, n_tol: int = 4):
+    def __init__(self, modules: Chainable | None, tol: float | None = None, n_tol: int = 10):
         defaults = dict(tol=tol, n_tol=n_tol)
         super().__init__(defaults, modules)
 
@@ -82,6 +82,7 @@ class RestartOnStuck(RestartStrategyBase):
 
         params = TensorList(var.params)
         tol = self.defaults['tol']
+        if tol is None: tol = torch.finfo(params[0].dtype).tiny * 2
         n_tol = self.defaults['n_tol']
         n_bad = self.global_state.get('n_bad', 0)
 

torchzero/modules/second_order/__init__.py

@@ -1,4 +1,7 @@
-from .newton import Newton, InverseFreeNewton
+from .ifn import InverseFreeNewton
+from .inm import INM
+from .multipoint import SixthOrder3P, SixthOrder3PM2, SixthOrder5P, TwoPointNewton
+from .newton import Newton
 from .newton_cg import NewtonCG, NewtonCGSteihaug
-from .nystrom import NystromSketchAndSolve, NystromPCG
-from .multipoint import SixthOrder3P, SixthOrder5P, TwoPointNewton, SixthOrder3PM2
+from .nystrom import NystromPCG, NystromSketchAndSolve
+from .rsn import RSN