torchzero 0.1.8__py3-none-any.whl → 0.3.1__py3-none-any.whl

torchzero/core/preconditioner.py

@@ -0,0 +1,137 @@
+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 = 0
+        if scale_first and step == 0:
+            # initial step size guess from pytorch LBFGS
+            scale_factor = TensorList(tensors).abs().sum()
+
+        # 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_div_(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 = 0
+        if scale_first and step == 0:
+            # initial step size guess from pytorch LBFGS
+            scale_factor = tensors_vec.abs().sum()
+
+        # 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:
+            if scale_factor >= torch.finfo(tensors_vec.dtype).eps:
+                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/core/transform.py

@@ -0,0 +1,252 @@
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
+from typing import Any, Literal
+
+import torch
+
+from ..utils import set_storage_
+from .module import Module, Vars, Chain, Chainable
+
+Target = Literal['grad', 'update', 'closure', 'params_direct', 'params_difference', 'update_difference']
+
+class Transform(Module, ABC):
+    """Base class for a transform.
+
+    This is an abstract class, to use it, subclass it and override `transform`.
+
+    Args:
+        defaults (dict[str,Any] | None): dict with default values.
+        uses_grad (bool):
+            Set this to True if `transform` method uses the `grad` argument. This will ensure
+            `grad` is always computed and can't be None. Otherwise set to False.
+        target (Target, optional):
+            what to set on vars. Defaults to 'update'.
+    """
+    def __init__(self, defaults: dict[str,Any] | None, uses_grad: bool, target: Target = 'update'):
+        super().__init__(defaults)
+        self._target: Target = target
+        self._uses_grad = uses_grad
+
+    @abstractmethod
+    def transform(self, tensors: list[torch.Tensor], params: list[torch.Tensor], grads: list[torch.Tensor] | None, vars: Vars) -> Iterable[torch.Tensor]:
+        """applies the update rule to `target`."""
+
+    def step(self, vars: Vars) -> Vars:
+        # vars may change, therefore current params and grads have to be extracted and passed explicitly
+        if self._uses_grad: vars.get_grad()
+        params=vars.params; grad = vars.grad
+
+        # ---------------------------------- update ---------------------------------- #
+        if self._target == 'update':
+            vars.update = list(self.transform(vars.get_update(), params, grad, vars))
+            return vars
+
+        # ----------------------------------- grad ----------------------------------- #
+        if self._target == 'grad':
+            vars.grad = list(self.transform(vars.get_grad(), params, grad, vars))
+            return vars
+
+        # ------------------------------- params_direct ------------------------------ #
+        if self._target == 'params_direct':
+            new_params = self.transform(vars.params, params, grad, vars)
+            for p, new_p in zip(vars.params, new_params): set_storage_(p, new_p)
+            return vars
+
+        # ----------------------------- params_differnce ----------------------------- #
+        if self._target == 'params_difference':
+            new_params = tuple(self.transform([p.clone() for p in vars.params], params, grad, vars))
+            vars.update = list(torch._foreach_sub(vars.params, new_params))
+            return vars
+
+        # ----------------------------- update_difference ---------------------------- #
+        if self._target == 'update_difference':
+            update = vars.get_update()
+            new_update = tuple(self.transform([u.clone() for u in update], params, grad, vars))
+            vars.update = list(torch._foreach_sub(update, new_update))
+            return vars
+
+        # ---------------------------------- closure --------------------------------- #
+        if self._target == 'closure':
+            original_closure = vars.closure
+            if original_closure is None: raise ValueError('Target = "closure", but closure is None')
+
+            params = vars.params
+            def transformed_closure(backward=True):
+                if backward:
+                    loss = original_closure()
+                    current_grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+                    transformed_grad = list(self.transform(current_grad, params, grad, vars))
+                    for p, g in zip(params, transformed_grad):
+                        p.grad = g
+
+                else:
+                    loss = original_closure(False)
+
+                return loss
+
+            vars.closure = transformed_closure
+            return vars
+
+        # ---------------------------------- invalid --------------------------------- #
+        raise ValueError(f'Invalid target: {self._target}')
+
+
+class TensorwiseTransform(Module, ABC):
+    """Base class for a parameter-wise transform.
+
+    This is an abstract class, to use it, subclass it and override `transform`.
+
+    Args:
+        defaults (dict[str,Any] | None): dict with default values.
+        uses_grad (bool):
+            Set this to True if `transform` method uses the `grad` argument. This will ensure
+            `grad` is always computed and can't be None. Otherwise set to False.
+        target (Target, optional):
+            what to set on vars. Defaults to 'update'.
+    """
+    def __init__(self, defaults: dict[str,Any] | None, uses_grad: bool, target: Target = 'update'):
+        super().__init__(defaults)
+        self._target: Target = target
+        self._uses_grad: bool = uses_grad
+
+    @abstractmethod
+    def transform(
+        self,
+        tensor: torch.Tensor,
+        param: torch.Tensor,
+        grad: torch.Tensor | None,
+        vars: Vars,
+    ) -> torch.Tensor:
+        """applies the update rule to `target`"""
+
+    def step(self, vars: Vars) -> Vars:
+        params = vars.params
+        if self._uses_grad and vars.grad is None: vars.get_grad()
+
+        # ---------------------------------- update ---------------------------------- #
+        if self._target == 'update':
+            update = vars.get_update()
+            grad = vars.grad if vars.grad is not None else [None] * len(params)
+            transformed_update = []
+
+            for p, g, u in zip(params, grad, update):
+                # settings = self.settings[p] # couldn't make typing work with this
+                #, self.transform(target=u, param=p, grad=g, vars=vars, **{k:settings[k] for k in self.defaults})
+                transformed_update.append(self.transform(tensor=u, param=p, grad=g, vars=vars))
+
+            vars.update = transformed_update
+            return vars
+
+        # ----------------------------------- grad ----------------------------------- #
+        if self._target == 'grad':
+            grad = vars.get_grad()
+            transformed_grad = []
+
+            for p, g in zip(params, grad):
+                transformed_grad.append(self.transform(tensor=g, param=p, grad=g, vars=vars))
+
+            vars.grad = transformed_grad
+            return vars
+
+        # ------------------------------- params_direct ------------------------------ #
+        if self._target == 'params_direct':
+            grad = vars.grad if vars.grad is not None else [None] * len(params)
+
+            for p, g in zip(params, grad):
+                set_storage_(p, self.transform(tensor=p, param=p, grad=g, vars=vars))
+
+            return vars
+
+        # ----------------------------- params_difference ---------------------------- #
+        if self._target == 'params_difference':
+            grad = vars.grad if vars.grad is not None else [None] * len(params)
+            transformed_params = []
+
+            for p, g in zip(params, grad):
+                transformed_params.append(
+                    self.transform(tensor=p.clone(), param=p, grad=g, vars=vars)
+                )
+
+            vars.update = list(torch._foreach_sub(params, transformed_params))
+            return vars
+
+        # ----------------------------- update_difference ---------------------------- #
+        if self._target == 'update_difference':
+            update = vars.get_update()
+            grad = vars.grad if vars.grad is not None else [None] * len(params)
+            transformed_update = []
+
+            for p, g, u in zip(params, grad, update):
+                transformed_update.append(
+                    self.transform(tensor=u.clone(), param=p, grad=g, vars=vars)
+                )
+
+            vars.update = list(torch._foreach_sub(update, transformed_update))
+            return vars
+
+        # ---------------------------------- closure --------------------------------- #
+        if self._target == 'closure':
+            original_closure = vars.closure
+            if original_closure is None: raise ValueError('Target = "closure", but closure is None')
+
+            params = vars.params
+            def transformed_closure(backward=True):
+                if backward:
+                    loss = original_closure()
+                    grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+                    transformed_grad = []
+
+                    for p, g in zip(params, grad):
+                        transformed_grad.append(self.transform(tensor=g, param=p, grad=g, vars=vars))
+
+                    for p, g in zip(params, transformed_grad):
+                        p.grad = g
+
+                else:
+                    loss = original_closure(False)
+
+                return loss
+
+            vars.closure = transformed_closure
+            return vars
+
+        # ---------------------------------- invalid --------------------------------- #
+        raise ValueError(f'Invalid target: {self._target}')
+
+
+
+def apply(
+    tfm: Chainable,
+    tensors: list[torch.Tensor],
+    params: list[torch.Tensor],
+    grads: list[torch.Tensor] | None,
+    vars: Vars | None = None,
+    current_step: int = 0,
+):
+    if vars is None: vars = Vars(params=params, closure=None, model=None, current_step=current_step)
+    if isinstance(tfm, Transform):
+        if tfm._uses_grad and grads is None: grads = vars.get_grad()
+        return list(tfm.transform(tensors, params, grads, vars))
+
+    if isinstance(tfm, TensorwiseTransform):
+        grads_list = grads
+        if grads_list is None:
+            if tfm._uses_grad: grads_list = vars.get_grad()
+            else: grads_list = [None] * len(tensors)
+        return [tfm.transform(t, p, g, vars) for t,p,g in zip(tensors,params,grads_list)]
+
+    if isinstance(tfm, Chain): tfm = tfm.get_children_sequence() # pyright: ignore[reportAssignmentType]
+    if isinstance(tfm, Sequence):
+        for module in tfm:
+            tensors = apply(module, tensors=tensors, params=params, grads=grads, vars=vars)
+        return tensors
+
+    if isinstance(tfm, Module):
+        cvars = vars.clone(clone_update=False)
+        cvars.update = tensors
+        cvars = tfm.step(cvars)
+        vars.update_attrs_from_clone_(cvars)
+        assert cvars.update is not None
+        return cvars.update
+
+    raise TypeError(type(tfm))

torchzero/modules/__init__.py

@@ -1,21 +1,13 @@
-r"""
-This submodule contains composable optimizer "building blocks".
-"""
-
-from ..core.module import OptimizerModule
-from . import experimental
-from .adaptive import *
-from .gradient_approximation import *
-from .line_search import *
-from .meta import *
-from .misc import *
-from .momentum import *
-from .operations import *
-from .optimizers import *
-from .orthogonalization import *
-from .quasi_newton import *
-from .regularization import *
-from .scheduling import *
-from .second_order import *
-from .smoothing import *
-from .weight_averaging import *
+from .clipping import *
+from .grad_approximation import *
+from .line_search import *
+from .lr import *
+from .momentum import *
+from .ops import *
+from .optimizers import *
+from .projections import *
+from .quasi_newton import *
+from .smoothing import *
+from .weight_decay import *
+from .wrappers import *
+from .second_order import *

torchzero/modules/clipping/__init__.py

@@ -0,0 +1,3 @@
+from .clipping import ClipValue, ClipNorm, Normalize, clip_grad_norm_, clip_grad_value_, normalize_grads_, Centralize
+from .growth_clipping import ClipNormGrowth, ClipValueGrowth
+from .ema_clipping import ClipNormByEMA, NormalizeByEMA, ClipValueByEMA