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

torchzero/modules/line_search/line_search.py

@@ -8,15 +8,16 @@ from typing import Any
 import numpy as np
 import torch
 
-from ...core import Module, Target, Vars
+from ...core import Module, Target, Var
 from ...utils import tofloat
 
 
 class MaxLineSearchItersReached(Exception): pass
 
 
-class LineSearch(Module, ABC):
+class LineSearchBase(Module, ABC):
     """Base class for line searches.
+
     This is an abstract class, to use it, subclass it and override `search`.
 
     Args:
@@ -26,6 +27,62 @@ class LineSearch(Module, ABC):
             the objective this many times, and step size with the lowest loss value will be used.
             This is useful when passing `make_objective` to an external library which
             doesn't have a maxiter option. Defaults to None.
+
+    Other useful methods:
+        * `evaluate_step_size` - returns loss with a given scalar step size
+        * `evaluate_step_size_loss_and_derivative` - returns loss and directional derivative with a given scalar step size
+        * `make_objective` - creates a function that accepts a scalar step size and returns loss. This can be passed to a scalar solver, such as scipy.optimize.minimize_scalar.
+        * `make_objective_with_derivative` - creates a function that accepts a scalar step size and returns a tuple with loss and directional derivative. This can be passed to a scalar solver.
+
+    Examples:
+        #### Basic line search
+
+        This evaluates all step sizes in a range by using the :code:`self.evaluate_step_size` method.
+
+        .. code-block:: python
+
+            class GridLineSearch(LineSearch):
+                def __init__(self, start, end, num):
+                    defaults = dict(start=start,end=end,num=num)
+                    super().__init__(defaults)
+
+                @torch.no_grad
+                def search(self, update, var):
+                    settings = self.settings[var.params[0]]
+                    start = settings["start"]
+                    end = settings["end"]
+                    num = settings["num"]
+
+                    lowest_loss = float("inf")
+                    best_step_size = best_step_size
+
+                    for step_size in torch.linspace(start,end,num):
+                        loss = self.evaluate_step_size(step_size.item(), var=var, backward=False)
+                        if loss < lowest_loss:
+                            lowest_loss = loss
+                            best_step_size = step_size
+
+                    return best_step_size
+
+        #### Using external solver via self.make_objective
+
+        Here we let :code:`scipy.optimize.minimize_scalar` solver find the best step size via :code:`self.make_objective`
+
+        .. code-block:: python
+
+            class ScipyMinimizeScalar(LineSearch):
+                def __init__(self, method: str | None = None):
+                    defaults = dict(method=method)
+                    super().__init__(defaults)
+
+                @torch.no_grad
+                def search(self, update, var):
+                    objective = self.make_objective(var=var)
+                    method = self.settings[var.params[0]]["method"]
+
+                    res = self.scopt.minimize_scalar(objective, method=method)
+                    return res.x
+
     """
     def __init__(self, defaults: dict[str, Any] | None, maxiter: int | None = None):
         super().__init__(defaults)
@@ -62,12 +119,12 @@ class LineSearch(Module, ABC):
         if any(a!=0 for a in alpha):
             torch._foreach_add_(params, torch._foreach_mul(update, alpha))
 
-    def _loss(self, step_size: float, vars: Vars, closure, params: list[torch.Tensor],
+    def _loss(self, step_size: float, var: Var, closure, params: list[torch.Tensor],
               update: list[torch.Tensor], backward:bool=False) -> float:
 
         # if step_size is 0, we might already know the loss
-        if (vars.loss is not None) and (step_size == 0):
-            return tofloat(vars.loss)
+        if (var.loss is not None) and (step_size == 0):
+            return tofloat(var.loss)
 
         # check max iter
         if self._maxiter is not None and self._current_iter >= self._maxiter: raise MaxLineSearchItersReached
@@ -85,23 +142,23 @@ class LineSearch(Module, ABC):
             self._lowest_loss = tofloat(loss)
             self._best_step_size = step_size
 
-        # if evaluated loss at step size 0, set it to vars.loss
+        # if evaluated loss at step size 0, set it to var.loss
         if step_size == 0:
-            vars.loss = loss
-            if backward: vars.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+            var.loss = loss
+            if backward: var.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
 
         return tofloat(loss)
 
-    def _loss_derivative(self, step_size: float, vars: Vars, closure,
+    def _loss_derivative(self, step_size: float, var: Var, closure,
                          params: list[torch.Tensor], update: list[torch.Tensor]):
         # if step_size is 0, we might already know the derivative
-        if (vars.grad is not None) and (step_size == 0):
-            loss = self._loss(step_size=step_size,vars=vars,closure=closure,params=params,update=update,backward=False)
-            derivative = - sum(t.sum() for t in torch._foreach_mul(vars.grad, update))
+        if (var.grad is not None) and (step_size == 0):
+            loss = self._loss(step_size=step_size,var=var,closure=closure,params=params,update=update,backward=False)
+            derivative = - sum(t.sum() for t in torch._foreach_mul(var.grad, update))
 
         else:
             # loss with a backward pass sets params.grad
-            loss = self._loss(step_size=step_size,vars=vars,closure=closure,params=params,update=update,backward=True)
+            loss = self._loss(step_size=step_size,var=var,closure=closure,params=params,update=update,backward=True)
 
             # directional derivative
             derivative = - sum(t.sum() for t in torch._foreach_mul([p.grad if p.grad is not None
@@ -109,73 +166,74 @@ class LineSearch(Module, ABC):
 
         return loss, tofloat(derivative)
 
-    def evaluate_step_size(self, step_size: float, vars: Vars, backward:bool=False):
-        closure = vars.closure
+    def evaluate_step_size(self, step_size: float, var: Var, backward:bool=False):
+        closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return self._loss(step_size=step_size, vars=vars, closure=closure, params=vars.params,update=vars.get_update(),backward=backward)
+        return self._loss(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update(),backward=backward)
 
-    def evaluate_step_size_loss_and_derivative(self, step_size: float, vars: Vars):
-        closure = vars.closure
+    def evaluate_step_size_loss_and_derivative(self, step_size: float, var: Var):
+        closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return self._loss_derivative(step_size=step_size, vars=vars, closure=closure, params=vars.params,update=vars.get_update())
+        return self._loss_derivative(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update())
 
-    def make_objective(self, vars: Vars, backward:bool=False):
-        closure = vars.closure
+    def make_objective(self, var: Var, backward:bool=False):
+        closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return partial(self._loss, vars=vars, closure=closure, params=vars.params, update=vars.get_update(), backward=backward)
+        return partial(self._loss, var=var, closure=closure, params=var.params, update=var.get_update(), backward=backward)
 
-    def make_objective_with_derivative(self, vars: Vars):
-        closure = vars.closure
+    def make_objective_with_derivative(self, var: Var):
+        closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return partial(self._loss_derivative, vars=vars, closure=closure, params=vars.params, update=vars.get_update())
+        return partial(self._loss_derivative, var=var, closure=closure, params=var.params, update=var.get_update())
 
     @abstractmethod
-    def search(self, update: list[torch.Tensor], vars: Vars) -> float:
+    def search(self, update: list[torch.Tensor], var: Var) -> float:
         """Finds the step size to use"""
 
     @torch.no_grad
-    def step(self, vars: Vars) -> Vars:
+    def step(self, var: Var) -> Var:
         self._reset()
-        params = vars.params
-        update = vars.get_update()
+        params = var.params
+        update = var.get_update()
 
         try:
-            step_size = self.search(update=update, vars=vars)
+            step_size = self.search(update=update, var=var)
         except MaxLineSearchItersReached:
             step_size = self._best_step_size
 
         # set loss_approx
-        if vars.loss_approx is None: vars.loss_approx = self._lowest_loss
+        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 vars.is_last:
+        if var.is_last:
 
-            if vars.last_module_lrs is None:
+            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 vars.last_module_lrs], params=params, update=update)
+                self._set_per_parameter_step_size_([step_size*lr for lr in var.last_module_lrs], params=params, update=update)
 
-            vars.stop = True; vars.skip_update = True
-            return vars
+            var.stop = True; var.skip_update = True
+            return var
 
         # revert parameters and multiply update by step size
         self.set_step_size_(0, params=params, update=update)
-        torch._foreach_mul_(vars.update, step_size)
-        return vars
+        torch._foreach_mul_(var.update, step_size)
+        return var
 
 
-class GridLineSearch(LineSearch):
-    """Mostly for testing, this is not practical"""
-    def __init__(self, start, end, num):
-        defaults = dict(start=start,end=end,num=num)
-        super().__init__(defaults)
 
-    @torch.no_grad
-    def search(self, update, vars):
-        start,end,num=itemgetter('start','end','num')(self.settings[vars.params[0]])
+# class GridLineSearch(LineSearch):
+#     """Mostly for testing, this is not practical"""
+#     def __init__(self, start, end, num):
+#         defaults = dict(start=start,end=end,num=num)
+#         super().__init__(defaults)
+
+#     @torch.no_grad
+#     def search(self, update, var):
+#         start,end,num=itemgetter('start','end','num')(self.settings[var.params[0]])
 
-        for lr in torch.linspace(start,end,num):
-            self.evaluate_step_size(lr.item(), vars=vars, backward=False)
+#         for lr in torch.linspace(start,end,num):
+#             self.evaluate_step_size(lr.item(), var=var, backward=False)
 
-        return self._best_step_size
+#         return self._best_step_size

torchzero/modules/line_search/polynomial.py

@@ -0,0 +1,233 @@
+import numpy as np
+import torch
+
+from .line_search import LineSearchBase
+
+
+# polynomial interpolation
+# this code is from https://github.com/hjmshi/PyTorch-LBFGS/blob/master/functions/LBFGS.py
+# PyTorch-LBFGS: A PyTorch Implementation of L-BFGS
+def polyinterp(points, x_min_bound=None, x_max_bound=None, plot=False):
+    """
+    Gives the minimizer and minimum of the interpolating polynomial over given points
+    based on function and derivative information. Defaults to bisection if no critical
+    points are valid.
+
+    Based on polyinterp.m Matlab function in minFunc by Mark Schmidt with some slight
+    modifications.
+
+    Implemented by: Hao-Jun Michael Shi and Dheevatsa Mudigere
+    Last edited 12/6/18.
+
+    Inputs:
+        points (nparray): two-dimensional array with each point of form [x f g]
+        x_min_bound (float): minimum value that brackets minimum (default: minimum of points)
+        x_max_bound (float): maximum value that brackets minimum (default: maximum of points)
+        plot (bool): plot interpolating polynomial
+
+    Outputs:
+        x_sol (float): minimizer of interpolating polynomial
+        F_min (float): minimum of interpolating polynomial
+
+    Note:
+      . Set f or g to np.nan if they are unknown
+
+    """
+    no_points = points.shape[0]
+    order = np.sum(1 - np.isnan(points[:, 1:3]).astype('int')) - 1
+
+    x_min = np.min(points[:, 0])
+    x_max = np.max(points[:, 0])
+
+    # compute bounds of interpolation area
+    if x_min_bound is None:
+        x_min_bound = x_min
+    if x_max_bound is None:
+        x_max_bound = x_max
+
+    # explicit formula for quadratic interpolation
+    if no_points == 2 and order == 2 and plot is False:
+        # Solution to quadratic interpolation is given by:
+        # a = -(f1 - f2 - g1(x1 - x2))/(x1 - x2)^2
+        # x_min = x1 - g1/(2a)
+        # if x1 = 0, then is given by:
+        # x_min = - (g1*x2^2)/(2(f2 - f1 - g1*x2))
+
+        if points[0, 0] == 0:
+            x_sol = -points[0, 2] * points[1, 0] ** 2 / (2 * (points[1, 1] - points[0, 1] - points[0, 2] * points[1, 0]))
+        else:
+            a = -(points[0, 1] - points[1, 1] - points[0, 2] * (points[0, 0] - points[1, 0])) / (points[0, 0] - points[1, 0]) ** 2
+            x_sol = points[0, 0] - points[0, 2]/(2*a)
+
+        x_sol = np.minimum(np.maximum(x_min_bound, x_sol), x_max_bound)
+
+    # explicit formula for cubic interpolation
+    elif no_points == 2 and order == 3 and plot is False:
+        # Solution to cubic interpolation is given by:
+        # d1 = g1 + g2 - 3((f1 - f2)/(x1 - x2))
+        # d2 = sqrt(d1^2 - g1*g2)
+        # x_min = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2))
+        d1 = points[0, 2] + points[1, 2] - 3 * ((points[0, 1] - points[1, 1]) / (points[0, 0] - points[1, 0]))
+        d2 = np.sqrt(d1 ** 2 - points[0, 2] * points[1, 2])
+        if np.isreal(d2):
+            x_sol = points[1, 0] - (points[1, 0] - points[0, 0]) * ((points[1, 2] + d2 - d1) / (points[1, 2] - points[0, 2] + 2 * d2))
+            x_sol = np.minimum(np.maximum(x_min_bound, x_sol), x_max_bound)
+        else:
+            x_sol = (x_max_bound + x_min_bound)/2
+
+    # solve linear system
+    else:
+        # define linear constraints
+        A = np.zeros((0, order + 1))
+        b = np.zeros((0, 1))
+
+        # add linear constraints on function values
+        for i in range(no_points):
+            if not np.isnan(points[i, 1]):
+                constraint = np.zeros((1, order + 1))
+                for j in range(order, -1, -1):
+                    constraint[0, order - j] = points[i, 0] ** j
+                A = np.append(A, constraint, 0)
+                b = np.append(b, points[i, 1])
+
+        # add linear constraints on gradient values
+        for i in range(no_points):
+            if not np.isnan(points[i, 2]):
+                constraint = np.zeros((1, order + 1))
+                for j in range(order):
+                    constraint[0, j] = (order - j) * points[i, 0] ** (order - j - 1)
+                A = np.append(A, constraint, 0)
+                b = np.append(b, points[i, 2])
+
+        # check if system is solvable
+        if A.shape[0] != A.shape[1] or np.linalg.matrix_rank(A) != A.shape[0]:
+            x_sol = (x_min_bound + x_max_bound)/2
+            f_min = np.inf
+        else:
+            # solve linear system for interpolating polynomial
+            coeff = np.linalg.solve(A, b)
+
+            # compute critical points
+            dcoeff = np.zeros(order)
+            for i in range(len(coeff) - 1):
+                dcoeff[i] = coeff[i] * (order - i)
+
+            crit_pts = np.array([x_min_bound, x_max_bound])
+            crit_pts = np.append(crit_pts, points[:, 0])
+
+            if not np.isinf(dcoeff).any():
+                roots = np.roots(dcoeff)
+                crit_pts = np.append(crit_pts, roots)
+
+            # test critical points
+            f_min = np.inf
+            x_sol = (x_min_bound + x_max_bound) / 2 # defaults to bisection
+            for crit_pt in crit_pts:
+                if np.isreal(crit_pt) and crit_pt >= x_min_bound and crit_pt <= x_max_bound:
+                    F_cp = np.polyval(coeff, crit_pt)
+                    if np.isreal(F_cp) and F_cp < f_min:
+                        x_sol = np.real(crit_pt)
+                        f_min = np.real(F_cp)
+
+            if(plot):
+                import matplotlib.pyplot as plt
+                plt.figure()
+                x = np.arange(x_min_bound, x_max_bound, (x_max_bound - x_min_bound)/10000)
+                f = np.polyval(coeff, x)
+                plt.plot(x, f)
+                plt.plot(x_sol, f_min, 'x')
+
+    return x_sol
+
+
+
+# class PolynomialLineSearch(LineSearch):
+#     """TODO
+
+#     Line search via polynomial interpolation.
+
+#     Args:
+#         init (float, optional): Initial step size. Defaults to 1.0.
+#         c1 (float, optional): Acceptance value for weak wolfe condition. Defaults to 1e-4.
+#         c2 (float, optional): Acceptance value for strong wolfe condition (set to 0.1 for conjugate gradient). Defaults to 0.9.
+#         maxiter (int, optional): Maximum number of line search iterations. Defaults to 25.
+#         maxzoom (int, optional): Maximum number of zoom iterations. Defaults to 10.
+#         expand (float, optional): Expansion factor (multipler to step size when weak condition not satisfied). Defaults to 2.0.
+#         adaptive (bool, optional):
+#             when enabled, if line search failed, initial step size is reduced.
+#             Otherwise it is reset to initial value. Defaults to True.
+#         plus_minus (bool, optional):
+#             If enabled and the direction is not descent direction, performs line search in opposite direction. Defaults to False.
+
+
+#     Examples:
+#         Conjugate gradient method with strong wolfe line search. Nocedal, Wright recommend setting c2 to 0.1 for CG.
+
+#         .. code-block:: python
+
+#             opt = tz.Modular(
+#                 model.parameters(),
+#                 tz.m.PolakRibiere(),
+#                 tz.m.StrongWolfe(c2=0.1)
+#             )
+
+#         LBFGS strong wolfe line search:
+
+#         .. code-block:: python
+
+#             opt = tz.Modular(
+#                 model.parameters(),
+#                 tz.m.LBFGS(),
+#                 tz.m.StrongWolfe()
+#             )
+
+#     """
+#     def __init__(
+#         self,
+#         init: float = 1.0,
+#         c1: float = 1e-4,
+#         c2: float = 0.9,
+#         maxiter: int = 25,
+#         maxzoom: int = 10,
+#         # a_max: float = 1e10,
+#         expand: float = 2.0,
+#         adaptive = True,
+#         plus_minus = False,
+#     ):
+#         defaults=dict(init=init,c1=c1,c2=c2,maxiter=maxiter,maxzoom=maxzoom,
+#                       expand=expand, adaptive=adaptive, plus_minus=plus_minus)
+#         super().__init__(defaults=defaults)
+
+#         self.global_state['initial_scale'] = 1.0
+#         self.global_state['beta_scale'] = 1.0
+
+#     @torch.no_grad
+#     def search(self, update, var):
+#         objective = self.make_objective_with_derivative(var=var)
+
+#         init, c1, c2, maxiter, maxzoom, expand, adaptive, plus_minus = itemgetter(
+#             'init', 'c1', 'c2', 'maxiter', 'maxzoom',
+#             'expand', 'adaptive', 'plus_minus')(self.settings[var.params[0]])
+
+#         f_0, g_0 = objective(0)
+
+#         step_size,f_a = strong_wolfe(
+#             objective,
+#             f_0=f_0, g_0=g_0,
+#             init=init * self.global_state.setdefault("initial_scale", 1),
+#             c1=c1,
+#             c2=c2,
+#             maxiter=maxiter,
+#             maxzoom=maxzoom,
+#             expand=expand,
+#             plus_minus=plus_minus,
+#         )
+
+#         if f_a is not None and (f_a > f_0 or _notfinite(f_a)): step_size = None
+#         if step_size is not None and step_size != 0 and not _notfinite(step_size):
+#             self.global_state['initial_scale'] = min(1.0, self.global_state['initial_scale'] * math.sqrt(2))
+#             return step_size
+
+#         # fallback to backtracking on fail
+#         if adaptive: self.global_state['initial_scale'] *= 0.5
+#         return 0

torchzero/modules/line_search/scipy.py

@@ -3,10 +3,25 @@ from operator import itemgetter
 
 import torch
 
-from .line_search import LineSearch
+from .line_search import LineSearchBase
 
 
-class ScipyMinimizeScalar(LineSearch):
+class ScipyMinimizeScalar(LineSearchBase):
+    """Line search via :code:`scipy.optimize.minimize_scalar` which implements brent, golden search and bounded brent methods.
+
+    Args:
+        method (str | None, optional): "brent", "golden" or "bounded". Defaults to None.
+        maxiter (int | None, optional): maximum number of function evaluations the line search is allowed to perform. Defaults to None.
+        bracket (Sequence | None, optional):
+            Either a triple (xa, xb, xc) satisfying xa < xb < xc and func(xb) < func(xa) and  func(xb) < func(xc), or a pair (xa, xb) to be used as initial points for a downhill bracket search. Defaults to None.
+        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.
+        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
+
+    """
     def __init__(
         self,
         method: str | None = None,
@@ -24,10 +39,10 @@ class ScipyMinimizeScalar(LineSearch):
 
 
     @torch.no_grad
-    def search(self, update, vars):
-        objective = self.make_objective(vars=vars)
+    def search(self, update, var):
+        objective = self.make_objective(var=var)
         method, bracket, bounds, tol, options, maxiter = itemgetter(
-            'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.settings[vars.params[0]])
+            'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.settings[var.params[0]])
 
         if maxiter is not None:
             options = dict(options) if isinstance(options, Mapping) else {}

torchzero/modules/line_search/strong_wolfe.py

@@ -1,3 +1,4 @@
+"""this needs to be reworked maybe but it also works"""
 import math
 import warnings
 from operator import itemgetter
@@ -5,8 +6,7 @@ from operator import itemgetter
 import torch
 from torch.optim.lbfgs import _cubic_interpolate
 
-from .line_search import LineSearch
-from .backtracking import backtracking_line_search
+from .line_search import LineSearchBase
 from ...utils import totensor
 
 
@@ -182,7 +182,47 @@ def _notfinite(x):
     if isinstance(x, torch.Tensor): return not torch.isfinite(x).all()
     return not math.isfinite(x)
 
-class StrongWolfe(LineSearch):
+class StrongWolfe(LineSearchBase):
+    """Cubic interpolation line search satisfying Strong Wolfe condition.
+
+    Args:
+        init (float, optional): Initial step size. Defaults to 1.0.
+        c1 (float, optional): Acceptance value for weak wolfe condition. Defaults to 1e-4.
+        c2 (float, optional): Acceptance value for strong wolfe condition (set to 0.1 for conjugate gradient). Defaults to 0.9.
+        maxiter (int, optional): Maximum number of line search iterations. Defaults to 25.
+        maxzoom (int, optional): Maximum number of zoom iterations. Defaults to 10.
+        expand (float, optional): Expansion factor (multipler to step size when weak condition not satisfied). Defaults to 2.0.
+        use_prev (bool, optional):
+            if True, previous step size is used as the initial step size on the next step.
+        adaptive (bool, optional):
+            when enabled, if line search failed, initial step size is reduced.
+            Otherwise it is reset to initial value. Defaults to True.
+        plus_minus (bool, optional):
+            If enabled and the direction is not descent direction, performs line search in opposite direction. Defaults to False.
+
+
+    Examples:
+        Conjugate gradient method with strong wolfe line search. Nocedal, Wright recommend setting c2 to 0.1 for CG.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.PolakRibiere(),
+                tz.m.StrongWolfe(c2=0.1)
+            )
+
+        LBFGS strong wolfe line search:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.LBFGS(),
+                tz.m.StrongWolfe()
+            )
+
+    """
     def __init__(
         self,
         init: float = 1.0,
@@ -192,26 +232,27 @@ class StrongWolfe(LineSearch):
         maxzoom: int = 10,
         # a_max: float = 1e10,
         expand: float = 2.0,
+        use_prev: bool = False,
         adaptive = True,
-        fallback = False,
         plus_minus = False,
     ):
         defaults=dict(init=init,c1=c1,c2=c2,maxiter=maxiter,maxzoom=maxzoom,
-                      expand=expand, adaptive=adaptive, fallback=fallback, plus_minus=plus_minus)
+                      expand=expand, adaptive=adaptive, plus_minus=plus_minus,use_prev=use_prev)
         super().__init__(defaults=defaults)
 
         self.global_state['initial_scale'] = 1.0
         self.global_state['beta_scale'] = 1.0
 
     @torch.no_grad
-    def search(self, update, vars):
-        objective = self.make_objective_with_derivative(vars=vars)
+    def search(self, update, var):
+        objective = self.make_objective_with_derivative(var=var)
 
-        init, c1, c2, maxiter, maxzoom, expand, adaptive, fallback, plus_minus = itemgetter(
+        init, c1, c2, maxiter, maxzoom, expand, adaptive, plus_minus, use_prev = itemgetter(
             'init', 'c1', 'c2', 'maxiter', 'maxzoom',
-            'expand', 'adaptive', 'fallback', 'plus_minus')(self.settings[vars.params[0]])
+            'expand', 'adaptive', 'plus_minus', 'use_prev')(self.settings[var.params[0]])
 
         f_0, g_0 = objective(0)
+        if use_prev: init = self.global_state.get('prev_alpha', init)
 
         step_size,f_a = strong_wolfe(
             objective,
@@ -228,33 +269,8 @@ class StrongWolfe(LineSearch):
         if f_a is not None and (f_a > f_0 or _notfinite(f_a)): step_size = None
         if step_size is not None and step_size != 0 and not _notfinite(step_size):
             self.global_state['initial_scale'] = min(1.0, self.global_state['initial_scale'] * math.sqrt(2))
+            self.global_state['prev_alpha'] = step_size
             return step_size
 
-        # fallback to backtracking on fail
         if adaptive: self.global_state['initial_scale'] *= 0.5
-        if not fallback: return 0
-
-        objective = self.make_objective(vars=vars)
-
-        # # directional derivative
-        g_0 = -sum(t.sum() for t in torch._foreach_mul(vars.get_grad(), vars.get_update()))
-
-        step_size = backtracking_line_search(
-            objective,
-            g_0,
-            init=init * self.global_state["initial_scale"],
-            beta=0.5 * self.global_state["beta_scale"],
-            c=c1,
-            maxiter=maxiter * 2,
-            a_min=None,
-            try_negative=plus_minus,
-        )
-
-        # found an alpha that reduces loss
-        if step_size is not None:
-            self.global_state['beta_scale'] = min(1.0, self.global_state.get('beta_scale', 1) * math.sqrt(1.5))
-            return step_size
-
-        # on fail reduce beta scale value
-        self.global_state['beta_scale'] /= 1.5
-        return 0
+        return 0

torchzero/modules/misc/__init__.py

@@ -0,0 +1,27 @@
+from .debug import PrintLoss, PrintParams, PrintShape, PrintUpdate
+from .escape import EscapeAnnealing
+from .gradient_accumulation import GradientAccumulation
+from .misc import (
+    DivByLoss,
+    FillLoss,
+    GradSign,
+    GraftGradToUpdate,
+    GraftToGrad,
+    GraftToParams,
+    HpuEstimate,
+    LastAbsoluteRatio,
+    LastDifference,
+    LastGradDifference,
+    LastProduct,
+    LastRatio,
+    MulByLoss,
+    NoiseSign,
+    Previous,
+    RandomHvp,
+    Relative,
+    UpdateSign,
+)
+from .multistep import Multistep, NegateOnLossIncrease, Online, Sequential
+from .regularization import Dropout, PerturbWeights, WeightDropout
+from .split import Split
+from .switch import Alternate, Switch