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

torchzero/modules/quasi_newton/damping.py

@@ -0,0 +1,105 @@
+import math
+from typing import Literal, Protocol, overload
+
+import torch
+
+from ...utils import TensorList
+from ...utils.linalg.linear_operator import DenseInverse, LinearOperator
+from ..functional import safe_clip
+
+
+class DampingStrategy(Protocol):
+    def __call__(
+        self,
+        s: torch.Tensor,
+        y: torch.Tensor,
+        g: torch.Tensor,
+        H: LinearOperator,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return s, y
+
+def _sy_Hs_sHs(s:torch.Tensor, y:torch.Tensor, H:LinearOperator):
+    if isinstance(H, DenseInverse):
+        Hs = H.solve(y)
+        sHs = y.dot(Hs)
+    else:
+        Hs = H.matvec(s)
+        sHs = s.dot(Hs)
+
+    return s.dot(y), Hs, sHs
+
+
+
+def powell_damping(s:torch.Tensor, y:torch.Tensor, g:torch.Tensor, H:LinearOperator, u=0.2):
+    # here H is hessian! not the inverse
+
+    sy, Hs, sHs = _sy_Hs_sHs(s, y, H)
+    if sy < u*sHs:
+        phi = ((1-u) * sHs) / safe_clip((sHs - sy))
+        s = phi * s + (1 - phi) * Hs
+
+    return s, y
+
+def double_damping(s:torch.Tensor, y:torch.Tensor, g:torch.Tensor, H:LinearOperator, u1=0.2, u2=1/3):
+    # Goldfarb, Donald, Yi Ren, and Achraf Bahamou. "Practical quasi-newton methods for training deep neural networks." Advances in Neural Information Processing Systems 33 (2020): 2386-2396.
+
+    # Powell’s damping on H
+    sy, Hs, sHs = _sy_Hs_sHs(s, y, H)
+    if sy < u1*sHs:
+        phi = ((1-u1) * sHs) / safe_clip(sHs - sy)
+        s = phi * s + (1 - phi) * Hs
+
+    # Powell’s damping with B = I
+    sy = s.dot(y)
+    ss = s.dot(s)
+
+    if sy < u2*ss:
+        phi = ((1-u2) * ss) / safe_clip(ss - sy)
+        y = phi * y + (1 - phi) * s
+
+    return s, y
+
+
+
+_DAMPING_KEYS = Literal["powell", "double"]
+_DAMPING_STRATEGIES: dict[_DAMPING_KEYS, DampingStrategy] = {
+    "powell": powell_damping,
+    "double": double_damping,
+}
+
+
+DampingStrategyType = _DAMPING_KEYS | DampingStrategy | None
+
+@overload
+def apply_damping(
+    strategy: DampingStrategyType,
+    s: torch.Tensor,
+    y: torch.Tensor,
+    g: torch.Tensor,
+    H: LinearOperator,
+) -> tuple[torch.Tensor, torch.Tensor]: ...
+@overload
+def apply_damping(
+    strategy: DampingStrategyType,
+    s: TensorList,
+    y: TensorList,
+    g: TensorList,
+    H: LinearOperator,
+) -> tuple[TensorList, TensorList]: ...
+def apply_damping(
+    strategy: DampingStrategyType,
+    s,
+    y,
+    g,
+    H: LinearOperator,
+):
+    if strategy is None: return s, y
+    if isinstance(strategy, str): strategy = _DAMPING_STRATEGIES[strategy]
+
+    if isinstance(s, TensorList):
+        assert isinstance(y, TensorList) and isinstance(g, TensorList)
+        s_vec, y_vec = strategy(s.to_vec(), y.to_vec(), g.to_vec(), H)
+        return s.from_vec(s_vec), y.from_vec(y_vec)
+
+    assert isinstance(y, torch.Tensor) and isinstance(g, torch.Tensor)
+    return strategy(s, y, g, H)

torchzero/modules/quasi_newton/diagonal_quasi_newton.py

@@ -1,163 +1,167 @@
-from collections.abc import Callable
-
-import torch
-
-from .quasi_newton import (
-    HessianUpdateStrategy,
-    _HessianUpdateStrategyDefaults,
-    _InverseHessianUpdateStrategyDefaults,
-    _safe_clip,
-)
-
-
-def _diag_Bv(self: HessianUpdateStrategy):
-    B, is_inverse = self.get_B()
-
-    if is_inverse:
-        H=B
-        def Hxv(v): return v/H
-        return Hxv
-
-    def Bv(v): return B*v
-    return Bv
-
-def _diag_Hv(self: HessianUpdateStrategy):
-    H, is_inverse = self.get_H()
-
-    if is_inverse:
-        B=H
-        def Bxv(v): return v/B
-        return Bxv
-
-    def Hv(v): return H*v
-    return Hv
-
-def diagonal_bfgs_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
-    sy = s.dot(y)
-    if sy < tol: return H
-
-    sy_sq = _safe_clip(sy**2)
-
-    num1 = (sy + (y * H * y)) * s*s
-    term1 = num1.div_(sy_sq)
-    num2 = (H * y * s).add_(s * y * H)
-    term2 = num2.div_(sy)
-    H += term1.sub_(term2)
-    return H
-
-class DiagonalBFGS(_InverseHessianUpdateStrategyDefaults):
-    """Diagonal BFGS. This is simply BFGS with only the diagonal being updated and used. It doesn't satisfy the secant equation but may still be useful."""
-    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, setting):
-        return diagonal_bfgs_H_(H=H, s=s, y=y, tol=setting['tol'])
-
-    def _init_M(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
-    def make_Bv(self): return _diag_Bv(self)
-    def make_Hv(self): return _diag_Hv(self)
-
-def diagonal_sr1_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol:float):
-    z = s - H*y
-    denom = z.dot(y)
-
-    z_norm = torch.linalg.norm(z) # pylint:disable=not-callable
-    y_norm = torch.linalg.norm(y) # pylint:disable=not-callable
-
-    # if y_norm*z_norm < tol: return H
-
-    # check as in Nocedal, Wright. “Numerical optimization” 2nd p.146
-    if denom.abs() <= tol * y_norm * z_norm: return H # pylint:disable=not-callable
-    H += (z*z).div_(_safe_clip(denom))
-    return H
-class DiagonalSR1(_InverseHessianUpdateStrategyDefaults):
-    """Diagonal SR1. This is simply SR1 with only the diagonal being updated and used. It doesn't satisfy the secant equation but may still be useful."""
-    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, setting):
-        return diagonal_sr1_(H=H, s=s, y=y, tol=setting['tol'])
-    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
-        return diagonal_sr1_(H=B, s=y, y=s, tol=setting['tol'])
-
-    def _init_M(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
-    def make_Bv(self): return _diag_Bv(self)
-    def make_Hv(self): return _diag_Hv(self)
-
-
-
-# Zhu M., Nazareth J. L., Wolkowicz H. The quasi-Cauchy relation and diagonal updating //SIAM Journal on Optimization. – 1999. – Т. 9. – №. 4. – С. 1192-1204.
-def diagonal_qc_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor):
-    denom = _safe_clip((s**4).sum())
-    num = s.dot(y) - (s*B).dot(s)
-    B += s**2 * (num/denom)
-    return B
-
-class DiagonalQuasiCauchi(_HessianUpdateStrategyDefaults):
-    """Diagonal quasi-cauchi method.
-
-    Reference:
-        Zhu M., Nazareth J. L., Wolkowicz H. The quasi-Cauchy relation and diagonal updating //SIAM Journal on Optimization. – 1999. – Т. 9. – №. 4. – С. 1192-1204.
-    """
-    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
-        return diagonal_qc_B_(B=B, s=s, y=y)
-
-    def _init_M(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
-    def make_Bv(self): return _diag_Bv(self)
-    def make_Hv(self): return _diag_Hv(self)
-
-# Leong, Wah June, Sharareh Enshaei, and Sie Long Kek. "Diagonal quasi-Newton methods via least change updating principle with weighted Frobenius norm." Numerical Algorithms 86 (2021): 1225-1241.
-def diagonal_wqc_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor):
-    E_sq = s**2 * B**2
-    denom = _safe_clip((s*E_sq).dot(s))
-    num = s.dot(y) - (s*B).dot(s)
-    B += E_sq * (num/denom)
-    return B
-
-class DiagonalWeightedQuasiCauchi(_HessianUpdateStrategyDefaults):
-    """Diagonal quasi-cauchi method.
-
-    Reference:
-        Leong, Wah June, Sharareh Enshaei, and Sie Long Kek. "Diagonal quasi-Newton methods via least change updating principle with weighted Frobenius norm." Numerical Algorithms 86 (2021): 1225-1241.
-    """
-    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
-        return diagonal_wqc_B_(B=B, s=s, y=y)
-
-    def _init_M(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
-    def make_Bv(self): return _diag_Bv(self)
-    def make_Hv(self): return _diag_Hv(self)
-
-
-# Andrei, Neculai. "A diagonal quasi-Newton updating method for unconstrained optimization." Numerical Algorithms 81.2 (2019): 575-590.
-def dnrtr_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor):
-    denom = _safe_clip((s**4).sum())
-    num = s.dot(y) + s.dot(s) - (s*B).dot(s)
-    B += s**2 * (num/denom) - 1
-    return B
-
-class DNRTR(_HessianUpdateStrategyDefaults):
-    """Diagonal quasi-newton method.
-
-    Reference:
-        Andrei, Neculai. "A diagonal quasi-Newton updating method for unconstrained optimization." Numerical Algorithms 81.2 (2019): 575-590.
-    """
-    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
-        return diagonal_wqc_B_(B=B, s=s, y=y)
-
-    def _init_M(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
-    def make_Bv(self): return _diag_Bv(self)
-    def make_Hv(self): return _diag_Hv(self)
-
-# Nosrati, Mahsa, and Keyvan Amini. "A new diagonal quasi-Newton algorithm for unconstrained optimization problems." Applications of Mathematics 69.4 (2024): 501-512.
-def new_dqn_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor):
-    denom = _safe_clip((s**4).sum())
-    num = s.dot(y)
-    B += s**2 * (num/denom)
-    return B
-
-class NewDQN(_HessianUpdateStrategyDefaults):
-    """Diagonal quasi-newton method.
-
-    Reference:
-        Nosrati, Mahsa, and Keyvan Amini. "A new diagonal quasi-Newton algorithm for unconstrained optimization problems." Applications of Mathematics 69.4 (2024): 501-512.
-    """
-    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
-        return new_dqn_B_(B=B, s=s, y=y)
-
-    def _init_M(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
-    def make_Bv(self): return _diag_Bv(self)
-    def make_Hv(self): return _diag_Hv(self)
+from typing import Any, Literal
+from collections.abc import Callable
+
+import torch
+from ...core import Chainable
+from .quasi_newton import (
+    HessianUpdateStrategy,
+    _HessianUpdateStrategyDefaults,
+    _InverseHessianUpdateStrategyDefaults,
+)
+
+from ..functional import safe_clip
+
+
+def diagonal_bfgs_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
+    sy = s.dot(y)
+    if sy < tol: return H
+
+    sy_sq = safe_clip(sy**2)
+
+    num1 = (sy + (y * H * y)) * s*s
+    term1 = num1.div_(sy_sq)
+    num2 = (H * y * s).add_(s * y * H)
+    term2 = num2.div_(sy)
+    H += term1.sub_(term2)
+    return H
+
+class DiagonalBFGS(_InverseHessianUpdateStrategyDefaults):
+    """Diagonal BFGS. This is simply BFGS with only the diagonal being updated and used. It doesn't satisfy the secant equation but may still be useful."""
+    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, setting):
+        return diagonal_bfgs_H_(H=H, s=s, y=y, tol=setting['tol'])
+
+    def initialize_P(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
+
+def diagonal_sr1_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol:float):
+    z = s - H*y
+    denom = z.dot(y)
+
+    z_norm = torch.linalg.norm(z) # pylint:disable=not-callable
+    y_norm = torch.linalg.norm(y) # pylint:disable=not-callable
+
+    # if y_norm*z_norm < tol: return H
+
+    # check as in Nocedal, Wright. “Numerical optimization” 2nd p.146
+    if denom.abs() <= tol * y_norm * z_norm: return H # pylint:disable=not-callable
+    H += (z*z).div_(safe_clip(denom))
+    return H
+class DiagonalSR1(_InverseHessianUpdateStrategyDefaults):
+    """Diagonal SR1. This is simply SR1 with only the diagonal being updated and used. It doesn't satisfy the secant equation but may still be useful."""
+    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, setting):
+        return diagonal_sr1_(H=H, s=s, y=y, tol=setting['tol'])
+    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
+        return diagonal_sr1_(H=B, s=y, y=s, tol=setting['tol'])
+
+    def initialize_P(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
+
+
+
+# Zhu M., Nazareth J. L., Wolkowicz H. The quasi-Cauchy relation and diagonal updating //SIAM Journal on Optimization. – 1999. – Т. 9. – №. 4. – С. 1192-1204.
+def diagonal_qc_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor):
+    denom = safe_clip((s**4).sum())
+    num = s.dot(y) - (s*B).dot(s)
+    B += s**2 * (num/denom)
+    return B
+
+class DiagonalQuasiCauchi(_HessianUpdateStrategyDefaults):
+    """Diagonal quasi-cauchi method.
+
+    Reference:
+        Zhu M., Nazareth J. L., Wolkowicz H. The quasi-Cauchy relation and diagonal updating //SIAM Journal on Optimization. – 1999. – Т. 9. – №. 4. – С. 1192-1204.
+    """
+    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
+        return diagonal_qc_B_(B=B, s=s, y=y)
+
+    def initialize_P(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
+
+# Leong, Wah June, Sharareh Enshaei, and Sie Long Kek. "Diagonal quasi-Newton methods via least change updating principle with weighted Frobenius norm." Numerical Algorithms 86 (2021): 1225-1241.
+def diagonal_wqc_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor):
+    E_sq = s**2 * B**2
+    denom = safe_clip((s*E_sq).dot(s))
+    num = s.dot(y) - (s*B).dot(s)
+    B += E_sq * (num/denom)
+    return B
+
+class DiagonalWeightedQuasiCauchi(_HessianUpdateStrategyDefaults):
+    """Diagonal quasi-cauchi method.
+
+    Reference:
+        Leong, Wah June, Sharareh Enshaei, and Sie Long Kek. "Diagonal quasi-Newton methods via least change updating principle with weighted Frobenius norm." Numerical Algorithms 86 (2021): 1225-1241.
+    """
+    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
+        return diagonal_wqc_B_(B=B, s=s, y=y)
+
+    def initialize_P(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
+
+def _truncate(B: torch.Tensor, lb, ub):
+    return torch.where((B>lb).logical_and(B<ub), B, 1)
+
+# Andrei, Neculai. "A diagonal quasi-Newton updating method for unconstrained optimization." Numerical Algorithms 81.2 (2019): 575-590.
+def dnrtr_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor):
+    denom = safe_clip((s**4).sum())
+    num = s.dot(y) + s.dot(s) - (s*B).dot(s)
+    B += s**2 * (num/denom) - 1
+    return B
+
+class DNRTR(HessianUpdateStrategy):
+    """Diagonal quasi-newton method.
+
+    Reference:
+        Andrei, Neculai. "A diagonal quasi-Newton updating method for unconstrained optimization." Numerical Algorithms 81.2 (2019): 575-590.
+    """
+    def __init__(
+        self,
+        lb: float = 1e-2,
+        ub: float = 1e5,
+        init_scale: float | Literal["auto"] = "auto",
+        tol: float = 1e-32,
+        ptol: float | None = 1e-32,
+        ptol_restart: bool = False,
+        gtol: float | None = 1e-32,
+        restart_interval: int | None | Literal['auto'] = None,
+        beta: float | None = None,
+        update_freq: int = 1,
+        scale_first: bool = False,
+        concat_params: bool = True,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(lb=lb, ub=ub)
+        super().__init__(
+            defaults=defaults,
+            init_scale=init_scale,
+            tol=tol,
+            ptol=ptol,
+            ptol_restart=ptol_restart,
+            gtol=gtol,
+            restart_interval=restart_interval,
+            beta=beta,
+            update_freq=update_freq,
+            scale_first=scale_first,
+            concat_params=concat_params,
+            inverse=False,
+            inner=inner,
+        )
+
+    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
+        return diagonal_wqc_B_(B=B, s=s, y=y)
+
+    def modify_B(self, B, state, setting):
+        return _truncate(B, setting['lb'], setting['ub'])
+
+    def initialize_P(self, size:int, device, dtype, is_inverse:bool): return torch.ones(size, device=device, dtype=dtype)
+
+# Nosrati, Mahsa, and Keyvan Amini. "A new diagonal quasi-Newton algorithm for unconstrained optimization problems." Applications of Mathematics 69.4 (2024): 501-512.
+def new_dqn_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor):
+    denom = safe_clip((s**4).sum())
+    num = s.dot(y)
+    B += s**2 * (num/denom)
+    return B
+
+class NewDQN(DNRTR):
+    """Diagonal quasi-newton method.
+
+    Reference:
+        Nosrati, Mahsa, and Keyvan Amini. "A new diagonal quasi-Newton algorithm for unconstrained optimization problems." Applications of Mathematics 69.4 (2024): 501-512.
+    """
+    def update_B(self, B, s, y, p, g, p_prev, g_prev, state, setting):
+        return new_dqn_B_(B=B, s=s, y=y)