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

torchzero/modules/line_search/adaptive.py

@@ -1,58 +1,73 @@
 import math
+from bisect import insort
+from collections import deque
 from collections.abc import Callable
 from operator import itemgetter
 
+import numpy as np
 import torch
 
-from .line_search import LineSearchBase
-
+from .line_search import LineSearchBase, TerminationCondition, termination_condition
 
 
 def adaptive_tracking(
     f,
-    x_0,
+    a_init,
     maxiter: int,
     nplus: float = 2,
     nminus: float = 0.5,
+    f_0 = None,
 ):
-    f_0 = f(0)
+    niter = 0
+    if f_0 is None: f_0 = f(0)
 
-    t = x_0
-    f_t = f(t)
+    a = a_init
+    f_a = f(a)
 
     # backtrack
-    if f_t > f_0:
-        while f_t > f_0:
+    a_prev = a
+    f_prev = math.inf
+    if (f_a > f_0) or (not math.isfinite(f_a)):
+        while (f_a < f_prev) or not math.isfinite(f_a):
+            a_prev, f_prev = a, f_a
             maxiter -= 1
-            if maxiter < 0: return 0, f_0
-            t = t*nminus
-            f_t = f(t)
-        return t, f_t
+            if maxiter < 0: break
+
+            a = a*nminus
+            f_a = f(a)
+            niter += 1
+
+        if f_prev < f_0: return a_prev, f_prev, niter
+        return 0, f_0, niter
 
     # forwardtrack
-    f_prev = f_t
-    t *= nplus
-    f_t = f(t)
-    if f_prev < f_t: return t / nplus, f_prev
-    while f_prev >= f_t:
+    a_prev = a
+    f_prev = math.inf
+    while (f_a <= f_prev) and math.isfinite(f_a):
+        a_prev, f_prev = a, f_a
         maxiter -= 1
-        if maxiter < 0: return t, f_t
-        f_prev = f_t
-        t *= nplus
-        f_t = f(t)
-    return t / nplus, f_prev
+        if maxiter < 0: break
+
+        a *= nplus
+        f_a = f(a)
+        niter+= 1
+
+    if f_prev < f_0: return a_prev, f_prev, niter
+    return 0, f_0, niter
+
 
-class AdaptiveLineSearch(LineSearchBase):
-    """Adaptive line search, similar to backtracking but also has forward tracking mode.
-    Currently doesn't check for weak curvature condition.
+class AdaptiveTracking(LineSearchBase):
+    """A line search that evaluates previous step size, if value increased, backtracks until the value stops decreasing,
+    otherwise forward-tracks until value stops decreasing.
 
     Args:
         init (float, optional): initial step size. Defaults to 1.0.
-        beta (float, optional): multiplies each consecutive step size by this value. Defaults to 0.5.
-        maxiter (int, optional): Maximum line search function evaluations. Defaults to 10.
+        nplus (float, optional): multiplier to step size if initial step size is optimal. Defaults to 2.
+        nminus (float, optional): multiplier to step size if initial step size is too big. Defaults to 0.5.
+        maxiter (int, optional): maximum number of function evaluations per step. Defaults to 10.
         adaptive (bool, optional):
-            when enabled, if line search failed, beta size is reduced.
-            Otherwise it is reset to initial value. Defaults to True.
+            when enabled, if line search failed, step size will continue decreasing on the next step.
+            Otherwise it will restart the line search from ``init`` step size. Defaults to True.
     """
     def __init__(
         self,
@@ -62,38 +77,48 @@ class AdaptiveLineSearch(LineSearchBase):
         maxiter: int = 10,
         adaptive=True,
     ):
-        defaults=dict(init=init,nplus=nplus,nminus=nminus,maxiter=maxiter,adaptive=adaptive,)
+        defaults=dict(init=init,nplus=nplus,nminus=nminus,maxiter=maxiter,adaptive=adaptive)
         super().__init__(defaults=defaults)
-        self.global_state['beta_scale'] = 1.0
 
     def reset(self):
         super().reset()
-        self.global_state['beta_scale'] = 1.0
 
     @torch.no_grad
     def search(self, update, var):
         init, nplus, nminus, maxiter, adaptive = itemgetter(
-            'init', 'nplus', 'nminus', 'maxiter', 'adaptive')(self.settings[var.params[0]])
+            'init', 'nplus', 'nminus', 'maxiter', 'adaptive')(self.defaults)
 
         objective = self.make_objective(var=var)
 
-        # # directional derivative
-        # d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), var.get_update()))
+        # scale a_prev
+        a_prev = self.global_state.get('a_prev', init)
+        if adaptive: a_prev = a_prev * self.global_state.get('init_scale', 1)
 
-        # scale beta (beta is multiplicative and i think may be better than scaling initial step size)
-        beta_scale = self.global_state.get('beta_scale', 1)
-        x_prev = self.global_state.get('prev_x', 1)
+        a_init = a_prev
+        if a_init < torch.finfo(var.params[0].dtype).tiny * 2:
+            a_init = torch.finfo(var.params[0].dtype).max / 2
 
-        if adaptive: nminus = nminus * beta_scale
-
-
-        step_size, f = adaptive_tracking(objective, x_prev, maxiter, nplus=nplus, nminus=nminus)
+        step_size, f, niter = adaptive_tracking(
+            objective,
+            a_init=a_init,
+            maxiter=maxiter,
+            nplus=nplus,
+            nminus=nminus,
+        )
 
         # found an alpha that reduces loss
         if step_size != 0:
-            self.global_state['beta_scale'] = min(1.0, self.global_state['beta_scale'] * math.sqrt(1.5))
+            assert (var.loss is None) or (math.isfinite(f) and f < var.loss)
+            self.global_state['init_scale'] = 1
+
+            # if niter == 1, forward tracking failed to decrease function value compared to f_a_prev
+            if niter == 1 and step_size >= a_init: step_size *= nminus
+
+            self.global_state['a_prev'] = step_size
             return step_size
 
         # on fail reduce beta scale value
-        self.global_state['beta_scale'] /= 1.5
+        self.global_state['init_scale'] = self.global_state.get('init_scale', 1) * nminus**maxiter
+        self.global_state['a_prev'] = init
         return 0
+

torchzero/modules/line_search/backtracking.py

@@ -4,7 +4,7 @@ from operator import itemgetter
 
 import torch
 
-from .line_search import LineSearchBase
+from .line_search import LineSearchBase, TerminationCondition, termination_condition
 
 
 def backtracking_line_search(
@@ -14,7 +14,7 @@ def backtracking_line_search(
     beta: float = 0.5,
     c: float = 1e-4,
     maxiter: int = 10,
-    try_negative: bool = False,
+    condition: TerminationCondition = 'armijo',
 ) -> float | None:
     """
 
@@ -31,16 +31,20 @@ def backtracking_line_search(
     """
 
     a = init
-    f_x = f(0)
+    f_0 = f(0)
     f_prev = None
 
     for iteration in range(maxiter):
         f_a = f(a)
+        if not math.isfinite(f_a):
+            a *= beta
+            continue
 
-        if (f_prev is not None) and (f_a > f_prev) and (f_prev < f_x): return a / beta
+        if (f_prev is not None) and (f_a > f_prev) and (f_prev < f_0):
+            return a / beta # new value is larger than previous value
         f_prev = f_a
 
-        if f_a < f_x + c * a * min(g_0, 0): # pyright: ignore[reportArgumentType]
+        if termination_condition(condition, f_0=f_0, g_0=g_0, f_a=f_a, g_a=None, a=a, c=c):
             # found an acceptable alpha
             return a
 
@@ -48,53 +52,45 @@ def backtracking_line_search(
         a *= beta
 
     # fail
-    if try_negative:
-        def inv_objective(alpha): return f(-alpha)
-
-        v = backtracking_line_search(
-            inv_objective,
-            g_0=-g_0,
-            beta=beta,
-            c=c,
-            maxiter=maxiter,
-            try_negative=False,
-        )
-        if v is not None: return -v
-
     return None
 
 class Backtracking(LineSearchBase):
-    """Backtracking line search satisfying the Armijo condition.
+    """Backtracking line search.
 
     Args:
         init (float, optional): initial step size. Defaults to 1.0.
         beta (float, optional): multiplies each consecutive step size by this value. Defaults to 0.5.
-        c (float, optional): acceptance value for Armijo condition. Defaults to 1e-4.
-        maxiter (int, optional): Maximum line search function evaluations. Defaults to 10.
+        c (float, optional): sufficient decrease condition. Defaults to 1e-4.
+        condition (TerminationCondition, optional):
+            termination condition, only ones that do not use gradient at f(x+a*d) can be specified.
+            - "armijo" - sufficient decrease condition.
+            - "decrease" - any decrease in objective function value satisfies the condition.
+
+            "goldstein" can techincally be specified but it doesn't make sense because there is not zoom stage.
+            Defaults to 'armijo'.
+        maxiter (int, optional): maximum number of function evaluations per step. Defaults to 10.
         adaptive (bool, optional):
-            when enabled, if line search failed, beta is reduced.
-            Otherwise it is reset to initial value. Defaults to True.
-        try_negative (bool, optional): Whether to perform line search in opposite direction on fail. Defaults to False.
+            when enabled, if line search failed, step size will continue decreasing on the next step.
+            Otherwise it will restart the line search from ``init`` step size. Defaults to True.
 
     Examples:
-        Gradient descent with backtracking line search:
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Backtracking()
-            )
-
-        LBFGS with backtracking line search:
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.LBFGS(),
-                tz.m.Backtracking()
-            )
+    Gradient descent with backtracking line search:
+
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Backtracking()
+    )
+    ```
+
+    L-BFGS with backtracking line search:
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LBFGS(),
+        tz.m.Backtracking()
+    )
+    ```
 
     """
     def __init__(
@@ -102,41 +98,47 @@ class Backtracking(LineSearchBase):
         init: float = 1.0,
         beta: float = 0.5,
         c: float = 1e-4,
+        condition: TerminationCondition = 'armijo',
         maxiter: int = 10,
         adaptive=True,
-        try_negative: bool = False,
     ):
-        defaults=dict(init=init,beta=beta,c=c,maxiter=maxiter,adaptive=adaptive, try_negative=try_negative)
+        defaults=dict(init=init,beta=beta,c=c,condition=condition,maxiter=maxiter,adaptive=adaptive)
         super().__init__(defaults=defaults)
-        self.global_state['beta_scale'] = 1.0
 
     def reset(self):
         super().reset()
-        self.global_state['beta_scale'] = 1.0
 
     @torch.no_grad
     def search(self, update, var):
-        init, beta, c, maxiter, adaptive, try_negative = itemgetter(
-            'init', 'beta', 'c', 'maxiter', 'adaptive', 'try_negative')(self.settings[var.params[0]])
+        init, beta, c, condition, maxiter, adaptive = itemgetter(
+            'init', 'beta', 'c', 'condition', 'maxiter', 'adaptive')(self.defaults)
 
         objective = self.make_objective(var=var)
 
         # # directional derivative
-        d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), var.get_update()))
+        if c == 0: d = 0
+        else: d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), var.get_update()))
 
-        # scale beta (beta is multiplicative and i think may be better than scaling initial step size)
-        if adaptive: beta = beta * self.global_state['beta_scale']
+        # scale init
+        init_scale = self.global_state.get('init_scale', 1)
+        if adaptive: init = init * init_scale
 
-        step_size = backtracking_line_search(objective, d, init=init,beta=beta,
-                                        c=c,maxiter=maxiter, try_negative=try_negative)
+        step_size = backtracking_line_search(objective, d, init=init, beta=beta,c=c, condition=condition, maxiter=maxiter)
 
         # found an alpha that reduces loss
         if step_size is not None:
-            self.global_state['beta_scale'] = min(1.0, self.global_state['beta_scale'] * math.sqrt(1.5))
+            #self.global_state['beta_scale'] = min(1.0, self.global_state['beta_scale'] * math.sqrt(1.5))
+            self.global_state['init_scale'] = 1
             return step_size
 
-        # on fail reduce beta scale value
-        self.global_state['beta_scale'] /= 1.5
+        # on fail set init_scale to continue decreasing the step size
+        # or set to large step size when it becomes too small
+        if adaptive:
+            finfo = torch.finfo(var.params[0].dtype)
+            if init_scale <= finfo.tiny * 2:
+                self.global_state["init_scale"] = finfo.max / 2
+            else:
+                self.global_state['init_scale'] = init_scale * beta**maxiter
         return 0
 
 def _lerp(start,end,weight):
@@ -147,30 +149,37 @@ class AdaptiveBacktracking(LineSearchBase):
     such that optimal step size in the procedure would be found on the second line search iteration.
 
     Args:
-        init (float, optional): step size for the first step. Defaults to 1.0.
+        init (float, optional): initial step size. Defaults to 1.0.
         beta (float, optional): multiplies each consecutive step size by this value. Defaults to 0.5.
-        c (float, optional): acceptance value for Armijo condition. Defaults to 1e-4.
-        maxiter (int, optional): Maximum line search function evaluations. Defaults to 10.
+        c (float, optional): sufficient decrease condition. Defaults to 1e-4.
+        condition (TerminationCondition, optional):
+            termination condition, only ones that do not use gradient at f(x+a*d) can be specified.
+            - "armijo" - sufficient decrease condition.
+            - "decrease" - any decrease in objective function value satisfies the condition.
+
+            "goldstein" can techincally be specified but it doesn't make sense because there is not zoom stage.
+            Defaults to 'armijo'.
+        maxiter (int, optional): maximum number of function evaluations per step. Defaults to 10.
         target_iters (int, optional):
-            target number of iterations that would be performed until optimal step size is found. Defaults to 1.
+            sets next step size such that this number of iterations are expected
+            to be performed until optimal step size is found. Defaults to 1.
         nplus (float, optional):
-            Multiplier to initial step size if it was found to be the optimal step size. Defaults to 2.0.
+            if initial step size is optimal, it is multiplied by this value. Defaults to 2.0.
         scale_beta (float, optional):
-            Momentum for initial step size, at 0 disables momentum. Defaults to 0.0.
-        try_negative (bool, optional): Whether to perform line search in opposite direction on fail. Defaults to False.
+            momentum for initial step size, at 0 disables momentum. Defaults to 0.0.
     """
     def __init__(
         self,
         init: float = 1.0,
         beta: float = 0.5,
         c: float = 1e-4,
+        condition: TerminationCondition = 'armijo',
         maxiter: int = 20,
         target_iters = 1,
         nplus = 2.0,
         scale_beta = 0.0,
-        try_negative: bool = False,
     ):
-        defaults=dict(init=init,beta=beta,c=c,maxiter=maxiter,target_iters=target_iters,nplus=nplus,scale_beta=scale_beta, try_negative=try_negative)
+        defaults=dict(init=init,beta=beta,c=c,condition=condition,maxiter=maxiter,target_iters=target_iters,nplus=nplus,scale_beta=scale_beta)
         super().__init__(defaults=defaults)
 
         self.global_state['beta_scale'] = 1.0
@@ -183,8 +192,8 @@ class AdaptiveBacktracking(LineSearchBase):
 
     @torch.no_grad
     def search(self, update, var):
-        init, beta, c, maxiter, target_iters, nplus, scale_beta, try_negative=itemgetter(
-            'init','beta','c','maxiter','target_iters','nplus','scale_beta', 'try_negative')(self.settings[var.params[0]])
+        init, beta, c,condition, maxiter, target_iters, nplus, scale_beta=itemgetter(
+            'init','beta','c','condition', 'maxiter','target_iters','nplus','scale_beta')(self.defaults)
 
         objective = self.make_objective(var=var)
 
@@ -198,8 +207,7 @@ class AdaptiveBacktracking(LineSearchBase):
         # scale step size so that decrease is expected at target_iters
         init = init * self.global_state['initial_scale']
 
-        step_size = backtracking_line_search(objective, d, init=init, beta=beta,
-                                        c=c,maxiter=maxiter, try_negative=try_negative)
+        step_size = backtracking_line_search(objective, d, init=init, beta=beta, c=c, condition=condition, maxiter=maxiter)
 
         # found an alpha that reduces loss
         if step_size is not None:
@@ -208,7 +216,12 @@ class AdaptiveBacktracking(LineSearchBase):
             # initial step size satisfied conditions, increase initial_scale by nplus
             if step_size == init and target_iters > 0:
                 self.global_state['initial_scale'] *= nplus ** target_iters
-                self.global_state['initial_scale'] = min(self.global_state['initial_scale'], 1e32) # avoid overflow error
+
+                # clip by maximum possibel value to avoid overflow exception
+                self.global_state['initial_scale'] = min(
+                    self.global_state['initial_scale'],
+                    torch.finfo(var.params[0].dtype).max / 2,
+                )
 
             else:
                 # otherwise make initial_scale such that target_iters iterations will satisfy armijo