torchzero 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/core/preconditioner.py

@@ -1,138 +0,0 @@
-from abc import ABC, abstractmethod
-from collections import ChainMap, defaultdict
-from collections.abc import Mapping, Sequence
-from typing import Any, overload, final
-
-import torch
-
-from .module import Module, Chainable, Vars
-from .transform import apply, Transform, Target
-from ..utils import TensorList, vec_to_tensors
-
-class Preconditioner(Transform):
-    """Abstract class for a preconditioner."""
-    def __init__(
-        self,
-        defaults: dict | None,
-        uses_grad: bool,
-        concat_params: bool = False,
-        update_freq: int = 1,
-        scale_first: bool = False,
-        inner: Chainable | None = None,
-        target: Target = "update",
-    ):
-        if defaults is None: defaults = {}
-        defaults.update(dict(__update_freq=update_freq, __concat_params=concat_params, __scale_first=scale_first))
-        super().__init__(defaults, uses_grad=uses_grad, target=target)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
-    @abstractmethod
-    def update(self, tensors: list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, states: list[dict[str, Any]], settings: Sequence[Mapping[str, Any]]):
-        """updates the preconditioner with `tensors`, any internal state should be stored using `keys`"""
-
-    @abstractmethod
-    def apply(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, states: list[dict[str, Any]], settings: Sequence[Mapping[str, Any]]) -> list[torch.Tensor]:
-        """applies preconditioner to `tensors`, any internal state should be stored using `keys`"""
-
-
-    def _tensor_wise_transform(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, vars:Vars) -> list[torch.Tensor]:
-        step = self.global_state.get('__step', 0)
-        states = [self.state[p] for p in params]
-        settings = [self.settings[p] for p in params]
-        global_settings = settings[0]
-        update_freq = global_settings['__update_freq']
-
-        scale_first = global_settings['__scale_first']
-        scale_factor = 1
-        if scale_first and step == 0:
-            # 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:
-            self.update(tensors=tensors, params=params, grads=grads, states=states, settings=settings)
-
-        # step with inner
-        if 'inner' in self.children:
-            tensors = apply(self.children['inner'], tensors=tensors, params=params, grads=grads, vars=vars)
-
-        # apply preconditioner
-        tensors = self.apply(tensors=tensors, params=params, grads=grads, states=states, settings=settings)
-
-        # scale initial step, when preconditioner might not have been applied
-        if scale_first and step == 0:
-            torch._foreach_mul_(tensors, scale_factor)
-
-        self.global_state['__step'] = step + 1
-        return tensors
-
-    def _concat_transform(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, vars:Vars) -> list[torch.Tensor]:
-        step = self.global_state.get('__step', 0)
-        tensors_vec = torch.cat([t.ravel() for t in tensors])
-        params_vec = torch.cat([p.ravel() for p in params])
-        grads_vec = [torch.cat([g.ravel() for g in grads])] if grads is not None else None
-
-        states = [self.state[params[0]]]
-        settings = [self.settings[params[0]]]
-        global_settings = settings[0]
-        update_freq = global_settings['__update_freq']
-
-        scale_first = global_settings['__scale_first']
-        scale_factor = 1
-        if scale_first and step == 0:
-            # 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:
-            self.update(tensors=[tensors_vec], params=[params_vec], grads=grads_vec, states=states, settings=settings)
-
-        # step with inner
-        if 'inner' in self.children:
-            tensors = apply(self.children['inner'], tensors=tensors, params=params, grads=grads, vars=vars)
-            tensors_vec = torch.cat([t.ravel() for t in tensors]) # have to recat
-
-        # apply preconditioner
-        tensors_vec = self.apply(tensors=[tensors_vec], params=[params_vec], grads=grads_vec, states=states, settings=settings)[0]
-
-        # scale initial step, when preconditioner might not have been applied
-        if scale_first and step == 0:
-            tensors_vec *= scale_factor
-
-        tensors = vec_to_tensors(vec=tensors_vec, reference=tensors)
-        self.global_state['__step'] = step + 1
-        return tensors
-
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        concat_params = self.settings[params[0]]['__concat_params']
-        if concat_params: return self._concat_transform(tensors, params, grads, vars)
-        return self._tensor_wise_transform(tensors, params, grads, vars)
-
-class TensorwisePreconditioner(Preconditioner, ABC):
-    @abstractmethod
-    def update_tensor(self, tensor: torch.Tensor, param:torch.Tensor, grad: torch.Tensor | None, state: dict[str, Any], settings: Mapping[str, Any]):
-        """update preconditioner with `tensor`"""
-
-    @abstractmethod
-    def apply_tensor(self, tensor: torch.Tensor, param:torch.Tensor, grad: torch.Tensor | None, state: dict[str, Any], settings: Mapping[str, Any]) -> torch.Tensor:
-        """apply preconditioner to `tensor`"""
-
-    @final
-    def update(self, tensors, params, grads, states, settings):
-        if grads is None: grads = [None]*len(tensors)
-        for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
-            self.update_tensor(t, p, g, state, setting)
-
-    @final
-    def apply(self, tensors, params, grads, states, settings):
-        preconditioned = []
-        if grads is None: grads = [None]*len(tensors)
-        for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
-            preconditioned.append(self.apply_tensor(t, p, g, state, setting))
-        return preconditioned
-

torchzero/modules/experimental/algebraic_newton.py

@@ -1,145 +0,0 @@
-import warnings
-from functools import partial
-from typing import Literal
-from collections.abc import Callable
-import torch
-import torchalgebras as ta
-
-from ...core import Chainable, apply, Module
-from ...utils import vec_to_tensors, TensorList
-from ...utils.derivatives import (
-    hessian_list_to_mat,
-    hessian_mat,
-    jacobian_and_hessian_wrt,
-)
-
-class MaxItersReached(Exception): pass
-def tropical_lstsq(
-    H: torch.Tensor,
-    g: torch.Tensor,
-    solver,
-    maxiter,
-    tol,
-    algebra,
-    verbose,
-):
-    """it can run on any algebra with add despite it saying tropical"""
-    algebra = ta.get_algebra(algebra)
-
-    x = torch.zeros_like(g, requires_grad=True)
-    best_x = x.detach().clone()
-    best_loss = float('inf')
-    opt = solver([x])
-
-    niter = 0
-    def closure(backward=True):
-        nonlocal niter, best_x, best_loss
-        if niter == maxiter: raise MaxItersReached
-        niter += 1
-
-        g_hat = algebra.mm(H, x)
-        loss = torch.nn.functional.mse_loss(g_hat, g)
-        if loss < best_loss:
-            best_x = x.detach().clone()
-            best_loss = loss.detach()
-
-        if backward:
-            opt.zero_grad()
-            loss.backward()
-        return loss
-
-    loss = None
-    prev_loss = float('inf')
-    for i in range(maxiter):
-        try:
-            loss = opt.step(closure)
-            if loss == 0: break
-            if tol is not None and prev_loss - loss < tol: break
-            prev_loss = loss
-        except MaxItersReached:
-            break
-
-    if verbose: print(f'{best_loss = } after {niter} iters')
-    return best_x.detach()
-
-def tikhonov(H: torch.Tensor, reg: float, algebra: ta.Algebra = ta.TropicalSemiring()):
-    if reg!=0:
-        I = ta.AlgebraicTensor(torch.eye(H.size(-1), dtype=H.dtype, device=H.device), algebra)
-        I = I * reg
-        H = algebra.add(H, I.data)
-    return H
-
-
-class AlgebraicNewton(Module):
-    """newton in other algebras, not that it works."""
-    def __init__(
-        self,
-        reg: float | None = None,
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
-        solver=lambda p: torch.optim.LBFGS(p, line_search_fn='strong_wolfe'),
-        maxiter=1000,
-        tol: float | None = 1e-10,
-        algebra: ta.Algebra | str = 'tropical max',
-        verbose: bool = False,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(reg=reg, hessian_method=hessian_method, vectorize=vectorize)
-        super().__init__(defaults)
-
-        self.algebra = ta.get_algebra(algebra)
-        self.lstsq_args:dict = dict(solver=solver, maxiter=maxiter, tol=tol, algebra=algebra, verbose=verbose)
-
-        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']
-        hessian_method = settings['hessian_method']
-        vectorize = settings['vectorize']
-
-        # ------------------------ calculate grad and hessian ------------------------ #
-        if hessian_method == 'autograd':
-            with torch.enable_grad():
-                loss = vars.loss = vars.loss_approx = closure(False)
-                g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
-                g_list = [t[0] for t in g_list] # remove leading dim from loss
-                vars.grad = g_list
-                H = hessian_list_to_mat(H_list)
-
-        elif hessian_method in ('func', 'autograd.functional'):
-            strat = 'forward-mode' if vectorize else 'reverse-mode'
-            with torch.enable_grad():
-                g_list = vars.get_grad(retain_graph=True)
-                H: torch.Tensor = hessian_mat(partial(closure, backward=False), params,
-                                method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
-
-        else:
-            raise ValueError(hessian_method)
-
-        # -------------------------------- inner step -------------------------------- #
-        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])
-
-        # ------------------------------- regulazition ------------------------------- #
-        if reg is not None: H = tikhonov(H, reg)
-
-        # ----------------------------------- solve ---------------------------------- #
-        tropical_update = tropical_lstsq(H, g, **self.lstsq_args)
-        # what now? w - u is not defined, it is defined for max version if u < w
-        # w = params.to_vec()
-        # w_hat = self.algebra.sub(w, tropical_update)
-        # update = w_hat - w
-        # no
-        # it makes sense to solve tropical system and sub normally
-        # the only thing is that tropical system can have no solutions
-
-        vars.update = vec_to_tensors(tropical_update, params)
-        return vars

torchzero/modules/experimental/soapy.py

@@ -1,290 +0,0 @@
-from operator import itemgetter
-
-import torch
-
-from ...core import Chainable, Transform, apply
-from ..optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
-
-@torch.no_grad
-def update_soap_covariances_(
-    grad: torch.Tensor,
-    GGs_: list[torch.Tensor | None],
-    beta: float | None,
-):
-    for i, GG in enumerate(GGs_):
-        if GG is None: continue
-
-        axes = list(range(i)) + list(range(i + 1, grad.ndim)) # this works fine with 1d params
-        if beta is None: GG.add_(torch.tensordot(grad, grad, (axes, axes))) # pyright:ignore[reportArgumentType]
-        else: GG.lerp_(torch.tensordot(grad, grad, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
-
-@torch.no_grad
-def project(tensors: torch.Tensor, Q: list[torch.Tensor | None]):
-    """
-    Projects the gradient to the eigenbases of the preconditioner.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat, dims=[[0], [0]]) # pyright:ignore[reportArgumentType]
-        else:
-            # I don't understand this part but it is in https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
-
-@torch.no_grad
-def project_back(tensors: torch.Tensor, Q: list[torch.Tensor| None]):
-    """
-    Projects the gradient back to the original space.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
-        else:
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-@torch.no_grad
-def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
-    """
-    matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m in mat:
-        if m is None: continue
-        if len(m) == 0:
-            matrix.append([])
-            continue
-        if m.dtype != torch.float:
-            original_type = m.dtype
-            original_device = m.device
-            matrix.append(m.float())
-        else:
-            float_data = True
-            matrix.append(m)
-
-    final = []
-    for m in matrix:
-        if len(m) == 0:
-            final.append([])
-            continue
-        try:
-            _, Q = torch.linalg.eigh(m+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-        except Exception:
-            _, Q = torch.linalg.eigh(m.to(torch.float64)+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-            Q = Q.to(m.dtype)
-        Q = torch.flip(Q, [1])
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-    return final
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
-@torch.no_grad
-def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using one round of power iteration
-    followed by torch.linalg.qr decomposition.
-    """
-    matrix = []
-    orth_matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m,o in zip(GG, Q_list):
-        if m is None: continue
-        assert o is not None
-
-        if len(m) == 0:
-            matrix.append([])
-            orth_matrix.append([])
-            continue
-        if m.data.dtype != torch.float:
-            original_type = m.data.dtype
-            original_device = m.data.device
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-        else:
-            float_data = True
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-
-    final = []
-    for ind, (m,o) in enumerate(zip(matrix, orth_matrix)):
-        if len(m)==0:
-            final.append([])
-            continue
-        est_eig = torch.diag(o.T @ m @ o)
-        sort_idx = torch.argsort(est_eig, descending=True)
-        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
-        o = o[:,sort_idx]
-        power_iter = m @ o
-        Q, _ = torch.linalg.qr(power_iter) # pylint:disable=not-callable
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-
-    return final, exp_avg_sq
-
-class SOAPY(Transform):
-    """SOAP but uses scaled gradient differences
-
-    new args
-
-    scale by s whether to scale gradient differences by parameter differences
-
-    y_to_ema2 whether to use gradient differences for exponential moving average too
-    """
-    def __init__(
-        self,
-        beta1: float = 0.95,
-        beta2: float = 0.95,
-        shampoo_beta: float | None = 0.95,
-        precond_freq: int = 10,
-        merge_small: bool = True,
-        max_dim: int = 2_000,
-        precondition_1d: bool = True,
-        eps: float = 1e-8,
-        decay: float | None = None,
-        alpha: float = 1,
-        bias_correction: bool = True,
-        scale_by_s: bool = True,
-        y_to_ema2: bool = False,
-    ):
-        defaults = dict(
-            beta1=beta1,
-            beta2=beta2,
-            shampoo_beta=shampoo_beta,
-            precond_freq=precond_freq,
-            merge_small=merge_small,
-            max_dim=max_dim,
-            precondition_1d=precondition_1d,
-            eps=eps,
-            decay=decay,
-            bias_correction=bias_correction,
-            alpha=alpha,
-            scale_by_s=scale_by_s,
-            y_to_ema2=y_to_ema2,
-        )
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        updates = []
-        # update preconditioners
-        for i,(p,t) in enumerate(zip(params, tensors)):
-            state = self.state[p]
-            settings = self.settings[p]
-            beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, alpha = itemgetter(
-                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'alpha')(settings)
-            scale_by_s = settings['scale_by_s']
-            y_to_ema2 = settings['y_to_ema2']
-
-            if merge_small:
-                t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
-
-            if 'g_prev' not in state:
-                state['p_prev'] = p.clone()
-                state['g_prev'] = t.clone()
-                updates.append(tensors[i].clip(-0.1,0.1))
-                continue
-
-            p_prev = state['p_prev']
-            g_prev = state['g_prev']
-            s = p - p_prev
-            y = t - g_prev
-            if scale_by_s: y /= torch.linalg.norm(s).clip(min=1e-8) # pylint:disable=not-callable
-
-            state['p_prev'].copy_(p)
-            state['g_prev'].copy_(t)
-
-            # initialize state on 1st step
-            if 'GG' not in state:
-                state["exp_avg"] = torch.zeros_like(t)
-                if y_to_ema2: state["exp_avg_sq"] = torch.ones_like(t)
-                else: state["exp_avg_sq"] = torch.zeros_like(t)
-
-                if not precondition_1d and t.ndim <= 1:
-                    state['GG'] = []
-
-                else:
-                    state['GG'] = [torch.zeros(sh, sh, dtype=t.dtype, device=t.device) if 1<sh<max_dim else None for sh in t.shape]
-
-                # either scalar parameter, 1d with precondition_1d=False, or all dims are too big.
-                if len([i is not None for i in state['GG']]) == 0:
-                    state['GG'] = None
-
-                if state['GG'] is not None:
-                    update_soap_covariances_(y, GGs_=state['GG'], beta=shampoo_beta)
-                    state['Q'] = get_orthogonal_matrix(state['GG'])
-
-                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.
-
-            # Projecting gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            z_projected = None
-            if state['GG'] is not None:
-                if y_to_ema2: z_projected = project(y, state['Q'])
-                else: z_projected = project(t, state['Q'])
-
-            # exponential moving averages
-            # this part could be foreached but I will do that at some point its not a big difference compared to preconditioning
-            exp_avg: torch.Tensor = state["exp_avg"]
-            exp_avg_sq: torch.Tensor = state["exp_avg_sq"]
-
-            exp_avg.lerp_(t, 1-beta1)
-
-            if z_projected is None:
-                if y_to_ema2: exp_avg_sq.mul_(beta2).addcmul_(y, y, value=1-beta2)
-                else: exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
-            else:
-                exp_avg_sq.mul_(beta2).addcmul_(z_projected, z_projected, value=1-beta2)
-
-            # project exponential moving averages if they are accumulated unprojected
-            exp_avg_projected = exp_avg
-            if z_projected is not None:
-                exp_avg_projected = project(exp_avg, state['Q'])
-
-            exp_avg_sq_projected = exp_avg_sq
-
-            denom = exp_avg_sq_projected.sqrt().add_(eps)
-            # print(f'{t_projected = }, {exp_avg = }, {exp_avg_projected = }, {exp_avg_sq = }, {exp_avg_sq_projected = }, {denom = }')
-
-            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
-            # to the original space
-            update = exp_avg_projected / denom
-            if z_projected is not None:
-                update = project_back(update, state["Q"])
-
-            if settings['bias_correction']:
-                bias_correction1 = 1.0 - beta1 ** (state["step"]+1)
-                bias_correction2 = 1.0 - beta2 ** (state["step"]+1)
-                update *= ((bias_correction2 ** .5) / bias_correction1) * alpha
-            elif alpha is not None:
-                update *= alpha
-
-            if merge_small:
-                update = _unmerge_small_dims(update, state['flat_sizes'], state['sort_idxs'])
-
-            updates.append(update)
-            state["step"] += 1
-
-            # Update is done after the gradient step to avoid using current gradients in the projection.
-            if state['GG'] is not None:
-                update_soap_covariances_(y, state['GG'], shampoo_beta)
-                if state['step'] % settings['precond_freq'] == 0:
-                    state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, state['GG'], state['Q'])
-
-        return updates