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

torchzero/modules/adaptive/natural_gradient.py

@@ -1,12 +1,12 @@
 import torch
-from ...core import Module, Chainable, apply_transform
+from ...core import Transform
 
 from ...utils.derivatives import jacobian_wrt, flatten_jacobian
-from ...utils import vec_to_tensors, TensorList
-from ...utils.linalg import linear_operator
-from .lmadagrad import lm_adagrad_apply, lm_adagrad_update
+from ...utils import vec_to_tensors
+from ...linalg import linear_operator
+from .ggt import ggt_update
 
-class NaturalGradient(Module):
+class NaturalGradient(Transform):
     """Natural gradient approximated via empirical fisher information matrix.
 
     To use this, either pass vector of per-sample losses to the step method, or make sure
@@ -27,9 +27,9 @@ class NaturalGradient(Module):
             with a vector that isn't strictly per-sample gradients, but rather for example different losses.
         gn_grad (bool, optional):
             if True, uses Gauss-Newton G^T @ f as the gradient, which is effectively sum weighted by value
-            and is equivalent to squaring the values. This way you can solve least-squares
-            objectives with a NGD-like algorithm. If False, uses sum of per-sample gradients.
-            This has an effect when ``sqrt=True``, and affects the ``grad`` attribute.
+            and is equivalent to squaring the values. That makes the kernel trick solver incorrect, but for
+            some reason it still works. If False, uses sum of per-sample gradients.
+            This has an effect when ``sqrt=False``, and affects the ``grad`` attribute.
             Defaults to False.
         batched (bool, optional): whether to use vmapping. Defaults to True.
 
@@ -41,7 +41,7 @@ class NaturalGradient(Module):
     y = torch.randn(64, 10)
 
     model = nn.Sequential(nn.Linear(20, 64), nn.ELU(), nn.Linear(64, 10))
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.NaturalGradient(),
         tz.m.LR(3e-2)
@@ -61,7 +61,7 @@ class NaturalGradient(Module):
     y = torch.randn(64, 10)
 
     model = nn.Sequential(nn.Linear(20, 64), nn.ELU(), nn.Linear(64, 10))
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.NaturalGradient(),
         tz.m.LR(3e-2)
@@ -84,7 +84,7 @@ class NaturalGradient(Module):
         return torch.stack([(1 - x1).abs(), (10 * (x2 - x1**2).abs())])
 
     X = torch.tensor([-1.1, 2.5], requires_grad=True)
-    opt = tz.Modular([X], tz.m.NaturalGradient(sqrt=True, gn_grad=True), tz.m.LR(0.05))
+    opt = tz.Optimizer([X], tz.m.NaturalGradient(sqrt=True, gn_grad=True), tz.m.LR(0.05))
 
     for iter in range(200):
         losses = rosenbrock(X)
@@ -97,20 +97,27 @@ class NaturalGradient(Module):
         super().__init__(defaults=dict(batched=batched, reg=reg, sqrt=sqrt, gn_grad=gn_grad))
 
     @torch.no_grad
-    def update(self, var):
-        params = var.params
-        batched = self.defaults['batched']
-        gn_grad = self.defaults['gn_grad']
-
-        closure = var.closure
-        assert closure is not None
-
+    def update_states(self, objective, states, settings):
+        params = objective.params
+        closure = objective.closure
+        fs = settings[0]
+        batched = fs['batched']
+        gn_grad = fs['gn_grad']
+
+        # compute per-sample losses
+        f = objective.loss
+        if f is None:
+            assert closure is not None
+            with torch.enable_grad():
+                f = objective.get_loss(backward=False) # n_out
+                assert isinstance(f, torch.Tensor)
+
+        # compute per-sample gradients
         with torch.enable_grad():
-            f = var.get_loss(backward=False) # n_out
-            assert isinstance(f, torch.Tensor)
             G_list = jacobian_wrt([f.ravel()], params, batched=batched)
 
-        var.loss = f.sum()
+        # set scalar loss and it's grad to objective
+        objective.loss = f.sum()
         G = self.global_state["G"] = flatten_jacobian(G_list) # (n_samples, ndim)
 
         if gn_grad:
@@ -119,13 +126,15 @@ class NaturalGradient(Module):
         else:
             g = self.global_state["g"] = G.sum(0)
 
-        var.grad = vec_to_tensors(g, params)
+        objective.grads = vec_to_tensors(g, params)
 
         # set closure to calculate scalar value for line searches etc
-        if var.closure is not None:
+        if closure is not None:
+
             def ngd_closure(backward=True):
+
                 if backward:
-                    var.zero_grad()
+                    objective.zero_grad()
                     with torch.enable_grad():
                         loss = closure(False)
                         if gn_grad: loss = loss.pow(2)
@@ -137,39 +146,52 @@ class NaturalGradient(Module):
                 if gn_grad: loss = loss.pow(2)
                 return loss.sum()
 
-            var.closure = ngd_closure
+            objective.closure = ngd_closure
 
     @torch.no_grad
-    def apply(self, var):
-        params = var.params
-        reg = self.defaults['reg']
-        sqrt = self.defaults['sqrt']
+    def apply_states(self, objective, states, settings):
+        params = objective.params
+        fs = settings[0]
+        reg = fs['reg']
+        sqrt = fs['sqrt']
 
         G: torch.Tensor = self.global_state['G'] # (n_samples, n_dim)
 
         if sqrt:
             # this computes U, S <- SVD(M), then calculate update as U S^-1 Uᵀg,
             # but it computes it through eigendecompotision
-            U, L = lm_adagrad_update(G.H, reg, 0)
-            if U is None or L is None: return var
+            L, U = ggt_update(G.H, damping=reg, rdamping=1e-16, truncate=0, eig_tol=1e-12)
+
+            if U is None or L is None:
+
+                # fallback to element-wise
+                g = self.global_state["g"]
+                g /= G.square().mean(0).sqrt().add(reg)
+                objective.updates = vec_to_tensors(g, params)
+                return objective
 
-            v = lm_adagrad_apply(self.global_state["g"], U, L)
-            var.update = vec_to_tensors(v, params)
-            return var
+            # whiten
+            z = U.T @ self.global_state["g"]
+            v = (U * L.rsqrt()) @ z
+            objective.updates = vec_to_tensors(v, params)
+            return objective
 
-        GGT = G @ G.H # (n_samples, n_samples)
+        # we need (G^T G)v = g
+        # where g = G^T
+        # so we need to solve (G^T G)v = G^T
+        GGt = G @ G.H # (n_samples, n_samples)
 
         if reg != 0:
-            GGT.add_(torch.eye(GGT.size(0), device=GGT.device, dtype=GGT.dtype).mul_(reg))
+            GGt.add_(torch.eye(GGt.size(0), device=GGt.device, dtype=GGt.dtype).mul_(reg))
 
-        z, _ = torch.linalg.solve_ex(GGT, torch.ones_like(GGT[0])) # pylint:disable=not-callable
+        z, _ = torch.linalg.solve_ex(GGt, torch.ones_like(GGt[0])) # pylint:disable=not-callable
         v = G.H @ z
 
-        var.update = vec_to_tensors(v, params)
-        return var
+        objective.updates = vec_to_tensors(v, params)
+        return objective
 
 
-    def get_H(self, var):
+    def get_H(self, objective=...):
         if "G" not in self.global_state: return linear_operator.ScaledIdentity()
         G = self.global_state['G']
         return linear_operator.AtA(G)

torchzero/modules/adaptive/orthograd.py

@@ -1,13 +1,9 @@
-from operator import itemgetter
-import math
-import warnings
-from collections.abc import Iterable, Sequence
-from typing import Literal
+from collections.abc import Iterable
 
 import torch
 
-from ...core import Target, Transform
-from ...utils import as_tensorlist
+from ...core import  TensorTransform
+from ...utils import TensorList
 
 def orthograd_(params: Iterable[torch.Tensor], eps: float = 1e-30):
     """Applies ⟂Grad - projects gradient of an iterable of parameters to be orthogonal to the weights.
@@ -19,29 +15,29 @@ def orthograd_(params: Iterable[torch.Tensor], eps: float = 1e-30):
     reference
         https://arxiv.org/abs/2501.04697
     """
-    params = as_tensorlist(params).with_grad()
+    params = TensorList(params).with_grad()
     grad = params.grad
     grad -= (params.dot(grad)/(params.dot(params) + eps)) * params
 
 
-class OrthoGrad(Transform):
+class OrthoGrad(TensorTransform):
     """Applies ⟂Grad - projects gradient of an iterable of parameters to be orthogonal to the weights.
 
     Args:
         eps (float, optional): epsilon added to the denominator for numerical stability (default: 1e-30)
         renormalize (bool, optional): whether to graft projected gradient to original gradient norm. Defaults to True.
-        target (Target, optional): what to set on var. Defaults to 'update'.
     """
-    def __init__(self, eps: float = 1e-8, renormalize=True, target: Target = 'update'):
+    def __init__(self, eps: float = 1e-8, renormalize=True):
         defaults = dict(eps=eps, renormalize=renormalize)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    @torch.no_grad
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         eps = settings[0]['eps']
         renormalize = settings[0]['renormalize']
 
-        params = as_tensorlist(params)
-        target = as_tensorlist(tensors)
+        params = TensorList(params)
+        target = TensorList(tensors)
 
         scale = params.dot(target)/(params.dot(params) + eps)
         if renormalize:

torchzero/modules/adaptive/psgd/__init__.py

@@ -0,0 +1,5 @@
+from .psgd_dense_newton import PSGDDenseNewton
+from .psgd_kron_newton import PSGDKronNewton
+from .psgd_kron_whiten import PSGDKronWhiten
+from .psgd_lra_newton import PSGDLRANewton
+from .psgd_lra_whiten import PSGDLRAWhiten

torchzero/modules/adaptive/psgd/_psgd_utils.py

@@ -0,0 +1,37 @@
+# pylint:disable=not-callable
+import warnings
+
+import torch
+
+from .psgd import lift2single
+
+
+def _initialize_lra_state_(tensor: torch.Tensor, state, setting):
+    n = tensor.numel()
+    rank = max(min(setting["rank"], n-1), 1)
+    dtype=tensor.dtype
+    device=tensor.device
+
+    U = torch.randn((n, rank), dtype=dtype, device=device)
+    U *= 0.1**0.5 / torch.linalg.vector_norm(U)
+
+    V = torch.randn((n, rank), dtype=dtype, device=device)
+    V *= 0.1**0.5 / torch.linalg.vector_norm(V)
+
+    if setting["init_scale"] is None:
+        # warnings.warn("FYI: Will set the preconditioner initial scale on the fly. Recommend to set it manually.")
+        d = None
+    else:
+        d = torch.ones(n, 1, dtype=dtype, device=device) * setting["init_scale"]
+
+    state["UVd"] = [U, V, d]
+    state["Luvd"] = [lift2single(torch.zeros([], dtype=dtype, device=device)) for _ in range(3)]
+
+
+
+def _wrap_with_no_backward(opt):
+    """to use original psgd opts with visualbench"""
+    class _Wrapped:
+        def step(self, closure):
+            return opt.step(lambda: closure(False))
+    return _Wrapped()