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

torchzero/modules/quasi_newton/cg.py

@@ -1,14 +1,47 @@
 from abc import ABC, abstractmethod
+from typing import Literal
 
 import torch
 
-from ...core import Chainable, Transform, apply
-from ...utils import TensorList, as_tensorlist
+from ...core import Chainable, TensorwiseTransform, Transform, apply_transform
+from ...utils import TensorList, as_tensorlist, unpack_dicts, unpack_states
+from .quasi_newton import _safe_clip, HessianUpdateStrategy
 
 
 class ConguateGradientBase(Transform, ABC):
-    """all CGs are the same except beta calculation"""
-    def __init__(self, defaults = None, clip_beta: bool = False, reset_interval: int | None = None, inner: Chainable | None = None):
+    """Base class for conjugate gradient methods. The only difference between them is how beta is calculated.
+
+    This is an abstract class, to use it, subclass it and override `get_beta`.
+
+
+    Args:
+        defaults (dict | None, optional): dictionary of settings defaults. Defaults to None.
+        clip_beta (bool, optional): whether to clip beta to be no less than 0. Defaults to False.
+        reset_interval (int | None | Literal["auto"], optional):
+            interval between resetting the search direction.
+            "auto" means number of dimensions + 1, None means no reset. Defaults to None.
+        inner (Chainable | None, optional): previous direction is added to the output of this module. Defaults to None.
+
+    Example:
+
+        .. code-block:: python
+
+            class PolakRibiere(ConguateGradientBase):
+                def __init__(
+                    self,
+                    clip_beta=True,
+                    reset_interval: int | None = None,
+                    inner: Chainable | None = None
+                ):
+                    super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
+
+                def get_beta(self, p, g, prev_g, prev_d):
+                    denom = prev_g.dot(prev_g)
+                    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
+                    return g.dot(g - prev_g) / denom
+
+    """
+    def __init__(self, defaults = None, clip_beta: bool = False, reset_interval: int | None | Literal['auto'] = None, inner: Chainable | None = None):
         if defaults is None: defaults = {}
         defaults['reset_interval'] = reset_interval
         defaults['clip_beta'] = clip_beta
@@ -17,6 +50,15 @@ class ConguateGradientBase(Transform, ABC):
         if inner is not None:
             self.set_child('inner', inner)
 
+    def reset(self):
+        super().reset()
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_grad')
+        self.global_state.pop('stage', None)
+        self.global_state['step'] = self.global_state.get('step', 1) - 1
+
     def initialize(self, p: TensorList, g: TensorList):
         """runs on first step when prev_grads and prev_dir are not available"""
 
@@ -25,38 +67,55 @@ class ConguateGradientBase(Transform, ABC):
         """returns beta"""
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
         tensors = as_tensorlist(tensors)
         params = as_tensorlist(params)
 
-        step = self.global_state.get('step', 0)
-        prev_dir, prev_grads = self.get_state('prev_dir', 'prev_grad', params=params, cls=TensorList)
+        step = self.global_state.get('step', 0) + 1
+        self.global_state['step'] = step
 
         # initialize on first step
-        if step == 0:
+        if self.global_state.get('stage', 0) == 0:
+            g_prev, d_prev = unpack_states(states, tensors, 'g_prev', 'd_prev', cls=TensorList)
+            d_prev.copy_(tensors)
+            g_prev.copy_(tensors)
             self.initialize(params, tensors)
-            prev_dir.copy_(tensors)
-            prev_grads.copy_(tensors)
-            self.global_state['step'] = step + 1
+            self.global_state['stage'] = 1
+
+        else:
+            # if `update_tensors` was called multiple times before `apply_tensors`,
+            # stage becomes 2
+            self.global_state['stage'] = 2
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = as_tensorlist(tensors)
+        step = self.global_state['step']
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads))
+
+        assert self.global_state['stage'] != 0
+        if self.global_state['stage'] == 1:
+            self.global_state['stage'] = 2
             return tensors
 
+        params = as_tensorlist(params)
+        g_prev, d_prev = unpack_states(states, tensors, 'g_prev', 'd_prev', cls=TensorList)
+
         # get beta
-        beta = self.get_beta(params, tensors, prev_grads, prev_dir)
-        if self.settings[params[0]]['clip_beta']: beta = max(0, beta) # pyright:ignore[reportArgumentType]
-        prev_grads.copy_(tensors)
+        beta = self.get_beta(params, tensors, g_prev, d_prev)
+        if settings[0]['clip_beta']: beta = max(0, beta) # pyright:ignore[reportArgumentType]
 
         # inner step
-        if 'inner' in self.children:
-            tensors = as_tensorlist(apply(self.children['inner'], tensors, params, grads, vars))
-
         # calculate new direction with beta
-        dir = tensors.add_(prev_dir.mul_(beta))
-        prev_dir.copy_(dir)
+        dir = tensors.add_(d_prev.mul_(beta))
+        d_prev.copy_(dir)
 
         # resetting
-        self.global_state['step'] = step + 1
-        reset_interval = self.settings[params[0]]['reset_interval']
-        if reset_interval is not None and (step+1) % reset_interval == 0:
+        reset_interval = settings[0]['reset_interval']
+        if reset_interval == 'auto': reset_interval = tensors.global_numel() + 1
+        if reset_interval is not None and step % reset_interval == 0:
             self.reset()
 
         return dir
@@ -68,7 +127,11 @@ def polak_ribiere_beta(g: TensorList, prev_g: TensorList):
     return g.dot(g - prev_g) / denom
 
 class PolakRibiere(ConguateGradientBase):
-    """Polak-Ribière-Polyak nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this."""
+    """Polak-Ribière-Polyak nonlinear conjugate gradient method.
+
+    .. note::
+        - This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this.
+    """
     def __init__(self, clip_beta=True, reset_interval: int | None = None, inner: Chainable | None = None):
         super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
 
@@ -81,8 +144,12 @@ def fletcher_reeves_beta(gg: torch.Tensor, prev_gg: torch.Tensor):
     return gg / prev_gg
 
 class FletcherReeves(ConguateGradientBase):
-    """Fletcher–Reeves nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None = None, clip_beta=False, inner: Chainable | None = None):
+    """Fletcher–Reeves nonlinear conjugate gradient method.
+
+    .. note::
+        - This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this.
+    """
+    def __init__(self, reset_interval: int | None | Literal['auto'] = 'auto', clip_beta=False, inner: Chainable | None = None):
         super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
 
     def initialize(self, p, g):
@@ -103,8 +170,12 @@ def hestenes_stiefel_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
 
 
 class HestenesStiefel(ConguateGradientBase):
-    """Hestenes–Stiefel nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None = None, clip_beta=False, inner: Chainable | None = None):
+    """Hestenes–Stiefel nonlinear conjugate gradient method.
+
+    .. note::
+        - This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this.
+    """
+    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
         super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
 
     def get_beta(self, p, g, prev_g, prev_d):
@@ -118,8 +189,12 @@ def dai_yuan_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
     return (g.dot(g) / denom).neg()
 
 class DaiYuan(ConguateGradientBase):
-    """Dai–Yuan nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None = None, clip_beta=False, inner: Chainable | None = None):
+    """Dai–Yuan nonlinear conjugate gradient method.
+
+    .. note::
+        - This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this. Although Dai–Yuan formula provides an automatic step size scaling so it is technically possible to omit line search and instead use a small step size.
+    """
+    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
         super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
 
     def get_beta(self, p, g, prev_g, prev_d):
@@ -133,8 +208,12 @@ def liu_storey_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList, ):
     return g.dot(g - prev_g) / denom
 
 class LiuStorey(ConguateGradientBase):
-    """Liu-Storey nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None = None, clip_beta=False, inner: Chainable | None = None):
+    """Liu-Storey nonlinear conjugate gradient method.
+
+    .. note::
+        - This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this.
+    """
+    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
         super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
 
     def get_beta(self, p, g, prev_g, prev_d):
@@ -142,7 +221,11 @@ class LiuStorey(ConguateGradientBase):
 
 # ----------------------------- Conjugate Descent ---------------------------- #
 class ConjugateDescent(Transform):
-    """Conjugate Descent (CD). This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
+    """Conjugate Descent (CD).
+
+    .. note::
+        - This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this.
+    """
     def __init__(self, inner: Chainable | None = None):
         super().__init__(defaults={}, uses_grad=False)
 
@@ -151,10 +234,10 @@ class ConjugateDescent(Transform):
 
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         g = as_tensorlist(tensors)
 
-        prev_d = self.get_state('prev_dir', params=params, cls=TensorList, init = torch.zeros_like)
+        prev_d = unpack_states(states, tensors, 'prev_dir', cls=TensorList, init=torch.zeros_like)
         if 'denom' not in self.global_state:
             self.global_state['denom'] = torch.tensor(0.).to(g[0])
 
@@ -164,7 +247,7 @@ class ConjugateDescent(Transform):
 
         # inner step
         if 'inner' in self.children:
-            g = as_tensorlist(apply(self.children['inner'], g, params, grads, vars))
+            g = as_tensorlist(apply_transform(self.children['inner'], g, params, grads))
 
         dir = g.add_(prev_d.mul_(beta))
         prev_d.copy_(dir)
@@ -186,8 +269,11 @@ def hager_zhang_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList,):
 
 class HagerZhang(ConguateGradientBase):
     """Hager-Zhang nonlinear conjugate gradient method,
-    This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None = None, clip_beta=False, inner: Chainable | None = None):
+
+    .. note::
+        - This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this.
+    """
+    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
         super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
 
     def get_beta(self, p, g, prev_g, prev_d):
@@ -210,9 +296,74 @@ def hs_dy_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
 
 class HybridHS_DY(ConguateGradientBase):
     """HS-DY hybrid conjugate gradient method.
-    This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None = None, clip_beta=False, inner: Chainable | None = None):
+
+    .. note::
+        - This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this.
+    """
+    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
         super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
 
     def get_beta(self, p, g, prev_g, prev_d):
         return hs_dy_beta(g, prev_d, prev_g)
+
+
+def projected_gradient_(H:torch.Tensor, y:torch.Tensor):
+    Hy = H @ y
+    yHy = _safe_clip(y.dot(Hy))
+    H -= (Hy.outer(y) @ H) / yHy
+    return H
+
+class ProjectedGradientMethod(HessianUpdateStrategy): # this doesn't maintain hessian
+    """Projected gradient method.
+
+    .. note::
+        This method uses N^2 memory.
+
+    .. note::
+        This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this.
+
+    .. note::
+        This is not the same as projected gradient descent.
+
+    Reference:
+        Pearson, J. D. (1969). Variable metric methods of minimisation. The Computer Journal, 12(2), 171–178. doi:10.1093/comjnl/12.2.171.
+
+    """
+
+    def __init__(
+        self,
+        init_scale: float | Literal["auto"] = 1,
+        tol: float = 1e-8,
+        ptol: float | None = 1e-10,
+        ptol_reset: bool = False,
+        gtol: float | None = 1e-10,
+        reset_interval: int | None | Literal['auto'] = 'auto',
+        beta: float | None = None,
+        update_freq: int = 1,
+        scale_first: bool = False,
+        scale_second: bool = False,
+        concat_params: bool = True,
+        # inverse: bool = True,
+        inner: Chainable | None = None,
+    ):
+        super().__init__(
+            defaults=None,
+            init_scale=init_scale,
+            tol=tol,
+            ptol=ptol,
+            ptol_reset=ptol_reset,
+            gtol=gtol,
+            reset_interval=reset_interval,
+            beta=beta,
+            update_freq=update_freq,
+            scale_first=scale_first,
+            scale_second=scale_second,
+            concat_params=concat_params,
+            inverse=True,
+            inner=inner,
+        )
+
+
+
+    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, setting):
+        return projected_gradient_(H=H, y=y)

torchzero/modules/quasi_newton/diagonal_quasi_newton.py

@@ -0,0 +1,163 @@
+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)