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

torchzero/modules/second_order/rsn.py

@@ -6,10 +6,10 @@ from typing import Literal
 import torch
 
 from ...core import Chainable, Transform, HVPMethod
-from ...utils import vec_to_tensors
+from ...utils import vec_to_tensors_
 from ...linalg.linear_operator import Sketched
 
-from .newton import _newton_step
+from .newton import _newton_update_state_, _newton_solve
 
 def _qr_orthonormalize(A:torch.Tensor):
     m,n = A.shape
@@ -20,12 +20,10 @@ def _qr_orthonormalize(A:torch.Tensor):
     q, _ = torch.linalg.qr(A) # pylint:disable=not-callable
     return q
 
+
 def _orthonormal_sketch(m, n, dtype, device, generator):
     return _qr_orthonormalize(torch.randn(m, n, dtype=dtype, device=device, generator=generator))
 
-def _gaussian_sketch(m, n, dtype, device, generator):
-    return torch.randn(m, n, dtype=dtype, device=device, generator=generator) / math.sqrt(m)
-
 def _rademacher_sketch(m, n, dtype, device, generator):
     rademacher = torch.bernoulli(torch.full((m,n), 0.5), generator = generator).mul_(2).sub_(1)
     return rademacher.mul_(1 / math.sqrt(m))
@@ -37,11 +35,10 @@ class SubspaceNewton(Transform):
         sketch_size (int):
             size of the random sketch. This many hessian-vector products will need to be evaluated each step.
         sketch_type (str, optional):
+            - "common_directions" - uses history steepest descent directions as the basis[2]. It is orthonormalized on-line using Gram-Schmidt (default).
             - "orthonormal" - random orthonormal basis. Orthonormality is necessary to use linear operator based modules such as trust region, but it can be slower to compute.
-            - "rademacher" - approximately orthonormal scaled random rademacher basis.
-            - "gaussian" - random gaussian (not orthonormal) basis.
-            - "common_directions" - uses history steepest descent directions as the basis[2]. It is orthonormalized on-line using Gram-Schmidt.
-            - "mixed" - random orthonormal basis but with four directions set to gradient, slow and fast gradient EMAs, and previous update direction (default).
+            - "rademacher" - approximately orthonormal (if dimension is large) scaled random rademacher basis. It is recommended to use at least "orthonormal" - it requires QR but it is still very cheap.
+            - "mixed" - random orthonormal basis but with four directions set to gradient, slow and fast gradient EMAs, and previous update direction.
         damping (float, optional): hessian damping (scale of identity matrix added to hessian). Defaults to 0.
         hvp_method (str, optional):
             How to compute hessian-matrix product:
@@ -73,7 +70,7 @@ class SubspaceNewton(Transform):
 
     RSN with line search
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.RSN(),
         tz.m.Backtracking()
@@ -82,7 +79,7 @@ class SubspaceNewton(Transform):
 
     RSN with trust region
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LevenbergMarquardt(tz.m.RSN()),
     )
@@ -97,14 +94,14 @@ class SubspaceNewton(Transform):
     def __init__(
         self,
         sketch_size: int,
-        sketch_type: Literal["orthonormal", "gaussian", "common_directions", "mixed"] = "mixed",
+        sketch_type: Literal["orthonormal", "common_directions", "mixed", "rademacher"] = "common_directions",
         damping:float=0,
+        eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        update_freq: int = 1,
+        precompute_inverse: bool = False,
+        use_lstsq: bool = True,
         hvp_method: HVPMethod = "batched_autograd",
         h: float = 1e-2,
-        use_lstsq: bool = True,
-        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,
         seed: int | None = None,
         inner: Chainable | None = None,
     ):
@@ -128,10 +125,7 @@ class SubspaceNewton(Transform):
         sketch_type = fs["sketch_type"]
         hvp_method = fs["hvp_method"]
 
-        if sketch_type in ('normal', 'gaussian'):
-            S = _gaussian_sketch(ndim, sketch_size, device=device, dtype=dtype, generator=generator)
-
-        elif sketch_type == "rademacher":
+        if sketch_type == "rademacher":
             S = _rademacher_sketch(ndim, sketch_size, device=device, dtype=dtype, generator=generator)
 
         elif sketch_type == 'orthonormal':
@@ -187,7 +181,7 @@ class SubspaceNewton(Transform):
             # form and orthogonalize sketching matrix
             S = torch.stack([g, slow_ema, fast_ema, prev_dir], dim=1)
             if sketch_size > 4:
-                S_random = _gaussian_sketch(ndim, sketch_size - 3, device=device, dtype=dtype, generator=generator)
+                S_random = torch.randn(ndim, sketch_size - 3, device=device, dtype=dtype, generator=generator) / math.sqrt(ndim)
                 S = torch.cat([S, S_random], dim=1)
 
             S = _qr_orthonormalize(S)
@@ -200,38 +194,41 @@ class SubspaceNewton(Transform):
                                                  hvp_method=fs["hvp_method"], h=fs["h"])
         H_sketched = S.T @ HS
 
-        self.global_state["H_sketched"] = H_sketched
+        # update state
+        _newton_update_state_(
+            state = self.global_state,
+            H = H_sketched,
+            damping = fs["damping"],
+            eigval_fn = fs["eigval_fn"],
+            precompute_inverse = fs["precompute_inverse"],
+            use_lstsq = fs["use_lstsq"]
+
+        )
+
         self.global_state["S"] = S
 
     def apply_states(self, objective, states, settings):
-        S: torch.Tensor = self.global_state["S"]
+        updates = objective.get_updates()
+        fs = settings[0]
 
-        d_proj = _newton_step(
-            objective=objective,
-            H=self.global_state["H_sketched"],
-            damping=self.defaults["damping"],
-            H_tfm=self.defaults["H_tfm"],
-            eigval_fn=self.defaults["eigval_fn"],
-            use_lstsq=self.defaults["use_lstsq"],
-            g_proj = lambda g: S.T @ g
-        )
+        S = self.global_state["S"]
+        b = torch.cat([t.ravel() for t in updates])
+        b_proj = S.T @ b
+
+        d_proj = _newton_solve(b=b_proj, state=self.global_state, use_lstsq=fs["use_lstsq"])
 
         d = S @ d_proj
-        objective.updates = vec_to_tensors(d, objective.params)
+        vec_to_tensors_(d, updates)
         return objective
 
     def get_H(self, objective=...):
-        eigval_fn = self.defaults["eigval_fn"]
-        H_sketched: torch.Tensor = self.global_state["H_sketched"]
-        S: torch.Tensor = self.global_state["S"]
-
-        if eigval_fn is not None:
-            try:
-                L, Q = torch.linalg.eigh(H_sketched) # pylint:disable=not-callable
-                L: torch.Tensor = eigval_fn(L)
-                H_sketched = Q @ L.diag_embed() @ Q.mH
+        if "H" in self.global_state:
+            H_sketched = self.global_state["H"]
 
-            except torch.linalg.LinAlgError:
-                pass
+        else:
+            L = self.global_state["L"]
+            Q = self.global_state["Q"]
+            H_sketched = Q @ L.diag_embed() @ Q.mH
 
+        S: torch.Tensor = self.global_state["S"]
         return Sketched(S, H_sketched)

torchzero/modules/smoothing/laplacian.py

@@ -74,7 +74,7 @@ class LaplacianSmoothing(TensorTransform):
 
     ```python
 
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LaplacianSmoothing(),
         tz.m.LR(1e-2),

torchzero/modules/smoothing/sampling.py

@@ -7,15 +7,14 @@ from typing import Literal, cast
 
 import torch
 
-from ...core import Chainable, Modular, Module, Objective
+from ...core import Chainable, Optimizer, Module, Objective
 from ...core.reformulation import Reformulation
 from ...utils import Distributions, NumberList, TensorList
 from ..termination import TerminationCriteriaBase, make_termination_criteria
 
 
 def _reset_except_self(objective: Objective, modules, self: Module):
-    assert objective.modular is not None
-    for m in objective.modular.flat_modules:
+    for m in modules:
         if m is not self:
             m.reset()
 

torchzero/modules/step_size/adaptive.py

@@ -8,7 +8,7 @@ import torch
 from ...core import Chainable, TensorTransform
 from ...utils import NumberList, TensorList, tofloat, unpack_dicts, unpack_states
 from ...linalg.linear_operator import ScaledIdentity
-from ..functional import epsilon_step_size
+from ..opt_utils import epsilon_step_size
 
 def _acceptable_alpha(alpha, param:torch.Tensor):
     finfo = torch.finfo(param.dtype)
@@ -16,7 +16,7 @@ def _acceptable_alpha(alpha, param:torch.Tensor):
         return False
     return True
 
-def _get_H(self: TensorTransform, var):
+def _get_scaled_identity_H(self: TensorTransform, var):
     n = sum(p.numel() for p in var.params)
     p = var.params[0]
     alpha = self.global_state.get('alpha', 1)
@@ -87,7 +87,7 @@ class PolyakStepSize(TensorTransform):
         return tensors
 
     def get_H(self, objective):
-        return _get_H(self, objective)
+        return _get_scaled_identity_H(self, objective)
 
 
 def _bb_short(s: TensorList, y: TensorList, sy, eps):
@@ -176,7 +176,7 @@ class BarzilaiBorwein(TensorTransform):
         prev_g.copy_(g)
 
     def get_H(self, objective):
-        return _get_H(self, objective)
+        return _get_scaled_identity_H(self, objective)
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -288,7 +288,7 @@ class BBStab(TensorTransform):
         prev_g.copy_(g)
 
     def get_H(self, objective):
-        return _get_H(self, objective)
+        return _get_scaled_identity_H(self, objective)
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -384,4 +384,4 @@ class AdGD(TensorTransform):
         return tensors
 
     def get_H(self, objective):
-        return _get_H(self, objective)
+        return _get_scaled_identity_H(self, objective)

torchzero/modules/step_size/lr.py

@@ -51,7 +51,7 @@ class Warmup(TensorTransform):
 
         .. code-block:: python
 
-            opt = tz.Modular(
+            opt = tz.Optimizer(
                 model.parameters(),
                 tz.m.Adam(),
                 tz.m.LR(1e-2),
@@ -90,7 +90,7 @@ class WarmupNormClip(TensorTransform):
 
         .. code-block:: python
 
-            opt = tz.Modular(
+            opt = tz.Optimizer(
                 model.parameters(),
                 tz.m.Adam(),
                 tz.m.WarmupNormClip(steps=1000)

torchzero/modules/trust_region/cubic_regularization.py

@@ -109,7 +109,7 @@ class CubicRegularization(TrustRegionBase):
 
         .. code-block:: python
 
-            opt = tz.Modular(
+            opt = tz.Optimizer(
                 model.parameters(),
                 tz.m.CubicRegularization(tz.m.Newton()),
             )

torchzero/modules/trust_region/levenberg_marquardt.py

@@ -44,7 +44,7 @@ class LevenbergMarquardt(TrustRegionBase):
     Gauss-Newton with Levenberg-Marquardt trust-region
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LevenbergMarquardt(tz.m.GaussNewton()),
     )
@@ -52,7 +52,7 @@ class LevenbergMarquardt(TrustRegionBase):
 
     LM-SR1
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LevenbergMarquardt(tz.m.SR1(inverse=False)),
     )

torchzero/modules/trust_region/trust_cg.py

@@ -47,7 +47,7 @@ class TrustCG(TrustRegionBase):
 
         .. code-block:: python
 
-            opt = tz.Modular(
+            opt = tz.Optimizer(
                 model.parameters(),
                 tz.m.TrustCG(hess_module=tz.m.SR1(inverse=False)),
             )

torchzero/modules/variance_reduction/svrg.py

@@ -8,8 +8,7 @@ from ...utils import tofloat
 
 
 def _reset_except_self(objective: Objective, modules, self: Module):
-    assert objective.modular is not None
-    for m in objective.modular.flat_modules:
+    for m in modules:
         if m is not self:
             m.reset()
 
@@ -45,7 +44,7 @@ class SVRG(Module):
     ## Examples:
     SVRG-LBFGS
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SVRG(len(dataloader)),
         tz.m.LBFGS(),
@@ -55,7 +54,7 @@ class SVRG(Module):
 
     For extra variance reduction one can use Online versions of algorithms, although it won't always help.
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SVRG(len(dataloader)),
         tz.m.Online(tz.m.LBFGS()),
@@ -64,7 +63,7 @@ class SVRG(Module):
 
     Variance reduction can also be applied to gradient estimators.
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SPSA(),
         tz.m.SVRG(100),

torchzero/modules/weight_decay/reinit.py

@@ -6,8 +6,8 @@ from ...core import Module
 from ...utils import NumberList, TensorList
 
 
-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()
 

torchzero/modules/weight_decay/weight_decay.py

@@ -33,7 +33,7 @@ class WeightDecay(TensorTransform):
 
     Adam with non-decoupled weight decay
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.WeightDecay(1e-3),
         tz.m.Adam(),
@@ -44,7 +44,7 @@ class WeightDecay(TensorTransform):
     Adam with decoupled weight decay that still scales with learning rate
     ```python
 
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.WeightDecay(1e-3),
@@ -54,7 +54,7 @@ class WeightDecay(TensorTransform):
 
     Adam with fully decoupled weight decay that doesn't scale with learning rate
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.LR(1e-3),
@@ -93,7 +93,7 @@ class RelativeWeightDecay(TensorTransform):
 
     Adam with non-decoupled relative weight decay
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.RelativeWeightDecay(1e-1),
         tz.m.Adam(),
@@ -103,7 +103,7 @@ class RelativeWeightDecay(TensorTransform):
 
     Adam with decoupled relative weight decay
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.RelativeWeightDecay(1e-1),

torchzero/modules/wrappers/optim_wrapper.py

@@ -11,7 +11,7 @@ class Wrap(Module):
     Wraps a pytorch optimizer to use it as a module.
 
     Note:
-        Custom param groups are supported only by ``set_param_groups``, settings passed to Modular will be applied to all parameters.
+        Custom param groups are supported only by ``set_param_groups``, settings passed to Optimizer will be applied to all parameters.
 
     Args:
         opt_fn (Callable[..., torch.optim.Optimizer] | torch.optim.Optimizer):
@@ -21,7 +21,7 @@ class Wrap(Module):
         **kwargs:
             Extra args to be passed to opt_fn. The function is called as ``opt_fn(parameters, *args, **kwargs)``.
         use_param_groups:
-            Whether to pass settings passed to Modular to the wrapped optimizer.
+            Whether to pass settings passed to Optimizer to the wrapped optimizer.
 
             Note that settings to the first parameter are used for all parameters,
             so if you specified per-parameter settings, they will be ignored.
@@ -32,7 +32,7 @@ class Wrap(Module):
     ```python
 
     from pytorch_optimizer import StableAdamW
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Wrap(StableAdamW, lr=1),
         tz.m.Cautious(),
@@ -83,7 +83,7 @@ class Wrap(Module):
 
                 # settings passed in `set_param_groups` are the highest priority
                 # schedulers will override defaults but not settings passed in `set_param_groups`
-                # this is consistent with how Modular does it.
+                # this is consistent with how Optimizer does it.
                 if self._custom_param_groups is not None:
                     setting = {k:v for k,v in setting if k not in self._custom_param_groups[0]}
 

torchzero/modules/zeroth_order/cd.py

@@ -29,7 +29,7 @@ class CD(Module):
             whether to use three points (three function evaluatins) to determine descent direction.
             if False, uses two points, but then ``adaptive`` can't be used. Defaults to True.
     """
-    def __init__(self, h:float=1e-3, grad:bool=True, adaptive:bool=True, index:Literal['cyclic', 'cyclic2', 'random']="cyclic2", threepoint:bool=True,):
+    def __init__(self, h:float=1e-3, grad:bool=False, adaptive:bool=True, index:Literal['cyclic', 'cyclic2', 'random']="cyclic2", threepoint:bool=True,):
         defaults = dict(h=h, grad=grad, adaptive=adaptive, index=index, threepoint=threepoint)
         super().__init__(defaults)