torchzero 0.3.8__tar.gz → 0.3.9__tar.gz

{torchzero-0.3.8 → torchzero-0.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.8
+Version: 0.3.9
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 License: MIT License

{torchzero-0.3.8 → torchzero-0.3.9}/pyproject.toml

@@ -13,7 +13,7 @@ build-backend = "setuptools.build_meta"
 name = "torchzero"
 description = "Modular optimization library for PyTorch."
 
-version = "0.3.8"
+version = "0.3.9"
 dependencies = [
   "torch",
   "numpy",

{torchzero-0.3.8 → torchzero-0.3.9}/tests/test_opts.py

@@ -745,7 +745,7 @@ SSVM = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.SSVM(1), tz.m.StrongWolfe()),
     sphere_opt=lambda p: tz.Modular(p, tz.m.SSVM(1), tz.m.StrongWolfe()),
     needs_closure=True,
-    func='rosen', steps=50, loss=0.02, merge_invariant=True,
+    func='rosen', steps=50, loss=1e-10, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 

{torchzero-0.3.8 → torchzero-0.3.9}/tests/test_tensorlist.py

@@ -835,7 +835,7 @@ def test_global_reductions(simple_tl: TensorList, global_method, vec_equiv_metho
     expected = vec_equiv_func()
 
     if isinstance(result, bool): assert result == expected
-    else: assert torch.allclose(result, expected), f"Tensors not close: {result = }, {expected = }"
+    else: assert torch.allclose(result, expected, atol=1e-4), f"Tensors not close: {result = }, {expected = }"
 
 
 def test_global_vector_norm(simple_tl: TensorList):

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/core/preconditioner.py

@@ -45,12 +45,11 @@ class Preconditioner(Transform):
         update_freq = global_settings['__update_freq']
 
         scale_first = global_settings['__scale_first']
-        scale_factor = 0
+        scale_factor = 1
         if scale_first and step == 0:
-            # initial step size guess from pytorch LBFGS was too unstable
-            # I switched to norm
-            tensors = TensorList(tensors)
-            scale_factor = tensors.abs().global_mean().clip(min=1)
+            # initial step size guess from pytorch LBFGS
+            scale_factor = 1 / TensorList(tensors).abs().global_sum().clip(min=1)
+            scale_factor = scale_factor.clip(min=torch.finfo(tensors[0].dtype).eps)
 
         # update preconditioner
         if step % update_freq == 0:
@@ -65,7 +64,7 @@ class Preconditioner(Transform):
 
         # scale initial step, when preconditioner might not have been applied
         if scale_first and step == 0:
-            torch._foreach_div_(tensors, scale_factor)
+            torch._foreach_mul_(tensors, scale_factor)
 
         self.global_state['__step'] = step + 1
         return tensors
@@ -82,10 +81,11 @@ class Preconditioner(Transform):
         update_freq = global_settings['__update_freq']
 
         scale_first = global_settings['__scale_first']
-        scale_factor = 0
+        scale_factor = 1
         if scale_first and step == 0:
-            # initial step size guess from pytorch LBFGS was too unstable
-            scale_factor = tensors_vec.abs().mean().clip(min=1)
+            # initial step size guess from pytorch LBFGS
+            scale_factor = 1 / tensors_vec.abs().sum().clip(min=1)
+            scale_factor = scale_factor.clip(min=torch.finfo(tensors_vec.dtype).eps)
 
         # update preconditioner
         if step % update_freq == 0:
@@ -101,7 +101,7 @@ class Preconditioner(Transform):
 
         # scale initial step, when preconditioner might not have been applied
         if scale_first and step == 0:
-            tensors_vec /= scale_factor
+            tensors_vec *= scale_factor
 
         tensors = vec_to_tensors(vec=tensors_vec, reference=tensors)
         self.global_state['__step'] = step + 1

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/modules/experimental/adasoap.py

@@ -220,7 +220,7 @@ class AdaSOAP(Transform):
                 state['step'] = 0
                 updates.append(tensors[i].clip(-0.1,0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
-                # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
+                # that can mess with other modules scaling
 
             # Projecting gradients to the eigenbases of Shampoo's preconditioner
             # i.e. projecting to the eigenbases of matrices in state['GG']

torchzero-0.3.9/torchzero/modules/experimental/structured_newton.py

@@ -0,0 +1,111 @@
+# idea https://arxiv.org/pdf/2212.09841
+import warnings
+from collections.abc import Callable
+from functools import partial
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, apply
+from ...utils import TensorList, vec_to_tensors
+from ...utils.derivatives import (
+    hessian_list_to_mat,
+    hessian_mat,
+    hvp,
+    hvp_fd_central,
+    hvp_fd_forward,
+    jacobian_and_hessian_wrt,
+)
+
+
+class StructuredNewton(Module):
+    """TODO
+    Args:
+        structure (str, optional): structure.
+        reg (float, optional): tikhonov regularizer value. Defaults to 1e-6.
+        hvp_method (str):
+            how to calculate hvp_method. Defaults to "autograd".
+        inner (Chainable | None, optional): inner modules. Defaults to None.
+
+    """
+    def __init__(
+        self,
+        structure: Literal[
+            "diagonal",
+            "diagonal1",
+            "diagonal_abs",
+            "tridiagonal",
+            "circulant",
+            "toeplitz",
+            "toeplitz_like",
+            "hankel",
+            "rank1",
+            "rank2", # any rank
+        ]
+        | str = "diagonal",
+        reg: float = 1e-6,
+        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        h: float = 1e-3,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(reg=reg, hvp_method=hvp_method, structure=structure, h=h)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        reg = settings['reg']
+        hvp_method = settings['hvp_method']
+        structure = settings['structure']
+        h = settings['h']
+
+        # ------------------------ calculate grad and hessian ------------------------ #
+        if hvp_method == 'autograd':
+            grad = vars.get_grad(create_graph=True)
+            def Hvp_fn1(x):
+                return hvp(params, grad, x, retain_graph=True)
+            Hvp_fn = Hvp_fn1
+
+        elif hvp_method == 'forward':
+            grad = vars.get_grad()
+            def Hvp_fn2(x):
+                return hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1]
+            Hvp_fn = Hvp_fn2
+
+        elif hvp_method == 'central':
+            grad = vars.get_grad()
+            def Hvp_fn3(x):
+                return hvp_fd_central(closure, params, x, h=h, normalize=True)[1]
+            Hvp_fn = Hvp_fn3
+
+        else: raise ValueError(hvp_method)
+
+        # -------------------------------- inner step -------------------------------- #
+        update = vars.get_update()
+        if 'inner' in self.children:
+            update = apply(self.children['inner'], update, params=params, grads=grad, vars=vars)
+
+        # hessian
+        if structure.startswith('diagonal'):
+            H = Hvp_fn([torch.ones_like(p) for p in params])
+            if structure == 'diagonal1': torch._foreach_clamp_min_(H, 1)
+            if structure == 'diagonal_abs': torch._foreach_abs_(H)
+            torch._foreach_add_(H, reg)
+            torch._foreach_div_(update, H)
+            vars.update = update
+            return vars
+
+        # hessian
+        raise NotImplementedError(structure)
+
+
+
+
+

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/modules/quasi_newton/cg.py

@@ -64,7 +64,7 @@ class ConguateGradientBase(Transform, ABC):
 # ------------------------------- Polak-Ribière ------------------------------ #
 def polak_ribiere_beta(g: TensorList, prev_g: TensorList):
     denom = prev_g.dot(prev_g)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
     return g.dot(g - prev_g) / denom
 
 class PolakRibiere(ConguateGradientBase):
@@ -76,8 +76,8 @@ class PolakRibiere(ConguateGradientBase):
         return polak_ribiere_beta(g, prev_g)
 
 # ------------------------------ Fletcher–Reeves ----------------------------- #
-def fletcher_reeves_beta(gg, prev_gg):
-    if prev_gg == 0: return 0
+def fletcher_reeves_beta(gg: torch.Tensor, prev_gg: torch.Tensor):
+    if prev_gg.abs() <= torch.finfo(gg.dtype).eps: return 0
     return gg / prev_gg
 
 class FletcherReeves(ConguateGradientBase):
@@ -98,7 +98,7 @@ class FletcherReeves(ConguateGradientBase):
 def hestenes_stiefel_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
     grad_diff = g - prev_g
     denom = prev_d.dot(grad_diff)
-    if denom == 0: return 0
+    if denom.abs() < torch.finfo(g[0].dtype).eps: return 0
     return (g.dot(grad_diff) / denom).neg()
 
 
@@ -114,7 +114,7 @@ class HestenesStiefel(ConguateGradientBase):
 # --------------------------------- Dai–Yuan --------------------------------- #
 def dai_yuan_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
     denom = prev_d.dot(g - prev_g)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
     return (g.dot(g) / denom).neg()
 
 class DaiYuan(ConguateGradientBase):
@@ -129,7 +129,7 @@ class DaiYuan(ConguateGradientBase):
 # -------------------------------- Liu-Storey -------------------------------- #
 def liu_storey_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList, ):
     denom = prev_g.dot(prev_d)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
     return g.dot(g - prev_g) / denom
 
 class LiuStorey(ConguateGradientBase):
@@ -159,7 +159,7 @@ class ConjugateDescent(Transform):
             self.global_state['denom'] = torch.tensor(0.).to(g[0])
 
         prev_gd = self.global_state.get('prev_gd', 0)
-        if prev_gd == 0: beta = 0
+        if abs(prev_gd) <= torch.finfo(g[0].dtype).eps: beta = 0
         else: beta = g.dot(g) / prev_gd
 
         # inner step
@@ -176,7 +176,7 @@ class ConjugateDescent(Transform):
 def hager_zhang_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList,):
     g_diff = g - prev_g
     denom = prev_d.dot(g_diff)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
 
     term1 = 1/denom
     # term2
@@ -198,7 +198,7 @@ class HagerZhang(ConguateGradientBase):
 def hs_dy_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
     grad_diff = g - prev_g
     denom = prev_d.dot(grad_diff)
-    if denom == 0: return 0
+    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
 
     # Dai-Yuan
     dy_beta = (g.dot(g) / denom).neg().clamp(min=0)

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/modules/quasi_newton/lbfgs.py

@@ -38,9 +38,9 @@ def lbfgs(
     if len(s_history) == 0 or y_k is None or ys_k is None:
 
         # initial step size guess modified from pytorch L-BFGS
-        scale = 1 / tensors_.abs().global_sum()
-        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
-        return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
+        scale_factor = 1 / TensorList(tensors_).abs().global_sum().clip(min=1)
+        scale_factor = scale_factor.clip(min=torch.finfo(tensors_[0].dtype).eps)
+        return tensors_.mul_(scale_factor)
 
     else:
         # 1st loop

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/modules/quasi_newton/lsr1.py

@@ -17,9 +17,9 @@ def lsr1_(
 ):
     if step == 0 or not s_history:
         # initial step size guess from pytorch
-        scale = 1 / tensors_.abs().global_sum()
-        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
-        return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
+        scale_factor = 1 / TensorList(tensors_).abs().global_sum().clip(min=1)
+        scale_factor = scale_factor.clip(min=torch.finfo(tensors_[0].dtype).eps)
+        return tensors_.mul_(scale_factor)
 
     m = len(s_history)
 
@@ -65,9 +65,10 @@ def lsr1_(
         Hx.add_(w_k, alpha=w_k.dot(tensors_) / wy) # pyright:ignore[reportArgumentType]
 
     if scale_second and step == 1:
-        scale = 1 / tensors_.abs().global_sum()
-        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
-        Hx.mul_(min(1.0, scale)) # pyright:ignore[reportArgumentType]
+        scale_factor = 1 / TensorList(tensors_).abs().global_sum().clip(min=1)
+        scale_factor = scale_factor.clip(min=torch.finfo(tensors_[0].dtype).eps)
+        Hx.mul_(scale_factor)
+
     return Hx
 
 

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/modules/quasi_newton/quasi_newton.py

@@ -122,7 +122,9 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
         step = state.get('step', 0)
 
         if settings['scale_second'] and step == 2:
-            tensor = tensor / tensor.abs().mean().clip(min=1)
+            scale_factor = 1 / tensor.abs().sum().clip(min=1)
+            scale_factor = scale_factor.clip(min=torch.finfo(tensor.dtype).eps)
+            tensor = tensor * scale_factor
 
         inverse = settings['inverse']
         if inverse:

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/modules/second_order/newton.py

@@ -1,14 +1,18 @@
 import warnings
+from collections.abc import Callable
 from functools import partial
 from typing import Literal
-from collections.abc import Callable
+
 import torch
 
-from ...core import Chainable, apply, Module
-from ...utils import vec_to_tensors, TensorList
+from ...core import Chainable, Module, apply
+from ...utils import TensorList, vec_to_tensors
 from ...utils.derivatives import (
     hessian_list_to_mat,
     hessian_mat,
+    hvp,
+    hvp_fd_central,
+    hvp_fd_forward,
     jacobian_and_hessian_wrt,
 )
 
@@ -117,9 +121,10 @@ class Newton(Module):
             raise ValueError(hessian_method)
 
         # -------------------------------- inner step -------------------------------- #
+        update = vars.get_update()
         if 'inner' in self.children:
-            g_list = apply(self.children['inner'], list(g_list), params=params, grads=list(g_list), vars=vars)
-        g = torch.cat([t.view(-1) for t in g_list])
+            update = apply(self.children['inner'], update, params=params, grads=list(g_list), vars=vars)
+        g = torch.cat([t.view(-1) for t in update])
 
         # ------------------------------- regulazition ------------------------------- #
         if eig_reg: H = eig_tikhonov_(H, reg)
@@ -139,4 +144,4 @@ class Newton(Module):
         if update is None: update = least_squares_solve(H, g)
 
         vars.update = vec_to_tensors(update, params)
-        return vars
+        return vars

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/modules/second_order/newton_cg.py

@@ -66,9 +66,9 @@ class NewtonCG(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = grad
+        b = vars.get_update()
         if 'inner' in self.children:
-            b = as_tensorlist(apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars))
+            b = as_tensorlist(apply(self.children['inner'], b, params=params, grads=grad, vars=vars))
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero/modules/second_order/nystrom.py

@@ -15,7 +15,7 @@ class NystromSketchAndSolve(Module):
         rank: int,
         reg: float = 1e-3,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
-        h=1e-3,
+        h=1e-2,
         inner: Chainable | None = None,
         seed: int | None = None,
     ):
@@ -74,9 +74,9 @@ class NystromSketchAndSolve(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = grad
+        b = vars.get_update()
         if 'inner' in self.children:
-            b = apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars)
+            b = apply(self.children['inner'], b, params=params, grads=grad, vars=vars)
 
         # ------------------------------ sketch&n&solve ------------------------------ #
         x = nystrom_sketch_and_solve(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), rank=rank, reg=reg, generator=generator)
@@ -93,7 +93,7 @@ class NystromPCG(Module):
         tol=1e-3,
         reg: float = 1e-6,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
-        h=1e-3,
+        h=1e-2,
         inner: Chainable | None = None,
         seed: int | None = None,
     ):
@@ -156,9 +156,9 @@ class NystromPCG(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = grad
+        b = vars.get_update()
         if 'inner' in self.children:
-            b = apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars)
+            b = apply(self.children['inner'], b, params=params, grads=grad, vars=vars)
 
         # ------------------------------ sketch&n&solve ------------------------------ #
         x = nystrom_pcg(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), sketch_size=sketch_size, reg=reg, tol=tol, maxiter=maxiter, x0_=None, generator=generator)

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.8
+Version: 0.3.9
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 License: MIT License

{torchzero-0.3.8 → torchzero-0.3.9}/torchzero.egg-info/SOURCES.txt

@@ -36,6 +36,7 @@ torchzero/modules/experimental/newton_solver.py
 torchzero/modules/experimental/reduce_outward_lr.py
 torchzero/modules/experimental/soapy.py
 torchzero/modules/experimental/spectral.py
+torchzero/modules/experimental/structured_newton.py
 torchzero/modules/experimental/subspace_preconditioners.py
 torchzero/modules/experimental/tropical_newton.py
 torchzero/modules/grad_approximation/__init__.py