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

torchzero/modules/step_size/adaptive.py

@@ -1,122 +1,387 @@
 """Various step size strategies"""
-from typing import Any, Literal
+import math
 from operator import itemgetter
+from typing import Any, Literal
+
 import torch
 
-from ...core import Transform, Chainable
-from ...utils import TensorList, unpack_dicts, unpack_states, NumberList
+from ...core import Chainable, Transform
+from ...utils import NumberList, TensorList, tofloat, unpack_dicts, unpack_states
+from ...utils.linalg.linear_operator import ScaledIdentity
+from ..functional import epsilon_step_size
+
+def _acceptable_alpha(alpha, param:torch.Tensor):
+    finfo = torch.finfo(param.dtype)
+    if (alpha is None) or (alpha < finfo.tiny*2) or (not math.isfinite(alpha)) or (alpha > finfo.max/2):
+        return False
+    return True
+
+def _get_H(self: Transform, var):
+    n = sum(p.numel() for p in var.params)
+    p = var.params[0]
+    alpha = self.global_state.get('alpha', 1)
+    if not _acceptable_alpha(alpha, p): alpha = 1
+
+    return ScaledIdentity(1 / alpha, shape=(n,n), device=p.device, dtype=p.dtype)
 
 
 class PolyakStepSize(Transform):
-    """Polyak's subgradient method.
+    """Polyak's subgradient method with known or unknown f*.
 
     Args:
-        f_star (int, optional):
-            (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
+        f_star (float | Mone, optional):
+            minimal possible value of the objective function. If not known, set to ``None``. Defaults to 0.
+        y (float, optional):
+            when ``f_star`` is set to None, it is calculated as ``f_best - y``.
+        y_decay (float, optional):
+            ``y`` is multiplied by ``(1 - y_decay)`` after each step. Defaults to 1e-3.
         max (float | None, optional): maximum possible step size. Defaults to None.
         use_grad (bool, optional):
             if True, uses dot product of update and gradient to compute the step size.
-            Otherwise, dot product of update with itself is used, which has no geometric meaning so it probably won't work well.
-            Defaults to False.
+            Otherwise, dot product of update with itself is used.
         alpha (float, optional): multiplier to Polyak step-size. Defaults to 1.
     """
-    def __init__(self, f_star: float = 0, max: float | None = None, use_grad=False, alpha: float = 1, inner: Chainable | None = None):
+    def __init__(self, f_star: float | None = 0, y: float = 1, y_decay: float = 1e-3, max: float | None = None, use_grad=True, alpha: float = 1, inner: Chainable | None = None):
 
-        defaults = dict(alpha=alpha, max=max, f_star=f_star, use_grad=use_grad)
+        defaults = dict(alpha=alpha, max=max, f_star=f_star, y=y, y_decay=y_decay)
         super().__init__(defaults, uses_grad=use_grad, uses_loss=True, inner=inner)
 
+    @torch.no_grad
     def update_tensors(self, tensors, params, grads, loss, states, settings):
         assert grads is not None and loss is not None
         tensors = TensorList(tensors)
         grads = TensorList(grads)
 
-        use_grad, max, f_star = itemgetter('use_grad', 'max', 'f_star')(settings[0])
+        # load variables
+        max, f_star, y, y_decay = itemgetter('max', 'f_star', 'y', 'y_decay')(settings[0])
+        y_val = self.global_state.get('y_val', y)
+        f_best = self.global_state.get('f_best', None)
 
-        if use_grad: gg = tensors.dot(grads)
+        # gg
+        if self._uses_grad: gg = tensors.dot(grads)
         else: gg = tensors.dot(tensors)
 
-        if gg.abs() <= torch.finfo(gg.dtype).eps: step_size = 0 # converged
-        else: step_size = (loss - f_star) / gg
+        # store loss
+        if f_best is None or loss < f_best: f_best = tofloat(loss)
+        if f_star is None: f_star = f_best - y_val
+
+        # calculate the step size
+        if gg <= torch.finfo(gg.dtype).tiny * 2: alpha = 0 # converged
+        else: alpha = (loss - f_star) / gg
 
+        # clip
         if max is not None:
-            if step_size > max: step_size = max
+            if alpha > max: alpha = max
 
-        self.global_state['step_size'] = step_size
+        # store state
+        self.global_state['f_best'] = f_best
+        self.global_state['y_val'] = y_val * (1 - y_decay)
+        self.global_state['alpha'] = alpha
 
     @torch.no_grad
     def apply_tensors(self, tensors, params, grads, loss, states, settings):
-        step_size = self.global_state.get('step_size', 1)
-        torch._foreach_mul_(tensors, step_size * unpack_dicts(settings, 'alpha', cls=NumberList))
+        alpha = self.global_state.get('alpha', 1)
+        if not _acceptable_alpha(alpha, tensors[0]): alpha = epsilon_step_size(TensorList(tensors))
+
+        torch._foreach_mul_(tensors, alpha * unpack_dicts(settings, 'alpha', cls=NumberList))
         return tensors
 
+    def get_H(self, var):
+        return _get_H(self, var)
 
 
-def _bb_short(s: TensorList, y: TensorList, sy, eps, fallback):
+def _bb_short(s: TensorList, y: TensorList, sy, eps):
     yy = y.dot(y)
     if yy < eps:
-        if sy < eps: return fallback # try to fallback on long
+        if sy < eps: return None # try to fallback on long
         ss = s.dot(s)
         return ss/sy
     return sy/yy
 
-def _bb_long(s: TensorList, y: TensorList, sy, eps, fallback):
+def _bb_long(s: TensorList, y: TensorList, sy, eps):
     ss = s.dot(s)
     if sy < eps:
         yy = y.dot(y) # try to fallback on short
-        if yy < eps: return fallback
+        if yy < eps: return None
         return sy/yy
     return ss/sy
 
-def _bb_geom(s: TensorList, y: TensorList, sy, eps, fallback):
-    short = _bb_short(s, y, sy, eps, fallback)
-    long = _bb_long(s, y, sy, eps, fallback)
+def _bb_geom(s: TensorList, y: TensorList, sy, eps, fallback:bool):
+    short = _bb_short(s, y, sy, eps)
+    long = _bb_long(s, y, sy, eps)
+    if long is None or short is None:
+        if fallback:
+            if short is not None: return short
+            if long is not None: return long
+        return None
     return (short * long) ** 0.5
 
 class BarzilaiBorwein(Transform):
-    """Barzilai-Borwein method.
+    """Barzilai-Borwein step size method.
 
     Args:
         type (str, optional):
             one of "short" with formula sᵀy/yᵀy, "long" with formula sᵀs/sᵀy, or "geom" to use geometric mean of short and long.
-            Defaults to 'geom'.
-        scale_first (bool, optional):
-            whether to make first step very small when previous gradient is not available. Defaults to True.
+            Defaults to "geom".
         fallback (float, optional): step size when denominator is less than 0 (will happen on negative curvature). Defaults to 1e-3.
         inner (Chainable | None, optional):
             step size will be applied to outputs of this module. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        type: Literal["long", "short", "geom", "geom-fallback"] = "geom",
+        alpha_0: float = 1e-7,
+        use_grad=True,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(type=type, alpha_0=alpha_0)
+        super().__init__(defaults, uses_grad=use_grad, inner=inner)
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_g')
+        self.global_state['reset'] = True
+
+    @torch.no_grad
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        prev_p, prev_g = unpack_states(states, tensors, 'prev_p', 'prev_g', cls=TensorList)
+        type = self.defaults['type']
+
+        g = grads if self._uses_grad else tensors
+        assert g is not None
+
+        reset = self.global_state.get('reset', False)
+        self.global_state.pop('reset', None)
+
+        if step != 0 and not reset:
+            s = params-prev_p
+            y = g-prev_g
+            sy = s.dot(y)
+            eps = torch.finfo(sy.dtype).tiny * 2
+
+            if type == 'short': alpha = _bb_short(s, y, sy, eps)
+            elif type == 'long': alpha = _bb_long(s, y, sy, eps)
+            elif type == 'geom': alpha = _bb_geom(s, y, sy, eps, fallback=False)
+            elif type == 'geom-fallback': alpha = _bb_geom(s, y, sy, eps, fallback=True)
+            else: raise ValueError(type)
+
+            # if alpha is not None:
+            self.global_state['alpha'] = alpha
+
+        prev_p.copy_(params)
+        prev_g.copy_(g)
+
+    def get_H(self, var):
+        return _get_H(self, var)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        alpha = self.global_state.get('alpha', None)
+
+        if not _acceptable_alpha(alpha, tensors[0]):
+            alpha = epsilon_step_size(TensorList(tensors), settings[0]['alpha_0'])
+
+        torch._foreach_mul_(tensors, alpha)
+        return tensors
+
+
+class BBStab(Transform):
+    """Stabilized Barzilai-Borwein method (https://arxiv.org/abs/1907.06409).
+
+    This clips the norm of the Barzilai-Borwein update by ``delta``, where ``delta`` can be adaptive if ``c`` is specified.
+
+    Args:
+        c (float, optional):
+            adaptive delta parameter. If ``delta`` is set to None, first ``inf_iters`` updates are performed
+            with non-stabilized Barzilai-Borwein step size. Then delta is set to norm of
+            the update that had the smallest norm, and multiplied by ``c``. Defaults to 0.2.
+        delta (float | None, optional):
+            Barzilai-Borwein update is clipped to this value. Set to ``None`` to use an adaptive choice. Defaults to None.
+        type (str, optional):
+            one of "short" with formula sᵀy/yᵀy, "long" with formula sᵀs/sᵀy, or "geom" to use geometric mean of short and long.
+            Defaults to "geom". Note that "long" corresponds to BB1stab and "short" to BB2stab,
+            however I found that "geom" works really well.
+        inner (Chainable | None, optional):
+            step size will be applied to outputs of this module. Defaults to None.
 
     """
-    def __init__(self, type: Literal['long', 'short', 'geom'] = 'geom', scale_first:bool=True, fallback:float=1e-3, inner:Chainable|None = None):
-        defaults = dict(type=type, fallback=fallback)
-        super().__init__(defaults, uses_grad=False, scale_first=scale_first, inner=inner)
+    def __init__(
+        self,
+        c=0.2,
+        delta:float | None = None,
+        type: Literal["long", "short", "geom", "geom-fallback"] = "geom",
+        alpha_0: float = 1e-7,
+        use_grad=True,
+        inf_iters: int = 3,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(type=type,alpha_0=alpha_0, c=c, delta=delta, inf_iters=inf_iters)
+        super().__init__(defaults, uses_grad=use_grad, inner=inner)
 
     def reset_for_online(self):
         super().reset_for_online()
-        self.clear_state_keys('prev_p', 'prev_g')
+        self.clear_state_keys('prev_g')
+        self.global_state['reset'] = True
 
     @torch.no_grad
     def update_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
         prev_p, prev_g = unpack_states(states, tensors, 'prev_p', 'prev_g', cls=TensorList)
-        fallback = unpack_dicts(settings, 'fallback', cls=NumberList)
-        type = settings[0]['type']
+        type = self.defaults['type']
+        c = self.defaults['c']
+        delta = self.defaults['delta']
+        inf_iters = self.defaults['inf_iters']
+
+        g = grads if self._uses_grad else tensors
+        assert g is not None
+        g = TensorList(g)
+
+        reset = self.global_state.get('reset', False)
+        self.global_state.pop('reset', None)
+
+        if step != 0 and not reset:
+            s = params-prev_p
+            y = g-prev_g
+            sy = s.dot(y)
+            eps = torch.finfo(sy.dtype).tiny
+
+            if type == 'short': alpha = _bb_short(s, y, sy, eps)
+            elif type == 'long': alpha = _bb_long(s, y, sy, eps)
+            elif type == 'geom': alpha = _bb_geom(s, y, sy, eps, fallback=False)
+            elif type == 'geom-fallback': alpha = _bb_geom(s, y, sy, eps, fallback=True)
+            else: raise ValueError(type)
+
+            if alpha is not None:
+
+                # adaptive delta
+                if delta is None:
+                    niters = self.global_state.get('niters', 0) # this accounts for skipped negative curvature steps
+                    self.global_state['niters'] = niters + 1
+
+
+                    if niters == 0: pass # 1st iteration is scaled GD step, shouldn't be used to find s_norm_min
+                    elif niters <= inf_iters:
+                        s_norm_min = self.global_state.get('s_norm_min', None)
+                        if s_norm_min is None: s_norm_min = s.global_vector_norm()
+                        else: s_norm_min = min(s_norm_min, s.global_vector_norm())
+                        self.global_state['s_norm_min'] = s_norm_min
+                        # first few steps use delta=inf, so delta remains None
 
-        s = params-prev_p
-        y = tensors-prev_g
-        sy = s.dot(y)
-        eps = torch.finfo(sy.dtype).eps
+                    else:
+                        delta = c * self.global_state['s_norm_min']
 
-        if type == 'short': step_size = _bb_short(s, y, sy, eps, fallback)
-        elif type == 'long': step_size = _bb_long(s, y, sy, eps, fallback)
-        elif type == 'geom': step_size = _bb_geom(s, y, sy, eps, fallback)
-        else: raise ValueError(type)
+                if delta is None: # delta is inf for first few steps
+                    self.global_state['alpha'] = alpha
 
-        self.global_state['step_size'] = step_size
+                # BBStab step size
+                else:
+                    a_stab = delta / g.global_vector_norm()
+                    self.global_state['alpha'] = min(alpha, a_stab)
 
         prev_p.copy_(params)
-        prev_g.copy_(tensors)
+        prev_g.copy_(g)
+
+    def get_H(self, var):
+        return _get_H(self, var)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        alpha = self.global_state.get('alpha', None)
+
+        if not _acceptable_alpha(alpha, tensors[0]):
+            alpha = epsilon_step_size(TensorList(tensors), settings[0]['alpha_0'])
+
+        torch._foreach_mul_(tensors, alpha)
+        return tensors
+
+
+class AdGD(Transform):
+    """AdGD and AdGD-2 (https://arxiv.org/abs/2308.02261)"""
+    def __init__(self, variant:Literal[1,2]=2, alpha_0:float = 1e-7, sqrt:bool=True, use_grad=True, inner: Chainable | None = None,):
+        defaults = dict(variant=variant, alpha_0=alpha_0, sqrt=sqrt)
+        super().__init__(defaults, uses_grad=use_grad, inner=inner,)
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_g')
+        self.global_state['reset'] = True
+
+    @torch.no_grad
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        variant = settings[0]['variant']
+        theta_0 = 0 if variant == 1 else 1/3
+        theta = self.global_state.get('theta', theta_0)
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        p = TensorList(params)
+        g = grads if self._uses_grad else tensors
+        assert g is not None
+        g = TensorList(g)
+
+        prev_p, prev_g = unpack_states(states, tensors, 'prev_p', 'prev_g', cls=TensorList)
+
+        # online
+        if self.global_state.get('reset', False):
+            del self.global_state['reset']
+            prev_p.copy_(p)
+            prev_g.copy_(g)
+            return
 
+        if step == 0:
+            alpha_0 = settings[0]['alpha_0']
+            if alpha_0 is None: alpha_0 = epsilon_step_size(g)
+            self.global_state['alpha']  = alpha_0
+            prev_p.copy_(p)
+            prev_g.copy_(g)
+            return
+
+        sqrt = settings[0]['sqrt']
+        alpha = self.global_state.get('alpha', math.inf)
+        L = (g - prev_g).global_vector_norm() / (p - prev_p).global_vector_norm()
+        eps = torch.finfo(L.dtype).tiny * 2
+
+        if variant == 1:
+            a1 = math.sqrt(1 + theta)*alpha
+            val = math.sqrt(2) if sqrt else 2
+            if L > eps: a2 = 1 / (val*L)
+            else: a2 = math.inf
+
+        elif variant == 2:
+            a1 = math.sqrt(2/3 + theta)*alpha
+            a2 = alpha / math.sqrt(max(eps, 2 * alpha**2 * L**2 - 1))
+
+        else:
+            raise ValueError(variant)
+
+        alpha_new = min(a1, a2)
+        if alpha_new < 0: alpha_new = max(a1, a2)
+        if alpha_new > eps:
+            self.global_state['theta'] = alpha_new/alpha
+            self.global_state['alpha'] = alpha_new
+
+        prev_p.copy_(p)
+        prev_g.copy_(g)
+
+    @torch.no_grad
     def apply_tensors(self, tensors, params, grads, loss, states, settings):
-        step_size = self.global_state.get('step_size', 1)
-        torch._foreach_mul_(tensors, step_size)
+        alpha = self.global_state.get('alpha', None)
+
+        if not _acceptable_alpha(alpha, tensors[0]):
+            # alpha isn't None on 1st step
+            self.state.clear()
+            self.global_state.clear()
+            alpha = epsilon_step_size(TensorList(tensors), settings[0]['alpha_0'])
+
+        torch._foreach_mul_(tensors, alpha)
         return tensors
 
+    def get_H(self, var):
+        return _get_H(self, var)

torchzero/modules/termination/__init__.py

@@ -0,0 +1,14 @@
+from .termination import (
+    TerminateAfterNEvaluations,
+    TerminateAfterNSeconds,
+    TerminateAfterNSteps,
+    TerminateAll,
+    TerminateAny,
+    TerminateByGradientNorm,
+    TerminateByUpdateNorm,
+    TerminateOnLossReached,
+    TerminateOnNoImprovement,
+    TerminationCriteriaBase,
+    TerminateNever,
+    make_termination_criteria
+)

torchzero/modules/termination/termination.py

@@ -0,0 +1,207 @@
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import cast
+
+import torch
+
+from ...core import Module, Var
+from ...utils import Metrics, TensorList, safe_dict_update_, tofloat
+
+
+class TerminationCriteriaBase(Module):
+    def __init__(self, defaults:dict | None = None, n: int = 1):
+        if defaults is None: defaults = {}
+        safe_dict_update_(defaults, {"_n": n})
+        super().__init__(defaults)
+
+    @abstractmethod
+    def termination_criteria(self, var: Var) -> bool:
+        ...
+
+    def should_terminate(self, var: Var) -> bool:
+        n_bad = self.global_state.get('_n_bad', 0)
+        n = self.defaults['_n']
+
+        if self.termination_criteria(var):
+            n_bad += 1
+            if n_bad >= n:
+                self.global_state['_n_bad'] = 0
+                return True
+
+        else:
+            n_bad = 0
+
+        self.global_state['_n_bad'] = n_bad
+        return False
+
+
+    def update(self, var):
+        var.should_terminate = self.should_terminate(var)
+        if var.should_terminate: self.global_state['_n_bad'] = 0
+
+    def apply(self, var):
+        return var
+
+
+class TerminateAfterNSteps(TerminationCriteriaBase):
+    def __init__(self, steps:int):
+        defaults = dict(steps=steps)
+        super().__init__(defaults)
+
+    def termination_criteria(self, var):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        max_steps = self.defaults['steps']
+        return step >= max_steps
+
+class TerminateAfterNEvaluations(TerminationCriteriaBase):
+    def __init__(self, maxevals:int):
+        defaults = dict(maxevals=maxevals)
+        super().__init__(defaults)
+
+    def termination_criteria(self, var):
+        maxevals = self.defaults['maxevals']
+        return var.modular.num_evaluations >= maxevals
+
+class TerminateAfterNSeconds(TerminationCriteriaBase):
+    def __init__(self, seconds:float, sec_fn = time.time):
+        defaults = dict(seconds=seconds, sec_fn=sec_fn)
+        super().__init__(defaults)
+
+    def termination_criteria(self, var):
+        max_seconds = self.defaults['seconds']
+        sec_fn = self.defaults['sec_fn']
+
+        if 'start' not in self.global_state:
+            self.global_state['start'] = sec_fn()
+            return False
+
+        seconds_passed = sec_fn() - self.global_state['start']
+        return seconds_passed >= max_seconds
+
+
+
+class TerminateByGradientNorm(TerminationCriteriaBase):
+    def __init__(self, tol:float = 1e-8, n: int = 3, ord: Metrics = 2):
+        defaults = dict(tol=tol, ord=ord)
+        super().__init__(defaults, n=n)
+
+    def termination_criteria(self, var):
+        tol = self.defaults['tol']
+        ord = self.defaults['ord']
+        return TensorList(var.get_grad()).global_metric(ord) <= tol
+
+
+class TerminateByUpdateNorm(TerminationCriteriaBase):
+    """update is calculated as parameter difference"""
+    def __init__(self, tol:float = 1e-8, n: int = 3, ord: Metrics = 2):
+        defaults = dict(tol=tol, ord=ord)
+        super().__init__(defaults, n=n)
+
+    def termination_criteria(self, var):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        tol = self.defaults['tol']
+        ord = self.defaults['ord']
+
+        p_prev = self.get_state(var.params, 'p_prev', cls=TensorList)
+        if step == 0:
+            p_prev.copy_(var.params)
+            return False
+
+        should_terminate = (p_prev - var.params).global_metric(ord) <= tol
+        p_prev.copy_(var.params)
+        return should_terminate
+
+
+class TerminateOnNoImprovement(TerminationCriteriaBase):
+    def __init__(self, tol:float = 1e-8, n: int = 10):
+        defaults = dict(tol=tol)
+        super().__init__(defaults, n=n)
+
+    def termination_criteria(self, var):
+        tol = self.defaults['tol']
+
+        f = tofloat(var.get_loss(False))
+        if 'f_min' not in self.global_state:
+            self.global_state['f_min'] = f
+            return False
+
+        f_min = self.global_state['f_min']
+        d = f_min - f
+        should_terminate = d <= tol
+        self.global_state['f_min'] = min(f, f_min)
+        return should_terminate
+
+class TerminateOnLossReached(TerminationCriteriaBase):
+    def __init__(self, value: float):
+        defaults = dict(value=value)
+        super().__init__(defaults)
+
+    def termination_criteria(self, var):
+        value = self.defaults['value']
+        return var.get_loss(False) <= value
+
+class TerminateAny(TerminationCriteriaBase):
+    def __init__(self, *criteria: TerminationCriteriaBase):
+        super().__init__()
+
+        self.set_children_sequence(criteria)
+
+    def termination_criteria(self, var: Var) -> bool:
+        for c in self.get_children_sequence():
+            if cast(TerminationCriteriaBase, c).termination_criteria(var): return True
+
+        return False
+
+class TerminateAll(TerminationCriteriaBase):
+    def __init__(self, *criteria: TerminationCriteriaBase):
+        super().__init__()
+
+        self.set_children_sequence(criteria)
+
+    def termination_criteria(self, var: Var) -> bool:
+        for c in self.get_children_sequence():
+            if not cast(TerminationCriteriaBase, c).termination_criteria(var): return False
+
+        return True
+
+class TerminateNever(TerminationCriteriaBase):
+    def __init__(self):
+        super().__init__()
+
+    def termination_criteria(self, var): return False
+
+def make_termination_criteria(
+    ftol: float | None = None,
+    gtol: float | None = None,
+    stol: float | None = None,
+    maxiter: int | None = None,
+    maxeval: int | None = None,
+    maxsec: float | None = None,
+    target_loss: float | None = None,
+    extra: TerminationCriteriaBase | Sequence[TerminationCriteriaBase] | None = None,
+    n: int = 3,
+):
+    criteria: list[TerminationCriteriaBase] = []
+
+    if ftol is not None: criteria.append(TerminateOnNoImprovement(ftol, n=n))
+    if gtol is not None: criteria.append(TerminateByGradientNorm(gtol, n=n))
+    if stol is not None: criteria.append(TerminateByUpdateNorm(stol, n=n))
+
+    if maxiter is not None: criteria.append(TerminateAfterNSteps(maxiter))
+    if maxeval is not None: criteria.append(TerminateAfterNEvaluations(maxeval))
+    if maxsec is not None: criteria.append(TerminateAfterNSeconds(maxsec))
+
+    if target_loss is not None: criteria.append(TerminateOnLossReached(target_loss))
+
+    if extra is not None:
+        if isinstance(extra, TerminationCriteriaBase): criteria.append(extra)
+        else: criteria.extend(extra)
+
+    if len(criteria) == 0: return TerminateNever()
+    if len(criteria) == 1: return criteria[0]
+    return TerminateAny(*criteria)

torchzero/modules/trust_region/__init__.py

@@ -0,0 +1,5 @@
+from .trust_region import TrustRegionBase
+from .cubic_regularization import CubicRegularization
+from .trust_cg import TrustCG
+from .levenberg_marquardt import LevenbergMarquardt
+from .dogleg import Dogleg