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

torchzero/modules/experimental/eigen_sr1.py

@@ -1,182 +0,0 @@
-import torch
-
-from ...core import Transform
-from ...linalg.orthogonalize import orthogonalize, OrthogonalizeMethod
-from ...linalg.eigh import eigh_plus_uuT, regularize_eigh
-from ...utils import TensorList, unpack_states, vec_to_tensors_
-from ..opt_utils import safe_clip
-from .eigengrad import _eigengrad_update_state_, eigengrad_apply
-
-
-def sr1_u(L: torch.Tensor, Q: torch.Tensor, s:torch.Tensor, y: torch.Tensor, tol:float):
-    """u from u u^T correction and its sign"""
-    r = y - torch.linalg.multi_dot([Q, L.diag_embed(), Q.T, s]) # pylint:disable=not-callable
-    rs = r.dot(s)
-
-    if rs.abs() < tol * torch.linalg.vector_norm(r) * torch.linalg.vector_norm(s): # pylint:disable=not-callable
-        return None, None
-
-    u = r / rs.abs().sqrt()
-    return u, torch.sign(rs)
-
-class EigenSR1(Transform):
-    def __init__(
-        self,
-        rank: int = 100,
-        tol: float = 1e-32,
-        eig_tol: float | None = None,
-        damping: float = 0,
-        rdamping: float = 0,
-        abs: bool = False,
-        mm_tol: float = 1e-7,
-        mm_truncate: int | None = None,
-        mm_damping: float = 1e-4,
-        mm_rdamping: float = 0,
-        mm_abs: bool = True,
-        id_reg: float | None = None,
-        column_space_tol=1e-9,
-        beta: float = 0.95,
-        balance_tol: float = 10,
-        balance_strength: float = 1e-1,
-
-        eigenbasis_optimizer = None,
-        update_freq: int = 1,
-        init_steps: int = 10,
-        orthogonalize_interval: int | None = 1,
-        orthogonalize_method: OrthogonalizeMethod = 'qr',
-
-        hvp_method = "autograd",
-        h = 1e-3,
-        inner = None,
-
-    ):
-        defaults = locals().copy()
-        for k in ["self", "inner"]:
-            del defaults[k]
-
-        super().__init__(defaults)
-
-    def update_states(self, objective, states, settings):
-        fs = settings[0]
-        step = self.increment_counter("step", 0)
-
-        if step % fs["update_freq"] == 0:
-
-            params = TensorList(objective.params)
-
-            # compute y as hessian-vector product with s (random vecs during init steps)
-            if ("p_prev" not in self.global_state) or (step < fs["init_steps"]):
-                s_list = params.sample_like('rademacher')
-
-            else:
-                p_prev = self.global_state["p_prev"]
-                s_list = params - p_prev
-
-            if s_list.dot(s_list) < torch.finfo(s_list[0].dtype).tiny * 2:
-                s_list = params.sample_like('rademacher')
-
-            self.global_state["p_prev"] = params
-
-            # compute y as hessian-vector product with s
-            Hz_list, _ = objective.hessian_vector_product(s_list, rgrad=None, at_x0=True, hvp_method=fs["hvp_method"], h=fs["h"])
-
-            s = torch.cat([t.ravel() for t in s_list])
-            y = torch.cat([t.ravel() for t in Hz_list])
-
-            # keep track of exponential moving average of hessian diagonal and balance eigenvalues
-            if (fs["balance_strength"] != 0) and (step > fs["init_steps"]) and ("L" in self.global_state):
-
-                D = s * y # hutchinson estimator
-                exp_avg = self.global_state.get("exp_avg", None)
-
-                if exp_avg is None:
-                    exp_avg = self.global_state["exp_avg"] = D
-
-                else:
-                    exp_avg.lerp_(D, weight=1-fs["beta"])
-
-                L = self.global_state["L"]
-                L_abs = L.abs()
-                tau = L_abs.amax() / exp_avg.abs().amax()
-
-                if tau > fs["balance_tol"]:
-                    L_balanced = L_abs.pow((1 / tau) ** (1 / fs["balance_strength"])).copysign(L)
-                    self.global_state["L"] = torch.where(L_abs > 1, L_balanced, L)
-
-            # initialize L and Q on 1st step
-            if "L" not in self.global_state:
-
-                L = torch.zeros(1, dtype=s.dtype, device=s.device) # rank, rank
-                Q = torch.zeros([s.numel(), 1], dtype=s.dtype, device=s.device) # ndim, rank
-
-                u, sign = sr1_u(L=L, Q=Q, s=s, y=y, tol=0)
-                assert u is not None and sign is not None
-
-                # for uu^T u is eigenvector and u^T u is eigenvalue
-                norm = torch.linalg.vector_norm(u).clip(min=torch.finfo(u.dtype).tiny * 2) # pylint:disable=not-callable
-
-                self.global_state["L"] = self.global_state["L_reg"] = (u.dot(u).unsqueeze(0) / norm) * sign # (rank,)
-                self.global_state["Q"] = self.global_state["Q_reg"] = u.unsqueeze(-1) / norm # (m, rank)
-
-            # update hessian
-            else:
-                try:
-                    L = self.global_state["L"]
-                    Q = self.global_state["Q"]
-
-                    H_step = self.increment_counter("H_step", start=0)
-                    if H_step % fs["orthogonalize_interval"] == 0:
-                        Q = orthogonalize(Q, method=fs["orthogonalize_method"])
-
-                    u, sign = sr1_u(L=L, Q=Q, s=s, y=y, tol=fs["tol"])
-
-                    if (u is not None) and (sign is not None):
-
-                        # compute new factors
-                        L_new, Q_new = eigh_plus_uuT(L, Q, u, tol=fs["column_space_tol"], alpha=sign.item(), retry_float64=True)
-
-                        # truncate/regularize new factors (those go into the accumulator)
-                        L_new, Q_new = regularize_eigh(L=L_new, Q=Q_new, truncate=min(fs["rank"], s.numel()),
-                                                      tol=fs["eig_tol"], damping=fs["damping"], rdamping=fs["rdamping"])
-
-                        _eigengrad_update_state_(state=self.global_state, setting=fs, L_new=L_new, Q_new=Q_new)
-
-                except torch.linalg.LinAlgError:
-                    pass
-
-
-
-    def apply_states(self, objective, states, settings):
-        fs = settings[0]
-        updates = objective.get_updates()
-
-        if "eigenbasis_state" not in self.global_state:
-            self.global_state["eigenbasis_state"] = {}
-
-        step = self.global_state["step"] # starts at 0
-        if step < fs["init_steps"]:
-
-            # skip update first init_steps to let hessian kick-start
-            objective.stop = True
-            objective.skip_update = True
-            return objective
-
-        if "L_reg" not in self.global_state:
-            TensorList(updates).clip_(-0.1, 0.1)
-            return objective
-
-        dir = eigengrad_apply(
-            tensor = torch.cat([t.ravel() for t in updates]),
-            L_reg = self.global_state["L_reg"],
-            Q_reg = self.global_state["Q_reg"],
-            beta = None,
-            step = None,
-            debias = False,
-            id_reg = fs["id_reg"],
-            eigenbasis_optimizer = fs["eigenbasis_optimizer"],
-            eigenbasis_state = self.global_state["eigenbasis_state"],
-            whiten_fn = lambda x: x
-        )
-
-        vec_to_tensors_(dir, updates)
-        return objective

torchzero/modules/experimental/eigengrad.py

@@ -1,207 +0,0 @@
-# pylint: disable = non-ascii-name
-from collections.abc import Mapping
-
-import torch
-
-from ...core import Chainable, TensorTransform
-from ...linalg.eigh import eigh_plus_uuT, regularize_eigh
-from ...linalg.orthogonalize import OrthogonalizeMethod, orthogonalize
-from ...linalg.linear_operator import Eigendecomposition
-from ..adaptive.lre_optimizers import LREOptimizerBase
-
-
-def _eigengrad_update_state_(state:dict, setting: Mapping, L_new: torch.Tensor | None, Q_new:torch.Tensor | None):
-    """stores L, Q, L_reg, Q_reg and reprojects eigenbasis opt (this is also used on other eigen based modules)"""
-    if (L_new is not None) and (Q_new is not None):
-
-        # re-orthogonalize
-        orthogonalize_interval = setting["orthogonalize_interval"]
-        if orthogonalize_interval is not None:
-            Q_step = state.get("Q_step", 0)
-            state["Q_step"] = Q_step + 1
-            if Q_step % orthogonalize_interval == 0:
-                Q_new = orthogonalize(Q_new, method=setting["orthogonalize_method"])
-
-        # take absolute value (for hessian)
-        if setting.get("abs", False):
-            L_new = L_new.abs()
-
-        # store
-        state["L"] = L_new
-        state["Q"] = Q_new
-
-        # absolute value for matmul
-        if setting.get("mm_abs", False):
-            L_new = L_new.abs()
-
-        # regularize for matmul
-        # this second round of regularization is only used for preconditioning
-        # and doesn't affect the accumulator
-        L_reg_new, Q_reg_new = regularize_eigh(L=L_new, Q=Q_new,
-            truncate=setting["mm_truncate"],
-            tol=setting["mm_tol"],
-            damping=setting["mm_damping"],
-            rdamping=setting["mm_rdamping"],
-        )
-
-        # print(f'{state["L_reg"] = }, {L_reg_new = }')
-
-        # reproject eigenbasis optimizer
-        if (L_reg_new is not None) and (Q_reg_new is not None):
-            eigenbasis_optimizer: LREOptimizerBase | None = setting["eigenbasis_optimizer"]
-            if eigenbasis_optimizer is not None:
-                eigenbasis_optimizer.reproject(L_old=state["L_reg"], Q_old=state["Q_reg"], L_new=L_reg_new,
-                                                Q_new=Q_reg_new, state=state["eigenbasis_state"])
-
-            state["L_reg"] = L_reg_new
-            state["Q_reg"] = Q_reg_new
-
-
-def eigengrad_apply(
-    tensor: torch.Tensor,
-    L_reg: torch.Tensor,
-    Q_reg: torch.Tensor,
-    beta: float | None,
-    step: int | None,
-    debias: bool,
-    id_reg: float | None,
-    eigenbasis_optimizer: LREOptimizerBase | None,
-    eigenbasis_state: dict,
-
-    whiten_fn = torch.sqrt
-):
-    # debias
-    if debias:
-        assert beta is not None and step is not None
-        L_reg = L_reg / (1 - beta **step)
-
-    # step with eigenbasis optimizer
-    if eigenbasis_optimizer is not None:
-        if (id_reg is not None) and (id_reg != 0):
-            raise RuntimeError("id_reg is not compatible with eigenbasis_optimizer")
-
-        update = eigenbasis_optimizer.step(tensor.ravel(), L=L_reg, Q=Q_reg, state=eigenbasis_state)
-        return update.view_as(tensor)
-
-    # or just whiten
-    # L_reg = L_reg.clip(min=torch.finfo(L_reg.dtype).tiny * 2)
-
-    if id_reg is None or id_reg == 0:
-        G = Eigendecomposition(whiten_fn(L_reg), Q_reg, use_nystrom=False)
-        dir = G.solve(tensor.ravel())
-
-    else:
-        G = Eigendecomposition(whiten_fn(L_reg), Q_reg, use_nystrom=True)
-        dir = G.solve_plus_diag(tensor.ravel(), diag=id_reg)
-
-    return dir.view_as(tensor)
-
-class Eigengrad(TensorTransform):
-    """we can easily compute rank 1 symmetric update to a low rank eigendecomposition.
-    So this stores covariance matrix as it.
-
-
-    Args:
-        rank (int): maximum allowed rank
-        beta (float, optional): beta for covariance matrix exponential moving average. Defaults to 0.95.
-        eig_tol (float, optional):
-            removes eigenvalues this much smaller than largest eigenvalue when updating the preconditioner. Defaults to 1e-7.
-        damping (float, optional):
-            added to eigenvalues when updating the preconditioner. Defaults to 1e-8.
-        rdamping (float, optional):
-            added to eigenvalues when updating the preconditioner, relative to largest eigenvalue. Defaults to 0.
-        mm_tol (float, optional):
-            removes eigenvalues this much smaller than largest eigenvalue when computing the update. Defaults to 1e-7.
-        mm_truncate (int | None, optional):
-            uses top k eigenvalues to compute the update. Defaults to None.
-        mm_damping (float, optional):
-            added to eigenvalues when computing the update. Defaults to 1e-4.
-        mm_rdamping (float, optional):
-            added to eigenvalues when computing the update, relative to largest eigenvalue. Defaults to 0.
-        id_reg (float, optional):
-            multiplier to identity matrix added to preconditioner before computing update
-            If this value is given, solution from Nyström sketch-and-solve will be used to compute the update.
-            This value can't be too small (i.e. less than 1e-5) or the solver will be very unstable. Defaults to None.
-        column_space_tol (float, optional):
-            tolerance for deciding if new eigenvector is within column space of the covariance matrix. Defaults to 1e-9.
-        concat_params (bool, optional):
-            whether to precondition all parameters at once if True, or each separately if False. Defaults to True.
-        update_freq (int, optional): update frequency. Defaults to 1.
-        inner (Chainable | None, optional): inner modules. Defaults to None.
-
-    """
-
-    def __init__(
-        self,
-        rank: int = 100,
-        beta=0.95,
-        eig_tol: float | None = 1e-5,
-        damping: float = 0,
-        rdamping: float = 0,
-        mm_tol: float = 0,
-        mm_truncate: int | None = None,
-        mm_damping: float = 1e-4,
-        mm_rdamping: float = 0,
-        id_reg: float | None = None,
-        column_space_tol = 1e-9,
-
-        orthogonalize_interval: int | None = None,
-        orthogonalize_method: OrthogonalizeMethod = 'qr',
-
-        eigenbasis_optimizer: LREOptimizerBase | None = None,
-        concat_params: bool = True,
-        update_freq: int = 1,
-        inner: Chainable | None = None,
-    ):
-        defaults = locals().copy()
-        for k in ["self", "concat_params", "inner", "update_freq"]:
-            del defaults[k]
-
-        super().__init__(defaults, concat_params=concat_params, inner=inner, update_freq=update_freq)
-
-    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
-        state["step"] = state.get("step", 0) + 1
-        beta = setting["beta"]
-
-        if "L" not in state:
-            # for uu^T u is eigenvector and u^T u is eigenvalue
-            norm = torch.linalg.vector_norm(tensor).clip(min=torch.finfo(tensor.dtype).tiny * 2) # pylint:disable=not-callable
-
-            state["L"] = state["L_reg"] = (tensor.dot(tensor).unsqueeze(0) / norm) # (rank,)
-            state["Q"] = state["Q_reg"] = tensor.unsqueeze(-1) / norm # (m, rank)
-
-        else:
-            try:
-                L = state["L"]
-                Q = state["Q"]
-
-                # compute new factors
-                L_new, Q_new = eigh_plus_uuT(L*beta, Q, tensor, alpha=(1-beta), tol=setting["column_space_tol"], retry_float64=True)
-
-                # truncate/regularize new factors (those go into the accumulator)
-                L_new, Q_new = regularize_eigh(L=L_new, Q=Q_new, truncate=setting["rank"], tol=setting["eig_tol"],
-                                              damping=setting["damping"], rdamping=setting["rdamping"])
-
-                _eigengrad_update_state_(state=state, setting=setting, L_new=L_new, Q_new=Q_new)
-
-            except torch.linalg.LinAlgError:
-                pass
-
-    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
-        if "L_reg" not in state:
-            return tensor.clip(-0.1, 0.1)
-
-        if "eigenbasis_state" not in state:
-            state["eigenbasis_state"] = {}
-
-        return eigengrad_apply(
-            tensor = tensor,
-            L_reg = state["L_reg"],
-            Q_reg = state["Q_reg"],
-            beta = setting["beta"],
-            step = state["step"],
-            debias = True,
-            id_reg = setting["id_reg"],
-            eigenbasis_optimizer = setting["eigenbasis_optimizer"],
-            eigenbasis_state = state["eigenbasis_state"]
-        )