torchzero 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

torchzero/modules/quasi_newton/sg2.py

@@ -1,29 +1,39 @@
 import torch
 
-from ...core import Module, Chainable, step
-from ...utils import TensorList, vec_to_tensors
-from ..second_order.newton import _newton_step, _get_H
+from ...core import Chainable, Transform
+from ...utils import TensorList, unpack_dicts, unpack_states, vec_to_tensors_
+from ...linalg.linear_operator import Dense
+
 
 def sg2_(
     delta_g: torch.Tensor,
     cd: torch.Tensor,
 ) -> torch.Tensor:
-    """cd is c * perturbation, and must be multiplied by two if hessian estimate is two-sided
-    (or divide delta_g by two)."""
+    """cd is c * perturbation."""
 
-    M = torch.outer(1.0 / cd, delta_g)
+    M = torch.outer(0.5 / cd, delta_g)
     H_hat = 0.5 * (M + M.T)
 
     return H_hat
 
 
 
-class SG2(Module):
+class SG2(Transform):
     """second-order stochastic gradient
 
+    2SPSA (second-order SPSA)
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.SPSA(),
+        tz.m.SG2(),
+        tz.m.LR(1e-2),
+    )
+    ```
+
     SG2 with line search
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SG2(),
         tz.m.Backtracking()
@@ -32,9 +42,9 @@ class SG2(Module):
 
     SG2 with trust region
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
-        tz.m.LevenbergMarquardt(tz.m.SG2()),
+        tz.m.LevenbergMarquardt(tz.m.SG2(beta=0.75. n_samples=4)),
     )
     ```
 
@@ -43,24 +53,22 @@ class SG2(Module):
     def __init__(
         self,
         n_samples: int = 1,
-        h: float = 1e-2,
+        n_first_step_samples: int = 10,
+        start_step: int = 10,
         beta: float | None = None,
-        damping: float = 0,
-        eigval_fn=None,
-        one_sided: bool = False, # one-sided hessian
-        use_lstsq: bool = True,
+        damping: float = 1e-4,
+        h: float = 1e-2,
         seed=None,
+        update_freq: int = 1,
         inner: Chainable | None = None,
     ):
-        defaults = dict(n_samples=n_samples, h=h, beta=beta, damping=damping, eigval_fn=eigval_fn, one_sided=one_sided, seed=seed, use_lstsq=use_lstsq)
-        super().__init__(defaults)
-
-        if inner is not None: self.set_child('inner', inner)
+        defaults = dict(n_samples=n_samples, h=h, beta=beta, damping=damping, seed=seed, start_step=start_step, n_first_step_samples=n_first_step_samples)
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def update(self, objective):
-        k = self.global_state.get('step', 0) + 1
-        self.global_state["step"] = k
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+        k = self.increment_counter("step", 0)
 
         params = TensorList(objective.params)
         closure = objective.closure
@@ -68,36 +76,28 @@ class SG2(Module):
             raise RuntimeError("closure is required for SG2")
         generator = self.get_generator(params[0].device, self.defaults["seed"])
 
-        h = self.get_settings(params, "h")
+        h = unpack_dicts(settings, "h")
         x_0 = params.clone()
-        n_samples = self.defaults["n_samples"]
+        n_samples = fs["n_samples"]
+        if k == 0: n_samples = fs["n_first_step_samples"]
         H_hat = None
 
+        # compute new approximation
         for i in range(n_samples):
             # generate perturbation
             cd = params.rademacher_like(generator=generator).mul_(h)
 
-            # one sided
-            if self.defaults["one_sided"]:
-                g_0 = TensorList(objective.get_grads())
-                params.add_(cd)
-                closure()
+            # two sided hessian approximation
+            params.add_(cd)
+            closure()
+            g_p = params.grad.fill_none_(params)
 
-                g_p = params.grad.fill_none_(params)
-                delta_g = (g_p - g_0) * 2
+            params.copy_(x_0)
+            params.sub_(cd)
+            closure()
+            g_n = params.grad.fill_none_(params)
 
-            # two sided
-            else:
-                params.add_(cd)
-                closure()
-                g_p = params.grad.fill_none_(params)
-
-                params.copy_(x_0)
-                params.sub_(cd)
-                closure()
-                g_n = params.grad.fill_none_(params)
-
-                delta_g = g_p - g_n
+            delta_g = g_p - g_n
 
             # restore params
             params.set_(x_0)
@@ -114,179 +114,43 @@ class SG2(Module):
         assert H_hat is not None
         if n_samples > 1: H_hat /= n_samples
 
+        # add damping
+        if fs["damping"] != 0:
+            reg = torch.eye(H_hat.size(0), device=H_hat.device, dtype=H_hat.dtype).mul_(fs["damping"])
+            H_hat += reg
+
         # update H
         H = self.global_state.get("H", None)
         if H is None: H = H_hat
         else:
-            beta = self.defaults["beta"]
-            if beta is None: beta = k / (k+1)
+            beta = fs["beta"]
+            if beta is None: beta = (k+1) / (k+2)
             H.lerp_(H_hat, 1-beta)
 
         self.global_state["H"] = H
 
 
     @torch.no_grad
-    def apply(self, objective):
-        dir = _newton_step(
-            objective=objective,
-            H = self.global_state["H"],
-            damping = self.defaults["damping"],
-            inner = self.children.get("inner", None),
-            H_tfm=None,
-            eigval_fn=self.defaults["eigval_fn"],
-            use_lstsq=self.defaults["use_lstsq"],
-            g_proj=None,
-        )
-
-        objective.updates = vec_to_tensors(dir, objective.params)
+    def apply_states(self, objective, states, settings):
+        fs = settings[0]
+        updates = objective.get_updates()
+
+        H: torch.Tensor = self.global_state["H"]
+        k = self.global_state["step"]
+        if k < fs["start_step"]:
+            # don't precondition yet
+            # I guess we can try using trace to scale the update
+            # because it will have horrible scaling otherwise
+            torch._foreach_div_(updates, H.trace())
+            return objective
+
+        b = torch.cat([t.ravel() for t in updates])
+        sol = torch.linalg.lstsq(H, b).solution # pylint:disable=not-callable
+
+        vec_to_tensors_(sol, updates)
         return objective
 
-    def get_H(self,objective=...):
-        return _get_H(self.global_state["H"], self.defaults["eigval_fn"])
-
-
-
-
-# two sided
-# we have g via x + d, x - d
-# H via g(x + d), g(x - d)
-# 1 is x, x+2d
-# 2 is x, x-2d
-# 5 evals in total
-
-# one sided
-# g via x, x + d
-# 1 is x, x + d
-# 2 is x, x - d
-# 3 evals and can use two sided for g_0
-
-class SPSA2(Module):
-    """second-order SPSA
-
-    SPSA2 with line search
-    ```python
-    opt = tz.Modular(
-        model.parameters(),
-        tz.m.SPSA2(),
-        tz.m.Backtracking()
-    )
-    ```
-
-    SPSA2 with trust region
-    ```python
-    opt = tz.Modular(
-        model.parameters(),
-        tz.m.LevenbergMarquardt(tz.m.SPSA2()),
-    )
-    ```
-    """
-
-    def __init__(
-        self,
-        n_samples: int = 1,
-        h: float = 1e-2,
-        beta: float | None = None,
-        damping: float = 0,
-        eigval_fn=None,
-        use_lstsq: bool = True,
-        seed=None,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(n_samples=n_samples, h=h, beta=beta, damping=damping, eigval_fn=eigval_fn, seed=seed, use_lstsq=use_lstsq)
-        super().__init__(defaults)
-
-        if inner is not None: self.set_child('inner', inner)
-
-    @torch.no_grad
-    def update(self, objective):
-        k = self.global_state.get('step', 0) + 1
-        self.global_state["step"] = k
-
-        params = TensorList(objective.params)
-        closure = objective.closure
-        if closure is None:
-            raise RuntimeError("closure is required for SPSA2")
-
-        generator = self.get_generator(params[0].device, self.defaults["seed"])
-
-        h = self.get_settings(params, "h")
-        x_0 = params.clone()
-        n_samples = self.defaults["n_samples"]
-        H_hat = None
-        g_0 = None
-
-        for i in range(n_samples):
-            # perturbations for g and H
-            cd_g = params.rademacher_like(generator=generator).mul_(h)
-            cd_H = params.rademacher_like(generator=generator).mul_(h)
-
-            # evaluate 4 points
-            x_p = x_0 + cd_g
-            x_n = x_0 - cd_g
-
-            params.set_(x_p)
-            f_p = closure(False)
-            params.add_(cd_H)
-            f_pp = closure(False)
-
-            params.set_(x_n)
-            f_n = closure(False)
-            params.add_(cd_H)
-            f_np = closure(False)
-
-            g_p_vec = (f_pp - f_p) / cd_H
-            g_n_vec = (f_np - f_n) / cd_H
-            delta_g = g_p_vec - g_n_vec
-
-            # restore params
-            params.set_(x_0)
-
-            # compute grad
-            g_i = (f_p - f_n) / (2 * cd_g)
-            if g_0 is None: g_0 = g_i
-            else: g_0 += g_i
-
-            # compute H hat
-            H_i = sg2_(
-                delta_g = delta_g.to_vec().div_(2.0),
-                cd = cd_g.to_vec(), # The interval is measured by the original 'cd'
-            )
-            if H_hat is None: H_hat = H_i
-            else: H_hat += H_i
-
-        assert g_0 is not None and H_hat is not None
-        if n_samples > 1:
-            g_0 /= n_samples
-            H_hat /= n_samples
-
-        # set grad to approximated grad
-        objective.grads = g_0
+    def get_H(self, objective=...):
+        return Dense(self.global_state["H"])
 
-        # update H
-        H = self.global_state.get("H", None)
-        if H is None: H = H_hat
-        else:
-            beta = self.defaults["beta"]
-            if beta is None: beta = k / (k+1)
-            H.lerp_(H_hat, 1-beta)
-
-        self.global_state["H"] = H
-
-    @torch.no_grad
-    def apply(self, objective):
-        dir = _newton_step(
-            objective=objective,
-            H = self.global_state["H"],
-            damping = self.defaults["damping"],
-            inner = self.children.get("inner", None),
-            H_tfm=None,
-            eigval_fn=self.defaults["eigval_fn"],
-            use_lstsq=self.defaults["use_lstsq"],
-            g_proj=None,
-        )
-
-        objective.updates = vec_to_tensors(dir, objective.params)
-        return objective
 
-    def get_H(self,objective=...):
-        return _get_H(self.global_state["H"], self.defaults["eigval_fn"])

torchzero/modules/restarts/restars.py

@@ -8,8 +8,8 @@ from ...core import Chainable, Module, Objective
 from ...utils import TensorList
 from ..termination import TerminationCriteriaBase
 
-def _reset_except_self(optimizer, var, self: Module):
-    for m in optimizer.unrolled_modules:
+def _reset_except_self(objective, modules, self: Module):
+    for m in modules:
         if m is not self:
             m.reset()
 
@@ -26,15 +26,15 @@ class RestartStrategyBase(Module, ABC):
             self.set_child('modules', modules)
 
     @abstractmethod
-    def should_reset(self, var: Objective) -> bool:
+    def should_reset(self, objective: Objective) -> bool:
         """returns whether reset should occur"""
 
-    def _reset_on_condition(self, var):
+    def _reset_on_condition(self, objective: Objective):
         modules = self.children.get('modules', None)
 
-        if self.should_reset(var):
+        if self.should_reset(objective):
             if modules is None:
-                var.post_step_hooks.append(partial(_reset_except_self, self=self))
+                objective.post_step_hooks.append(partial(_reset_except_self, self=self))
             else:
                 modules.reset()
 
@@ -78,11 +78,11 @@ class RestartOnStuck(RestartStrategyBase):
         super().__init__(defaults, modules)
 
     @torch.no_grad
-    def should_reset(self, var):
+    def should_reset(self, objective):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
-        params = TensorList(var.params)
+        params = TensorList(objective.params)
         tol = self.defaults['tol']
         if tol is None: tol = torch.finfo(params[0].dtype).tiny * 2
         n_tol = self.defaults['n_tol']
@@ -124,12 +124,12 @@ class RestartEvery(RestartStrategyBase):
         defaults = dict(steps=steps)
         super().__init__(defaults, modules)
 
-    def should_reset(self, var):
+    def should_reset(self, objective):
         step = self.global_state.get('step', 0) + 1
         self.global_state['step'] = step
 
         n = self.defaults['steps']
-        if isinstance(n, str): n = sum(p.numel() for p in var.params if p.requires_grad)
+        if isinstance(n, str): n = sum(p.numel() for p in objective.params if p.requires_grad)
 
         # reset every n steps
         if step % n == 0:
@@ -143,9 +143,9 @@ class RestartOnTerminationCriteria(RestartStrategyBase):
         super().__init__(None, modules)
         self.set_child('criteria', criteria)
 
-    def should_reset(self, var):
+    def should_reset(self, objective):
         criteria = cast(TerminationCriteriaBase, self.children['criteria'])
-        return criteria.should_terminate(var)
+        return criteria.should_terminate(objective)
 
 class PowellRestart(RestartStrategyBase):
     """Powell's two restarting criterions for conjugate gradient methods.
@@ -171,14 +171,14 @@ class PowellRestart(RestartStrategyBase):
         defaults=dict(cond1=cond1, cond2=cond2)
         super().__init__(defaults, modules)
 
-    def should_reset(self, var):
-        g = TensorList(var.get_grads())
+    def should_reset(self, objective):
+        g = TensorList(objective.get_grads())
         cond1 = self.defaults['cond1']; cond2 = self.defaults['cond2']
 
         # -------------------------------- initialize -------------------------------- #
         if 'initialized' not in self.global_state:
             self.global_state['initialized'] = 0
-            g_prev = self.get_state(var.params, 'g_prev', init=g)
+            g_prev = self.get_state(objective.params, 'g_prev', init=g)
             return False
 
         g_g = g.dot(g)
@@ -186,7 +186,7 @@ class PowellRestart(RestartStrategyBase):
         reset = False
         # ------------------------------- 1st condition ------------------------------ #
         if cond1 is not None:
-            g_prev = self.get_state(var.params, 'g_prev', must_exist=True, cls=TensorList)
+            g_prev = self.get_state(objective.params, 'g_prev', must_exist=True, cls=TensorList)
             g_g_prev = g_prev.dot(g)
 
             if g_g_prev.abs() >= cond1 * g_g:
@@ -194,7 +194,7 @@ class PowellRestart(RestartStrategyBase):
 
         # ------------------------------- 2nd condition ------------------------------ #
         if (cond2 is not None) and (not reset):
-            d_g = TensorList(var.get_updates()).dot(g)
+            d_g = TensorList(objective.get_updates()).dot(g)
             if (-1-cond2) * g_g < d_g < (-1 + cond2) * g_g:
                 reset = True
 

torchzero/modules/second_order/inm.py

@@ -3,9 +3,9 @@ from collections.abc import Callable
 import torch
 
 from ...core import Chainable, Transform, HessianMethod
-from ...utils import TensorList, vec_to_tensors, unpack_states
-from ..functional import safe_clip
-from .newton import _get_H, _newton_step
+from ...utils import TensorList, vec_to_tensors_, unpack_states
+from ..opt_utils import safe_clip
+from .newton import _newton_update_state_, _newton_solve, _newton_get_H
 
 @torch.no_grad
 def inm(f:torch.Tensor, J:torch.Tensor, s:torch.Tensor, y:torch.Tensor):
@@ -34,10 +34,10 @@ class ImprovedNewton(Transform):
     def __init__(
         self,
         damping: float = 0,
-        use_lstsq: bool = False,
-        update_freq: int = 1,
-        H_tfm: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, bool]] | Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
         eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        update_freq: int = 1,
+        precompute_inverse: bool | None = None,
+        use_lstsq: bool = False,
         hessian_method: HessianMethod = "batched_autograd",
         h: float = 1e-3,
         inner: Chainable | None = None,
@@ -65,37 +65,45 @@ class ImprovedNewton(Transform):
         x_prev, f_prev = unpack_states(states, objective.params, "x_prev", "f_prev", cls=TensorList)
 
         # initialize on 1st step, do Newton step
-        if "P" not in self.global_state:
+        if "H" not in self.global_state:
             x_prev.copy_(x_list)
             f_prev.copy_(f_list)
-            self.global_state["P"] = J
-            return
+            P = J
 
         # INM update
-        s_list = x_list - x_prev
-        y_list = f_list - f_prev
-        x_prev.copy_(x_list)
-        f_prev.copy_(f_list)
+        else:
+            s_list = x_list - x_prev
+            y_list = f_list - f_prev
+            x_prev.copy_(x_list)
+            f_prev.copy_(f_list)
 
-        self.global_state["P"] = inm(f, J, s=s_list.to_vec(), y=y_list.to_vec())
+            P = inm(f, J, s=s_list.to_vec(), y=y_list.to_vec())
 
+        # update state
+        precompute_inverse = fs["precompute_inverse"]
+        if precompute_inverse is None:
+            precompute_inverse = fs["__update_freq"] >= 10
+
+        _newton_update_state_(
+            H=P,
+            state = self.global_state,
+            damping = fs["damping"],
+            eigval_fn = fs["eigval_fn"],
+            precompute_inverse = precompute_inverse,
+            use_lstsq = fs["use_lstsq"]
+        )
 
     @torch.no_grad
     def apply_states(self, objective, states, settings):
+        updates = objective.get_updates()
         fs = settings[0]
 
-        update = _newton_step(
-            objective = objective,
-            H = self.global_state["P"],
-            damping = fs["damping"],
-            H_tfm = fs["H_tfm"],
-            eigval_fn = None, # it is applied in `update`
-            use_lstsq = fs["use_lstsq"],
-        )
-
-        objective.updates = vec_to_tensors(update, objective.params)
+        b = torch.cat([t.ravel() for t in updates])
+        sol = _newton_solve(b=b, state=self.global_state, use_lstsq=fs["use_lstsq"])
 
+        vec_to_tensors_(sol, updates)
         return objective
 
+
     def get_H(self,objective=...):
-        return _get_H(self.global_state["P"], eigval_fn=None)
+        return _newton_get_H(self.global_state)