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

torchzero/modules/ops/higher_level.py

@@ -30,6 +30,7 @@ class EMASquared(TensorTransform):
     def __init__(self, beta:float=0.999, amsgrad=False, pow:float=2):
         defaults = dict(beta=beta,pow=pow,amsgrad=amsgrad)
         super().__init__(defaults)
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -57,7 +58,7 @@ class SqrtEMASquared(TensorTransform):
     def __init__(self, beta:float=0.999, amsgrad=False, debiased: bool = False, pow:float=2,):
         defaults = dict(beta=beta,pow=pow,amsgrad=amsgrad,debiased=debiased)
         super().__init__(defaults)
-
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -141,6 +142,8 @@ class CenteredEMASquared(TensorTransform):
     def __init__(self, beta: float = 0.99, amsgrad=False, pow:float=2):
         defaults = dict(beta=beta, amsgrad=amsgrad, pow=pow)
         super().__init__(defaults, uses_grad=False)
+        self.add_projected_keys("grad", "exp_avg")
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -175,6 +178,8 @@ class CenteredSqrtEMASquared(TensorTransform):
     def __init__(self, beta: float = 0.99, amsgrad=False, debiased: bool = False, pow:float=2):
         defaults = dict(beta=beta, amsgrad=amsgrad, debiased=debiased, pow=pow)
         super().__init__(defaults, uses_grad=False)
+        self.add_projected_keys("grad", "exp_avg")
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):

torchzero/modules/second_order/inm.py

@@ -35,6 +35,8 @@ class ImprovedNewton(Transform):
         self,
         damping: float = 0,
         eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        eigv_tol: float | None = None,
+        truncate: int | None = None,
         update_freq: int = 1,
         precompute_inverse: bool | None = None,
         use_lstsq: bool = False,
@@ -89,6 +91,8 @@ class ImprovedNewton(Transform):
             state = self.global_state,
             damping = fs["damping"],
             eigval_fn = fs["eigval_fn"],
+            eigv_tol = fs["eigv_tol"],
+            truncate = fs["truncate"],
             precompute_inverse = precompute_inverse,
             use_lstsq = fs["use_lstsq"]
         )

torchzero/modules/second_order/newton.py

@@ -7,6 +7,7 @@ from ...core import Chainable, Transform, Objective, HessianMethod
 from ...utils import vec_to_tensors_
 from ...linalg.linear_operator import Dense, DenseWithInverse, Eigendecomposition
 from ...linalg import torch_linalg
+from ...linalg.eigh import regularize_eigh
 
 def _try_lu_solve(H: torch.Tensor, g: torch.Tensor):
     try:
@@ -30,6 +31,8 @@ def _newton_update_state_(
     H: torch.Tensor,
     damping: float,
     eigval_fn: Callable | None,
+    eigv_tol: float | None,
+    truncate: int | None,
     precompute_inverse: bool,
     use_lstsq: bool,
 ):
@@ -39,10 +42,11 @@ def _newton_update_state_(
         reg = torch.eye(H.size(0), device=H.device, dtype=H.dtype).mul_(damping)
         H += reg
 
-    # if eigval_fn is given, we don't need H or H_inv, we store factors
-    if eigval_fn is not None:
+    # if any args require eigendecomp, we don't need H or H_inv, we store factors
+    if any(i is not None for i in [eigval_fn, eigv_tol, truncate]):
         L, Q = torch_linalg.eigh(H, retry_float64=True)
-        L = eigval_fn(L)
+        if eigval_fn is not None: L = eigval_fn(L)
+        L, Q = regularize_eigh(L, Q, truncate=truncate, tol=eigv_tol)
         state["L"] = L
         state["Q"] = Q
         return
@@ -216,6 +220,8 @@ class Newton(Transform):
         self,
         damping: float = 0,
         eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        eigv_tol: float | None = None,
+        truncate: int | None = None,
         update_freq: int = 1,
         precompute_inverse: bool | None = None,
         use_lstsq: bool = False,
@@ -242,6 +248,8 @@ class Newton(Transform):
             H=H,
             damping = fs["damping"],
             eigval_fn = fs["eigval_fn"],
+            eigv_tol = fs["eigv_tol"],
+            truncate = fs["truncate"],
             precompute_inverse = precompute_inverse,
             use_lstsq = fs["use_lstsq"]
         )

torchzero/modules/second_order/newton_cg.py

@@ -226,7 +226,8 @@ class NewtonCGSteihaug(Transform):
         tol: float = 1e-8,
         reg: float = 1e-8,
         solver: Literal['cg', "minres"] = 'cg',
-        adapt_tol: bool = True,
+        adapt_tol: bool = False,
+        terminate_on_tr: bool = True,
         npc_terminate: bool = False,
 
         # hvp settings
@@ -272,7 +273,6 @@ class NewtonCGSteihaug(Transform):
         npc_terminate=fs["npc_terminate"]
         miniter=fs["miniter"]
         max_history=fs["max_history"]
-        adapt_tol=fs["adapt_tol"]
 
 
         # ------------------------------- trust region ------------------------------- #
@@ -294,9 +294,13 @@ class NewtonCGSteihaug(Transform):
             finfo = torch.finfo(orig_params[0].dtype)
             if trust_radius < finfo.tiny * 2:
                 trust_radius = self.global_state['trust_radius'] = init
-                if adapt_tol:
+
+                if fs["adapt_tol"]:
                     self.global_state["tol_mul"] = self.global_state.get("tol_mul", 1) * 0.1
 
+                if fs["terminate_on_tr"]:
+                    objective.should_terminate = True
+
             elif trust_radius > finfo.max / 2:
                 trust_radius = self.global_state['trust_radius'] = init
 

torchzero/modules/second_order/nystrom.py

@@ -5,7 +5,7 @@ import torch
 
 from ...core import Chainable, Transform, HVPMethod
 from ...utils import TensorList, vec_to_tensors
-from ...linalg import nystrom_pcg, nystrom_sketch_and_solve, nystrom_approximation, cg, regularize_eigh, OrthogonalizeMethod
+from ...linalg import nystrom_pcg, nystrom_sketch_and_solve, nystrom_approximation, cg, regularize_eigh, OrthogonalizeMethod, orthogonalize
 from ...linalg.linear_operator import Eigendecomposition, ScaledIdentity
 
 class NystromSketchAndSolve(Transform):
@@ -75,7 +75,7 @@ class NystromSketchAndSolve(Transform):
     """
     def __init__(
         self,
-        rank: int,
+        rank: int = 100,
         reg: float | None = 1e-2,
         eigv_tol: float = 0,
         truncate: int | None = None,
@@ -109,17 +109,15 @@ class NystromSketchAndSolve(Transform):
 
         generator = self.get_generator(params[0].device, seed=fs['seed'])
         try:
+            Omega = torch.randn([ndim, min(fs["rank"], ndim)], device=device, dtype=dtype, generator=generator)
+            Omega = orthogonalize(Omega, fs["orthogonalize_method"])
+            HOmega = H_mm(Omega)
+
             # compute the approximation
             L, Q = nystrom_approximation(
-                A_mv=H_mv,
-                A_mm=H_mm,
-                ndim=ndim,
-                rank=min(fs["rank"], ndim),
+                Omega=Omega,
+                AOmega=HOmega,
                 eigv_tol=fs["eigv_tol"],
-                orthogonalize_method=fs["orthogonalize_method"],
-                dtype=dtype,
-                device=device,
-                generator=generator,
             )
 
             # regularize
@@ -225,7 +223,7 @@ class NystromPCG(Transform):
     """
     def __init__(
         self,
-        rank: int,
+        rank: int = 100,
         maxiter=None,
         tol=1e-8,
         reg: float = 1e-6,
@@ -260,16 +258,13 @@ class NystromPCG(Transform):
             generator = self.get_generator(device, seed=fs['seed'])
 
             try:
+                Omega = torch.randn(ndim, min(fs["rank"], ndim), device=device, dtype=dtype, generator=generator)
+                HOmega = H_mm(orthogonalize(Omega, fs["orthogonalize_method"]))
+                # compute the approximation
                 L, Q = nystrom_approximation(
-                    A_mv=None,
-                    A_mm=H_mm,
-                    ndim=ndim,
-                    rank=min(fs["rank"], ndim),
+                    Omega=Omega,
+                    AOmega=HOmega,
                     eigv_tol=fs["eigv_tol"],
-                    orthogonalize_method=fs["orthogonalize_method"],
-                    dtype=dtype,
-                    device=device,
-                    generator=generator,
                 )
 
                 self.global_state["L"] = L

torchzero/modules/second_order/rsn.py

@@ -25,9 +25,23 @@ def _orthonormal_sketch(m, n, dtype, device, generator):
     return _qr_orthonormalize(torch.randn(m, n, dtype=dtype, device=device, generator=generator))
 
 def _rademacher_sketch(m, n, dtype, device, generator):
-    rademacher = torch.bernoulli(torch.full((m,n), 0.5), generator = generator).mul_(2).sub_(1)
+    rademacher = torch.bernoulli(torch.full((m,n), 0.5, device=device, dtype=dtype), generator = generator).mul_(2).sub_(1)
     return rademacher.mul_(1 / math.sqrt(m))
 
+def _row_sketch(m, n, dtype, device, generator):
+    weights = torch.ones(m, dtype=dtype, device=device)
+    indices = torch.multinomial(weights, n, replacement=False, generator=generator)
+
+    P = torch.zeros(m, n, dtype=dtype, device=device)
+    P[indices, range(n)] = 1
+    return P
+
+def _topk_rows(grad, m, n, dtype, device, generator):
+    _, indices = torch.topk(grad.abs(), n)
+    P = torch.zeros(m, n, dtype=dtype, device=device)
+    P[indices, range(n)] = 1
+    return P
+
 class SubspaceNewton(Transform):
     """Subspace Newton. Performs a Newton step in a subspace (random or spanned by past gradients).
 
@@ -37,7 +51,9 @@ class SubspaceNewton(Transform):
         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 (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.
+            - "rows" - samples random rows.
+            - "topk" - samples top-rank rows with largest gradient magnitude.
+            - "rademacher" - approximately orthonormal (if dimension is large) scaled random rademacher basis.
             - "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):
@@ -93,13 +109,15 @@ class SubspaceNewton(Transform):
 
     def __init__(
         self,
-        sketch_size: int,
-        sketch_type: Literal["orthonormal", "common_directions", "mixed", "rademacher"] = "common_directions",
+        sketch_size: int = 100,
+        sketch_type: Literal["orthonormal", "common_directions", "mixed", "rademacher", "rows", "topk"] = "common_directions",
         damping:float=0,
         eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        eigv_tol: float | None = None,
+        truncate: int | None = None,
         update_freq: int = 1,
         precompute_inverse: bool = False,
-        use_lstsq: bool = True,
+        use_lstsq: bool = False,
         hvp_method: HVPMethod = "batched_autograd",
         h: float = 1e-2,
         seed: int | None = None,
@@ -131,6 +149,14 @@ class SubspaceNewton(Transform):
         elif sketch_type == 'orthonormal':
             S = _orthonormal_sketch(ndim, sketch_size, device=device, dtype=dtype, generator=generator)
 
+        elif sketch_type == "rows":
+            S = _row_sketch(ndim, sketch_size, device=device, dtype=dtype, generator=generator)
+
+        elif sketch_type == "topk":
+            g_list = objective.get_grads(create_graph=hvp_method in ("batched_autograd", "autograd"))
+            g = torch.cat([t.ravel() for t in g_list])
+            S = _topk_rows(g, ndim, sketch_size, device=device, dtype=dtype, generator=generator)
+
         elif sketch_type == 'common_directions':
             # Wang, Po-Wei, Ching-pei Lee, and Chih-Jen Lin. "The common-directions method for regularized empirical risk minimization." Journal of Machine Learning Research 20.58 (2019): 1-49.
             g_list = objective.get_grads(create_graph=hvp_method in ("batched_autograd", "autograd"))
@@ -189,6 +215,10 @@ class SubspaceNewton(Transform):
         else:
             raise ValueError(f'Unknown sketch_type {sketch_type}')
 
+        # print(f'{S.shape = }')
+        # I = torch.eye(S.size(1), device=S.device, dtype=S.dtype)
+        # print(f'{torch.nn.functional.mse_loss(S.T @ S, I) = }')
+
         # form sketched hessian
         HS, _ = objective.hessian_matrix_product(S, rgrad=None, at_x0=True,
                                                  hvp_method=fs["hvp_method"], h=fs["h"])
@@ -200,9 +230,10 @@ class SubspaceNewton(Transform):
             H = H_sketched,
             damping = fs["damping"],
             eigval_fn = fs["eigval_fn"],
+            eigv_tol = fs["eigv_tol"],
+            truncate = fs["truncate"],
             precompute_inverse = fs["precompute_inverse"],
             use_lstsq = fs["use_lstsq"]
-
         )
 
         self.global_state["S"] = S

torchzero/modules/trust_region/trust_region.py

@@ -1,7 +1,7 @@
 import math
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence, MutableMapping
 from functools import partial
 from typing import Any, Literal, Protocol, cast, final, overload
 
@@ -203,6 +203,7 @@ def fixed_radius(
 ) -> tuple[float, bool]:
     return init, True
 
+
 _RADIUS_KEYS = Literal['default', 'fixed']
 _RADIUS_STRATEGIES: dict[_RADIUS_KEYS, _RadiusStrategy] = {
     "default": default_radius,

torchzero/utils/benchmarks/logistic.py

@@ -5,39 +5,54 @@ import numpy as np
 import torch
 import tqdm
 
-
-def generate_correlated_logistic_data(n_samples=2000, n_features=32, n_correlated_pairs=512, correlation=0.99, seed=0):
-    """Hard logistic regression dataset with correlated features"""
+def generate_correlated_logistic_data(
+    n_samples=100_000,
+    n_features=32,
+    n_classes=10,
+    n_correlated=768,
+    correlation=0.99,
+    seed=0
+) -> tuple[np.ndarray, np.ndarray]:
+    assert n_classes >= 2
     generator = np.random.default_rng(seed)
 
-    # ------------------------------------- X ------------------------------------ #
     X = generator.standard_normal(size=(n_samples, n_features))
-    weights = generator.uniform(-2, 2, n_features)
+    weights = generator.uniform(-2, 2, size=(n_features, n_classes))
+
+    used_pairs = set()
+    n_correlated = min(n_correlated, n_features * (n_features - 1) // 2)
 
-    used_pairs = []
-    for i in range(n_correlated_pairs):
+    for _ in range(n_correlated):
         idxs = None
         while idxs is None or idxs in used_pairs:
-            idxs = tuple(generator.choice(n_features, size=2, replace=False).tolist())
+            pair = generator.choice(n_features, size=2, replace=False)
+            pair.sort()
+            idxs = tuple(pair)
 
-        used_pairs.append(idxs)
+        used_pairs.add(idxs)
         idx1, idx2 = idxs
 
         noise = generator.standard_normal(n_samples) * np.sqrt(1 - correlation**2)
         X[:, idx2] = correlation * X[:, idx1] + noise
 
         w = generator.integers(1, 51)
-        weights[idx1] = w
-        weights[idx2] = -w
+        cls = generator.integers(0, n_classes)
+        weights[idx1, cls] = w
+        weights[idx2, cls] = -w
 
-    # ---------------------------------- logits ---------------------------------- #
     logits = X @ weights
-    probabilities = 1 / (1 + np.exp(-logits))
-    y = generator.binomial(1, probabilities).astype(np.float32)
 
-    X = X - X.mean(0, keepdims=True)
-    X = X / X.std(0, keepdims=True)
-    return X, y
+    logits -= logits.max(axis=1, keepdims=True)
+    exp_logits = np.exp(logits)
+    probabilities = exp_logits / exp_logits.sum(axis=1, keepdims=True)
+
+    y_one_hot = generator.multinomial(1, pvals=probabilities)
+    y = np.argmax(y_one_hot, axis=1)
+
+    X -= X.mean(0, keepdims=True)
+    X /= X.std(0, keepdims=True)
+
+    return X, y.astype(np.int64)
 
 
 # if __name__ == '__main__':
@@ -101,7 +116,7 @@ def run_logistic_regression(X: torch.Tensor, y: torch.Tensor, opt_fn, max_steps:
         # this is for tests
         if _assert_on_evaluated_same_params:
             for p in evaluated_params:
-                assert not _tensorlist_equal(p, model.parameters()), f"evaluated same parameters on epoch {epoch}"
+                assert not _tensorlist_equal(p, model.parameters()), f"{optimizer} evaluated same parameters on epoch {epoch}"
 
             evaluated_params.append([p.clone() for p in model.parameters()])
 

torchzero/utils/optuna_tools.py

@@ -27,7 +27,7 @@ def get_momentum(trial: optuna.Trial, prefix: str, conditional: bool=True) -> li
             m = NAG(beta, dampening, lerp)
             if debiased: m = Chain(m, Debias(beta1=beta))
         else:
-            m = EMA(beta, dampening, debiased=debiased, lerp=lerp)
+            m = EMA(beta, dampening, debias=debiased, lerp=lerp)
         return [m]
     return []
 

torchzero/utils/params.py

@@ -3,7 +3,7 @@ from collections.abc import Sequence, Iterable, Mapping
 import warnings
 import torch, numpy as np
 
-
+from .torch_tools import set_storage_
 
 Params = Iterable[torch.Tensor | tuple[str, torch.Tensor] | Mapping[str, Any]]
 
@@ -147,3 +147,15 @@ def _set_update_and_grad_(
 
     return param_groups
 
+
+def _set_fake_params_(fake_params: Iterable[torch.Tensor], storage: Iterable[torch.Tensor]):
+    """sets ``fake_params`` storage to ``storage`` while they remain the same python object"""
+    for fake_p, s in zip(fake_params, storage):
+        fake_p.set_(s.view_as(s).requires_grad_()) # pyright: ignore[reportArgumentType]
+
+def _empty_fake_param_storage_(fake_params: Iterable[torch.Tensor]):
+    """sets ``fake_params`` storage to empty while they remain the same python object"""
+    for p in fake_params:
+        set_storage_(p, torch.empty(0, device=p.device, dtype=p.dtype))
+
+

torchzero/utils/tensorlist.py

@@ -330,10 +330,10 @@ class TensorList(list[torch.Tensor | Any]):
 
     def global_vector_norm(self, ord:float = 2) -> torch.Tensor:
         # return torch.linalg.vector_norm(self.to_vec(), ord = ord) # pylint:disable = not-callable
-        if ord == 1: return self.global_sum()
-        if ord % 2 == 0: return self.pow(ord).global_sum().pow(1/ord)
         if ord == torch.inf: return self.abs().global_max()
         if ord == -torch.inf: return self.abs().global_min()
+        if ord == 1: return self.abs().global_sum()
+        if ord % 2 == 0: return self.pow(ord).global_sum().pow(1/ord)
         if ord == 0: return (self != 0).global_sum().to(self[0].dtype)
 
         return self.abs().pow_(ord).global_sum().pow(1/ord)

{torchzero-0.4.1.dist-info → torchzero-0.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.4.1
+Version: 0.4.3
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 Project-URL: Homepage, https://github.com/inikishev/torchzero