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

torchzero/modules/second_order/newton_cg.py

@@ -13,21 +13,12 @@ from ..trust_region.trust_region import default_radius
 class NewtonCG(Module):
     """Newton's method with a matrix-free conjugate gradient or minimial-residual solver.
 
-    This optimizer implements Newton's method using a matrix-free conjugate
-    gradient (CG) or a minimal-residual (MINRES) solver to approximate the search direction. Instead of
-    forming the full Hessian matrix, it only requires Hessian-vector products
-    (HVPs). These can be calculated efficiently using automatic
-    differentiation or approximated using finite differences.
+    Notes:
+        * In most cases NewtonCGSteihaug should be the first module in the chain because it relies on autograd. Use the ``inner`` argument if you wish to apply Newton preconditioning to another module's output.
 
-    .. note::
-        In most cases NewtonCG should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Newton preconditioning to another module's output.
+        * This module requires the a closure passed to the optimizer step, as it needs to re-evaluate the loss and gradients for calculating HVPs. The closure must accept a ``backward`` argument (refer to documentation).
 
-    .. note::
-        This module requires the a closure passed to the optimizer step,
-        as it needs to re-evaluate the loss and gradients for calculating HVPs.
-        The closure must accept a ``backward`` argument (refer to documentation).
-
-    .. warning::
+    Warning:
         CG may fail if hessian is not positive-definite.
 
     Args:
@@ -66,26 +57,24 @@ class NewtonCG(Module):
             NewtonCG will attempt to apply preconditioning to the output of this module.
 
     Examples:
-        Newton-CG with a backtracking line search:
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NewtonCG(),
-                tz.m.Backtracking()
-            )
-
-        Truncated Newton method (useful for large-scale problems):
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NewtonCG(maxiter=10, warm_start=True),
-                tz.m.Backtracking()
-            )
-
+    Newton-CG with a backtracking line search:
+
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NewtonCG(),
+        tz.m.Backtracking()
+    )
+    ```
+
+    Truncated Newton method (useful for large-scale problems):
+    ```
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NewtonCG(maxiter=10),
+        tz.m.Backtracking()
+    )
+    ```
 
     """
     def __init__(
@@ -95,7 +84,7 @@ class NewtonCG(Module):
         reg: float = 1e-8,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
         solver: Literal['cg', 'minres', 'minres_npc'] = 'cg',
-        h: float = 1e-3,
+        h: float = 1e-3, # tuned 1e-4 or 1e-3
         miniter:int = 1,
         warm_start=False,
         inner: Chainable | None = None,
@@ -187,96 +176,95 @@ class NewtonCG(Module):
 
 
 class NewtonCGSteihaug(Module):
-    """Trust region Newton's method with a matrix-free Steihaug-Toint conjugate gradient or MINRES solver.
+    """Newton's method with trust region and a matrix-free Steihaug-Toint conjugate gradient solver.
 
-    This optimizer implements Newton's method using a matrix-free conjugate
-    gradient (CG) solver to approximate the search direction. Instead of
-    forming the full Hessian matrix, it only requires Hessian-vector products
-    (HVPs). These can be calculated efficiently using automatic
-    differentiation or approximated using finite differences.
+    Notes:
+        * In most cases NewtonCGSteihaug should be the first module in the chain because it relies on autograd. Use the ``inner`` argument if you wish to apply Newton preconditioning to another module's output.
 
-    .. note::
-        In most cases NewtonCGSteihaug should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Newton preconditioning to another module's output.
-
-    .. note::
-        This module requires the a closure passed to the optimizer step,
-        as it needs to re-evaluate the loss and gradients for calculating HVPs.
-        The closure must accept a ``backward`` argument (refer to documentation).
-
-    .. warning::
-        CG may fail if hessian is not positive-definite.
+        * This module requires the a closure passed to the optimizer step, as it needs to re-evaluate the loss and gradients for calculating HVPs. The closure must accept a ``backward`` argument (refer to documentation).
 
     Args:
-        maxiter (int | None, optional):
-            Maximum number of iterations for the conjugate gradient solver.
-            By default, this is set to the number of dimensions in the
-            objective function, which is the theoretical upper bound for CG
-            convergence. Setting this to a smaller value (truncated Newton)
-            can still generate good search directions. Defaults to None.
         eta (float, optional):
-            whenever actual to predicted loss reduction ratio is larger than this, a step is accepted.
-        nplus (float, optional):
-            trust region multiplier on successful steps.
-        nminus (float, optional):
-            trust region multiplier on unsuccessful steps.
-        init (float, optional): initial trust region.
+            if ratio of actual to predicted rediction is larger than this, step is accepted. Defaults to 0.0.
+        nplus (float, optional): increase factor on successful steps. Defaults to 1.5.
+        nminus (float, optional): decrease factor on unsuccessful steps. Defaults to 0.75.
+        rho_good (float, optional):
+            if ratio of actual to predicted rediction is larger than this, trust region size is multiplied by `nplus`.
+        rho_bad (float, optional):
+            if ratio of actual to predicted rediction is less than this, trust region size is multiplied by `nminus`.
+        init (float, optional): Initial trust region value. Defaults to 1.
+        max_attempts (max_attempts, optional):
+            maximum number of trust radius reductions per step. A zero update vector is returned when
+            this limit is exceeded. Defaults to 10.
+        max_history (int, optional):
+            CG will store this many intermediate solutions, reusing them when trust radius is reduced
+            instead of re-running CG. Each solution storage requires 2N memory. Defaults to 100.
+        boundary_tol (float | None, optional):
+            The trust region only increases when suggested step's norm is at least `(1-boundary_tol)*trust_region`.
+            This prevents increasing trust region when solution is not on the boundary. Defaults to 1e-2.
+
+        maxiter (int | None, optional):
+            maximum number of CG iterations per step. Each iteration requies one backward pass if `hvp_method="forward"`, two otherwise. Defaults to None.
+        miniter (int, optional):
+            minimal number of CG iterations. This prevents making no progress
         tol (float, optional):
-            Relative tolerance for the conjugate gradient solver to determine
-            convergence. Defaults to 1e-4.
-        reg (float, optional):
-            Regularization parameter (damping) added to the Hessian diagonal.
-            This helps ensure the system is positive-definite. Defaults to 1e-8.
+            terminates CG when norm of the residual is less than this value. Defaults to 1e-8.
+            when initial guess is below tolerance. Defaults to 1.
+        reg (float, optional): hessian regularization. Defaults to 1e-8.
+        solver (str, optional): solver, "cg" or "minres". "cg" is recommended. Defaults to 'cg'.
+        adapt_tol (bool, optional):
+            if True, whenever trust radius collapses to smallest representable number,
+            the tolerance is multiplied by 0.1. Defaults to True.
+        npc_terminate (bool, optional):
+            whether to terminate CG/MINRES whenever negative curvature is detected. Defaults to False.
+
         hvp_method (str, optional):
-            Determines how Hessian-vector products are evaluated.
+            either "forward" to use forward formula which requires one backward pass per Hvp, or "central" to use a more accurate central formula which requires two backward passes. "forward" is usually accurate enough. Defaults to "forward".
+        h (float, optional): finite difference step size. Defaults to 1e-3.
 
-            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
-              This requires creating a graph for the gradient.
-            - ``"forward"``: Use a forward finite difference formula to
-              approximate the HVP. This requires one extra gradient evaluation.
-            - ``"central"``: Use a central finite difference formula for a
-              more accurate HVP approximation. This requires two extra
-              gradient evaluations.
-            Defaults to "autograd".
-        h (float, optional):
-            The step size for finite differences if :code:`hvp_method` is
-            ``"forward"`` or ``"central"``. Defaults to 1e-3.
         inner (Chainable | None, optional):
-            NewtonCG will attempt to apply preconditioning to the output of this module.
+            applies preconditioning to output of this module. Defaults to None.
 
-    Examples:
-        Trust-region Newton-CG:
-
-        .. code-block:: python
+    ### Examples:
+    Trust-region Newton-CG:
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NewtonCGSteihaug(),
-            )
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NewtonCGSteihaug(),
+    )
+    ```
 
-    Reference:
+    ### Reference:
         Steihaug, Trond. "The conjugate gradient method and trust regions in large scale optimization." SIAM Journal on Numerical Analysis 20.3 (1983): 626-637.
     """
     def __init__(
         self,
-        maxiter: int | None = None,
+        # trust region settings
         eta: float= 0.0,
         nplus: float = 3.5,
         nminus: float = 0.25,
         rho_good: float = 0.99,
         rho_bad: float = 1e-4,
         init: float = 1,
-        tol: float = 1e-8,
-        reg: float = 1e-8,
-        hvp_method: Literal["forward", "central"] = "forward",
-        solver: Literal['cg', "minres"] = 'cg',
-        h: float = 1e-3,
         max_attempts: int = 100,
         max_history: int = 100,
-        boundary_tol: float = 1e-1,
+        boundary_tol: float = 1e-6, # tuned
+
+        # cg settings
+        maxiter: int | None = None,
         miniter: int = 1,
-        rms_beta: float | None = None,
+        tol: float = 1e-8,
+        reg: float = 1e-8,
+        solver: Literal['cg', "minres"] = 'cg',
         adapt_tol: bool = True,
         npc_terminate: bool = False,
+
+        # hvp settings
+        hvp_method: Literal["forward", "central"] = "central",
+        h: float = 1e-3, # tuned 1e-4 or 1e-3
+
+        # inner
         inner: Chainable | None = None,
     ):
         defaults = locals().copy()
@@ -336,19 +324,8 @@ class NewtonCGSteihaug(Module):
                 raise ValueError(hvp_method)
 
 
-        # ------------------------- update RMS preconditioner ------------------------ #
-        b = var.get_update()
-        P_mm = None
-        rms_beta = self.defaults["rms_beta"]
-        if rms_beta is not None:
-            exp_avg_sq = self.get_state(params, "exp_avg_sq", init=b, cls=TensorList)
-            exp_avg_sq.mul_(rms_beta).addcmul(b, b, value=1-rms_beta)
-            exp_avg_sq_sqrt = exp_avg_sq.sqrt().add_(1e-8)
-            def _P_mm(x):
-                return x / exp_avg_sq_sqrt
-            P_mm = _P_mm
-
         # -------------------------------- inner step -------------------------------- #
+        b = var.get_update()
         if 'inner' in self.children:
             b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
         b = as_tensorlist(b)
@@ -392,7 +369,6 @@ class NewtonCGSteihaug(Module):
                         miniter=miniter,
                         npc_terminate=npc_terminate,
                         history_size=max_history,
-                        P_mm=P_mm,
                     )
 
                 elif solver == 'minres':

torchzero/modules/trust_region/levenberg_marquardt.py

@@ -14,13 +14,13 @@ class LevenbergMarquardt(TrustRegionBase):
         hess_module (Module | None, optional):
             A module that maintains a hessian approximation (not hessian inverse!).
             This includes all full-matrix quasi-newton methods, ``tz.m.Newton`` and ``tz.m.GaussNewton``.
-            When using quasi-newton methods, set `inverse=False` when constructing them.
+            When using quasi-newton methods, set ``inverse=False`` when constructing them.
         y (float, optional):
             when ``y=0``, identity matrix is added to hessian, when ``y=1``, diagonal of the hessian approximation
             is added. Values between interpolate. This should only be used with Gauss-Newton. Defaults to 0.
         eta (float, optional):
             if ratio of actual to predicted rediction is larger than this, step is accepted.
-            When :code:`hess_module` is GaussNewton, this can be set to 0. Defaults to 0.15.
+            When ``hess_module`` is ``Newton`` or ``GaussNewton``, this can be set to 0. Defaults to 0.15.
         nplus (float, optional): increase factor on successful steps. Defaults to 1.5.
         nminus (float, optional): decrease factor on unsuccessful steps. Defaults to 0.75.
         rho_good (float, optional):

torchzero/modules/trust_region/trust_cg.py

@@ -60,17 +60,19 @@ class TrustCG(TrustRegionBase):
         nminus: float = 0.25,
         rho_good: float = 0.99,
         rho_bad: float = 1e-4,
-        boundary_tol: float | None = 1e-1,
+        boundary_tol: float | None = 1e-6, # tuned
         init: float = 1,
         max_attempts: int = 10,
         radius_strategy: _RadiusStrategy | _RADIUS_KEYS = 'default',
         reg: float = 0,
-        cg_tol: float = 1e-4,
+        maxiter: int | None = None,
+        miniter: int = 1,
+        cg_tol: float = 1e-8,
         prefer_exact: bool = True,
         update_freq: int = 1,
         inner: Chainable | None = None,
     ):
-        defaults = dict(reg=reg, prefer_exact=prefer_exact, cg_tol=cg_tol)
+        defaults = dict(reg=reg, prefer_exact=prefer_exact, cg_tol=cg_tol, maxiter=maxiter, miniter=miniter)
         super().__init__(
             defaults=defaults,
             hess_module=hess_module,
@@ -93,5 +95,5 @@ class TrustCG(TrustRegionBase):
         if settings['prefer_exact'] and isinstance(H, linear_operator.ScaledIdentity):
             return H.solve_bounded(g, radius)
 
-        x, _ = cg(H.matvec, g, trust_radius=radius, reg=settings['reg'], tol=settings["cg_tol"])
+        x, _ = cg(H.matvec, g, trust_radius=radius, reg=settings['reg'], maxiter=settings["maxiter"], miniter=settings["miniter"], tol=settings["cg_tol"])
         return x

torchzero/modules/zeroth_order/__init__.py

@@ -1 +1 @@
-from .cd import CD, CCD, CCDLS
+from .cd import CD

torchzero/modules/zeroth_order/cd.py

@@ -9,11 +9,10 @@ import torch
 
 from ...core import Module
 from ...utils import NumberList, TensorList
-from ..line_search.adaptive import adaptive_tracking
 
 class CD(Module):
     """Coordinate descent. Proposes a descent direction along a single coordinate.
-    You can then put a line search such as ``tz.m.ScipyMinimizeScalar``, or a fixed step size.
+    A line search such as ``tz.m.ScipyMinimizeScalar(maxiter=8)`` or a fixed step size can be used after this.
 
     Args:
         h (float, optional): finite difference step size. Defaults to 1e-3.
@@ -121,239 +120,3 @@ class CD(Module):
         var.update = update
         return var
 
-
-def _icd_get_idx(self: Module, params: TensorList):
-    ndim = params.global_numel()
-    igrad = self.get_state(params, "igrad", cls=TensorList)
-
-    # -------------------------- 1st n steps fill igrad -------------------------- #
-    index = self.global_state.get('index', 0)
-    self.global_state['index'] = index + 1
-    if index < ndim:
-        return index, igrad
-
-    # ------------------ select randomly weighted by magnitudes ------------------ #
-    igrad_abs = igrad.abs()
-    gmin = igrad_abs.global_min()
-    gmax = igrad_abs.global_max()
-
-    pmin, pmax, pow = self.get_settings(params, "pmin", "pmax", "pow", cls=NumberList)
-
-    p: TensorList = ((igrad_abs - gmin) / (gmax - gmin)) ** pow # pyright:ignore[reportOperatorIssue]
-    p.mul_(pmax-pmin).add_(pmin)
-
-    if 'np_gen' not in self.global_state:
-        self.global_state['np_gen'] = np.random.default_rng(0)
-    np_gen = self.global_state['np_gen']
-
-    p_vec = p.to_vec()
-    p_sum = p_vec.sum()
-    if p_sum > 1e-12:
-        return np_gen.choice(ndim, p=p_vec.div_(p_sum).numpy(force=True)), igrad
-
-    # --------------------- sum is too small, do cycle again --------------------- #
-    self.global_state.clear()
-    self.clear_state_keys('h_vec', 'igrad', 'alphas')
-
-    if 'generator' not in self.global_state:
-        self.global_state['generator'] = random.Random(0)
-    generator = self.global_state['generator']
-    return generator.randrange(0, p_vec.numel()), igrad
-
-class CCD(Module):
-    """Cumulative coordinate descent. This updates one gradient coordinate at a time and accumulates it
-    to the update direction. The coordinate updated is random weighted by magnitudes of current update direction.
-    As update direction ceases to be a descent direction due to old accumulated coordinates, it is decayed.
-
-    Args:
-        pmin (float, optional): multiplier to probability of picking the lowest magnitude gradient. Defaults to 0.1.
-        pmax (float, optional): multiplier to probability of picking the largest magnitude gradient. Defaults to 1.0.
-        pow (int, optional): power transform to probabilities. Defaults to 2.
-        decay (float, optional): accumulated gradient decay on failed step. Defaults to 0.5.
-        decay2 (float, optional): decay multiplier decay on failed step. Defaults to 0.25.
-        nplus (float, optional): step size increase on successful steps. Defaults to 1.5.
-        nminus (float, optional): step size increase on unsuccessful steps. Defaults to 0.75.
-    """
-    def __init__(self, pmin=0.1, pmax=1.0, pow=2, decay:float=0.8, decay2:float=0.2, nplus=1.5, nminus=0.75):
-
-        defaults = dict(pmin=pmin, pmax=pmax, pow=pow, decay=decay, decay2=decay2, nplus=nplus, nminus=nminus)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        if closure is None:
-            raise RuntimeError("CD requires closure")
-
-        params = TensorList(var.params)
-        p_prev = self.get_state(params, "p_prev", init=params, cls=TensorList)
-
-        f_0 = var.get_loss(False)
-        step_size = self.global_state.get('step_size', 1)
-
-        # ------------------------ hard reset on infinite loss ----------------------- #
-        if not math.isfinite(f_0):
-            del self.global_state['f_prev']
-            var.update = params - p_prev
-            self.global_state.clear()
-            self.state.clear()
-            self.global_state["step_size"] = step_size / 10
-            return var
-
-        # ---------------------------- soft reset if stuck --------------------------- #
-        if "igrad" in self.state[params[0]]:
-            n_bad = self.global_state.get('n_bad', 0)
-
-            f_prev = self.global_state.get("f_prev", None)
-            if f_prev is not None:
-
-                decay2 = self.defaults["decay2"]
-                decay = self.global_state.get("decay", self.defaults["decay"])
-
-                if f_0 >= f_prev:
-
-                    igrad = self.get_state(params, "igrad", cls=TensorList)
-                    del self.global_state['f_prev']
-
-                    # undo previous update
-                    var.update = params - p_prev
-
-                    # increment n_bad
-                    self.global_state['n_bad'] = n_bad + 1
-
-                    # decay step size
-                    self.global_state['step_size'] = step_size * self.defaults["nminus"]
-
-                    # soft reset
-                    if n_bad > 0:
-                        igrad *= decay
-                        self.global_state["decay"] = decay*decay2
-                        self.global_state['n_bad'] = 0
-
-                    return var
-
-                else:
-                    # increase step size and reset n_bad
-                    self.global_state['step_size'] = step_size * self.defaults["nplus"]
-                    self.global_state['n_bad'] = 0
-                    self.global_state["decay"] = self.defaults["decay"]
-
-            self.global_state['f_prev'] = float(f_0)
-
-        # ------------------------------ determine index ----------------------------- #
-        idx, igrad = _icd_get_idx(self, params)
-
-        # -------------------------- find descent direction -------------------------- #
-        h_vec = self.get_state(params, 'h_vec', init=lambda x: torch.full_like(x, 1e-3), cls=TensorList)
-        h = float(h_vec.flat_get(idx))
-
-        params.flat_set_lambda_(idx, lambda x: x + h)
-        f_p = closure(False)
-
-        params.flat_set_lambda_(idx, lambda x: x - 2*h)
-        f_n = closure(False)
-        params.flat_set_lambda_(idx, lambda x: x + h)
-
-        # ---------------------------------- adapt h --------------------------------- #
-        if f_0 <= f_p and f_0 <= f_n:
-            h_vec.flat_set_lambda_(idx, lambda x: max(x/2, 1e-10))
-        else:
-            if abs(f_0 - f_n) < 1e-12 or abs((f_p - f_0) / (f_0 - f_n) - 1) < 1e-2:
-                h_vec.flat_set_lambda_(idx, lambda x: min(x*2, 1e10))
-
-        # ------------------------------- update igrad ------------------------------- #
-        if f_0 < f_p and f_0 < f_n: alpha = 0
-        else: alpha = (f_p - f_n) / (2*h)
-
-        igrad.flat_set_(idx, alpha)
-
-        # ----------------------------- create the update ---------------------------- #
-        var.update = igrad * step_size
-        p_prev.copy_(params)
-        return var
-
-
-class CCDLS(Module):
-    """CCD with line search instead of adaptive step size.
-
-    Args:
-        pmin (float, optional): multiplier to probability of picking the lowest magnitude gradient. Defaults to 0.1.
-        pmax (float, optional): multiplier to probability of picking the largest magnitude gradient. Defaults to 1.0.
-        pow (int, optional): power transform to probabilities. Defaults to 2.
-        decay (float, optional): accumulated gradient decay on failed step. Defaults to 0.5.
-        decay2 (float, optional): decay multiplier decay on failed step. Defaults to 0.25.
-        maxiter (int, optional): max number of line search iterations.
-    """
-    def __init__(self, pmin=0.1, pmax=1.0, pow=2, decay=0.8, decay2=0.2, maxiter=10, ):
-        defaults = dict(pmin=pmin, pmax=pmax, pow=pow, maxiter=maxiter, decay=decay, decay2=decay2)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        if closure is None:
-            raise RuntimeError("CD requires closure")
-
-        params = TensorList(var.params)
-        finfo = torch.finfo(params[0].dtype)
-        f_0 = var.get_loss(False)
-
-        # ------------------------------ determine index ----------------------------- #
-        idx, igrad = _icd_get_idx(self, params)
-
-        # -------------------------- find descent direction -------------------------- #
-        h_vec = self.get_state(params, 'h_vec', init=lambda x: torch.full_like(x, 1e-3), cls=TensorList)
-        h = float(h_vec.flat_get(idx))
-
-        params.flat_set_lambda_(idx, lambda x: x + h)
-        f_p = closure(False)
-
-        params.flat_set_lambda_(idx, lambda x: x - 2*h)
-        f_n = closure(False)
-        params.flat_set_lambda_(idx, lambda x: x + h)
-
-        # ---------------------------------- adapt h --------------------------------- #
-        if f_0 <= f_p and f_0 <= f_n:
-            h_vec.flat_set_lambda_(idx, lambda x: max(x/2, finfo.tiny * 2))
-        else:
-            # here eps, not tiny
-            if abs(f_0 - f_n) < finfo.eps or abs((f_p - f_0) / (f_0 - f_n) - 1) < 1e-2:
-                h_vec.flat_set_lambda_(idx, lambda x: min(x*2, finfo.max / 2))
-
-        # ------------------------------- update igrad ------------------------------- #
-        if f_0 < f_p and f_0 < f_n: alpha = 0
-        else: alpha = (f_p - f_n) / (2*h)
-
-        igrad.flat_set_(idx, alpha)
-
-        # -------------------------------- line search ------------------------------- #
-        x0 = params.clone()
-        def f(a):
-            params.sub_(igrad, alpha=a)
-            loss = closure(False)
-            params.copy_(x0)
-            return loss
-
-        a_prev = self.global_state.get('a_prev', 1)
-        a, f_a, niter = adaptive_tracking(f, a_prev, maxiter=self.defaults['maxiter'], f_0=f_0)
-        if (a is None) or (not math.isfinite(a)) or (not math.isfinite(f_a)):
-            a = 0
-
-        # -------------------------------- set a_prev -------------------------------- #
-        decay2 = self.defaults["decay2"]
-        decay = self.global_state.get("decay", self.defaults["decay"])
-
-        if abs(a) > finfo.tiny * 2:
-            assert f_a < f_0
-            self.global_state['a_prev'] = max(min(a, finfo.max / 2), finfo.tiny * 2)
-            self.global_state["decay"] = self.defaults["decay"]
-
-        # ---------------------------- soft reset on fail ---------------------------- #
-        else:
-            igrad *= decay
-            self.global_state["decay"] = decay*decay2
-            self.global_state['a_prev'] = a_prev / 2
-
-        # -------------------------------- set update -------------------------------- #
-        var.update = igrad * a
-        return var

torchzero/utils/optimizer.py

@@ -110,7 +110,7 @@ def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], param
         for i, param in enumerate(params):
             s = state[param]
             if key not in s:
-                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                if must_exist: raise KeyError(f"Key `{key}` doesn't exist in state with keys {tuple(s.keys())}")
                 s[key] = _make_initial_state_value(param, init, i)
             values.append(s[key])
         return values
@@ -125,7 +125,7 @@ def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], param
         s = state[param]
         for k_i, key in enumerate(keys):
             if key not in s:
-                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                if must_exist: raise KeyError(f"Key `{key}` doesn't exist in state with keys {tuple(s.keys())}")
                 k_init = init[k_i] if isinstance(init, (list,tuple)) else init
                 s[key] = _make_initial_state_value(param, k_init, i)
             values[k_i].append(s[key])

torchzero/utils/python_tools.py

@@ -67,3 +67,4 @@ def safe_dict_update_(d1_:dict, d2:dict):
     inter = set(d1_.keys()).intersection(d2.keys())
     if len(inter) > 0: raise RuntimeError(f"Duplicate keys {inter}")
     d1_.update(d2)
+

{torchzero-0.3.13.dist-info → torchzero-0.3.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.13
+Version: 0.3.14
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 Project-URL: Homepage, https://github.com/inikishev/torchzero