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

torchzero/modules/regularization/normalization.py

@@ -1,328 +0,0 @@
-from collections import abc
-import typing
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizerModule
-
-
-def _normalize_grad_(
-    grads: abc.Iterable[torch.Tensor],
-    norm_value: float = 1,
-    ord: float = 2,
-    min: float = 0,
-    mode: typing.Literal["global", "param", "channel"] = "param",
-    min_numel=2,
-):
-    if mode in ('param', 'channel'):
-        for grad in grads:
-            if grad.numel() >= min_numel:
-                if mode == 'channel' and grad.ndim >= 2:
-                    norm = torch.linalg.vector_norm(grad, ord, dim=tuple(range(1, grad.ndim)), keepdim=True) # pylint:disable=not-callable
-                    norm[norm<=min] = 1
-                    grad /= norm / norm_value
-                else: # mode = 'param' or 1d grad
-                    norm = torch.linalg.vector_norm(grad, ord) # pylint:disable=not-callable
-                    if norm > min:
-                        grad /= norm / norm_value
-    else:
-        if not isinstance(grads, TensorList): grads = TensorList(grads)
-        norm = grads.total_vector_norm(ord)
-        if norm > min:
-            grads /= norm / norm_value # type:ignore
-
-@torch.no_grad
-def normalize_grad_(
-    params: abc.Iterable[torch.Tensor],
-    norm_value: float = 1,
-    ord: float = 2,
-    min: float = 0,
-    mode: typing.Literal["global", "param", "channel"] = "global",
-    min_numel=2,
-):
-    """Normalizes gradients of an iterable of parameters.
-
-    Args:
-        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to normalize.
-        norm_value (float, optional): value to normalize to. Defaults to 1.
-        ord (float, optional): order of the norm. Defaults to 2.
-        min (float, optional):
-            won't normalize when gradient is below this norm, you can increase this
-            to avoid amplifying extremely small gradients. Defaults to 0.
-        mode (str, optional):
-            what to normalize.
-
-            - "global": normalize the entire gradient, as if it was a single vector.
-
-            - "param": normalize each param's gradient (default).
-
-            - "channel": normalize gradient of each channel of each param.
-        min_numel (int, optional):
-            skips parameters with less than this many elements. This avoids the issue where
-            parameters that have a single element always get set to the value of 1.
-            Ignored when mode is 'global'.
-
-    Example:
-        >>> normalize_grad_(model.parameters())
-    """
-    _normalize_grad_(
-        (p.grad for p in params if p.grad is not None),
-        norm_value = norm_value,
-        ord = ord,
-        min = min,
-        mode = mode,
-        min_numel = min_numel,
-    )
-
-class Normalize(OptimizerModule):
-    """Normalizes update to the given norm value.
-
-    Args:
-        norm_value (float, optional): value to normalize to. Defaults to 1.
-        ord (float, optional): order of the norm. Defaults to 2.
-        min (float, optional):
-            won't normalize when gradient is below this norm, you can increase this
-            to avoid amplifying extremely small gradients. Defaults to 0.
-        mode (str, optional):
-            what to normalize.
-
-            - "global": normalize the entire gradient, as if it was a single vector.
-
-            - "param": normalize each param's gradient (default).
-
-            - "channel": normalize gradient of each channel of each param.
-        min_numel (int, optional):
-            skips parameters with less than this many elements. This avoids the issue where
-            parameters that have a single element always get set to the value of 1.
-            Ignored when mode is 'global'.
-    """
-    def __init__(
-        self,
-        norm_value: float = 1,
-        ord: float = 2,
-        min: float = 0,
-        mode: typing.Literal["global", "param", "channel"] = "param",
-        min_numel=2,
-    ):
-        super().__init__({})
-        self.norm_value = norm_value
-        self.ord = ord
-        self.min = min
-        self.mode: typing.Literal["global", "param", "channel"] = mode
-        self.min_numel = min_numel
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        _normalize_grad_(
-            ascent,
-            norm_value = self.norm_value,
-            ord = self.ord,
-            min = self.min,
-            mode = self.mode,
-            min_numel = self.min_numel,
-        )
-        return ascent
-
-
-def _centralize_grad_(
-    grads: abc.Iterable[torch.Tensor],
-    mode: typing.Literal["global", "param", "channel"] = "channel",
-    min_ndim=2,
-    min_numel=2,
-):
-    if mode in ('param', 'channel'):
-        if mode == 'channel': min_ndim = max(min_ndim, 2)
-        for grad in grads:
-            if grad.numel() >= min_numel and grad.ndim > min_ndim:
-                if mode == 'channel':
-                    grad -= grad.mean(dim=tuple(range(1, grad.ndim)), keepdim=True)
-                else: # mode = 'param'
-                    grad -= grad.mean()
-    else:
-        if not isinstance(grads, TensorList): grads = TensorList(grads)
-        grads -= grads.mean()
-
-@torch.no_grad
-def centralize_grad_(
-    params: abc.Iterable[torch.Tensor],
-    mode: typing.Literal["global", "param", "channel"] = "channel",
-    min_ndim=2,
-    min_numel=2,
-):
-    """Centralizes gradients of an iterable of parameters.
-
-    Args:
-        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to centralize.
-        mode (str, optional):
-            what to centralize.
-
-            - "global": centralize the entire gradient (uses mean of entire gradient).
-
-            - "param": centralize each param's gradient.
-
-            - "channel": centralize gradient of each channel of each param (default).
-        min_numel (int, optional):
-            skips parameters with less than this many elements. This avoids negating updates for
-            parameters that have a single element since subtracting mean always makes it 0.
-            Ignored when mode is 'global'.
-        min_ndim (int, optional):
-            skips parameters with less than this many dimensions.
-            bias usually has 1 dimension and you don't want to centralize it.
-            Ignored when mode is 'global'.
-
-    reference
-        *Yong, H., Huang, J., Hua, X., & Zhang, L. (2020).
-        Gradient centralization: A new optimization technique for deep neural networks.
-        In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK,
-        August 23–28, 2020, Proceedings, Part I 16 (pp. 635-652). Springer International Publishing.*
-
-    Example:
-        >>> centralize_grad_(model.parameters())
-    """
-    _centralize_grad_(
-        (p.grad for p in params if p.grad is not None),
-        mode = mode,
-        min_ndim = min_ndim,
-        min_numel = min_numel,
-    )
-
-class Centralize(OptimizerModule):
-    """Centralizes the update.
-
-    Args:
-        mode (str, optional):
-            what to centralize.
-
-            - "global": centralize the entire gradient (uses mean of entire gradient).
-
-            - "param": centralize each param's gradient.
-
-            - "channel": centralize gradient of each channel of each param (default).
-        min_numel (int, optional):
-            skips parameters with less than this many elements. This avoids negating updates for
-            parameters that have a single element since subtracting mean always makes it 0.
-            Ignored when mode is 'global'.
-        min_ndim (int, optional):
-            skips parameters with less than this many dimensions.
-            bias usually has 1 dimension and you don't want to centralize it.
-            Ignored when mode is 'global'.
-
-    reference
-        *Yong, H., Huang, J., Hua, X., & Zhang, L. (2020).
-        Gradient centralization: A new optimization technique for deep neural networks.
-        In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK,
-        August 23–28, 2020, Proceedings, Part I 16 (pp. 635-652). Springer International Publishing.*
-    """
-    def __init__(
-        self,
-        mode: typing.Literal["global", "param", "channel"] = "channel",
-        min_ndim=2,
-        min_numel=2,
-    ):
-        super().__init__({})
-        self.mode: typing.Literal["global", "param", "channel"] = mode
-        self.min_ndim = min_ndim
-        self.min_numel = min_numel
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        _centralize_grad_(
-            ascent,
-            mode = self.mode,
-            min_ndim = self.min_ndim,
-            min_numel = self.min_numel,
-        )
-        return ascent
-
-
-def clip_grad_value_(params: abc.Iterable[torch.Tensor], value:float):
-    """Clip the gradients of an iterable of parameters at specified value.
-
-    Args:
-        params (abc.Iterable[torch.Tensor]): an iterable of Tensors or a single Tensor that will have gradients clipped.
-        value (float, optional):
-            maximum allowed magnitude of the gradients.
-            The gradients are clipped in the range `[-clip_value, clip_value]`
-    """
-    TensorList(params).get_existing_grads().clamp_(-value, value)
-
-class ClipValue(OptimizerModule):
-    """Clip the update at specified value.
-
-    Args:
-    value (float, optional): maximum allowed magnitude of the gradients.
-        The gradients are clipped in the range `[-clip_value, clip_value]`
-    """
-    def __init__(self, value: float):
-        defaults = dict(value = value)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        value = self.get_group_key('value')
-        ascent.clamp_(-value, value)
-        return ascent
-
-def clip_grad_norm_(
-    params: abc.Iterable[torch.Tensor],
-    max_norm: float,
-    ord: float = 2,
-    mode: typing.Literal["global", "param", "channel"] = "param",
-):
-    """Clip the gradient norm of an iterable of parameters.
-
-    Args:
-        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to clip the norm of.
-        max_norm (float, optional): norm value to clip to.
-        ord (float, optional): order of the norm. Defaults to 2.
-        mode (str, optional):
-            what to calculate the norm over.
-
-            - "global": calculates and clips the norm of the entire gradient, as if it was a single vector.
-
-            - "param": calculates and clips each param's gradient norm (default).
-
-            - "channel": calculate and clip the norm of gradient of each channel of each param.
-
-    Example:
-        >>> clip_grad_norm_(model.parameters())
-    """
-    _normalize_grad_(
-        (p.grad for p in params if p.grad is not None),
-        norm_value = max_norm,
-        min = max_norm,
-        ord = ord,
-        mode = mode,
-    )
-
-class ClipNorm(OptimizerModule):
-    """Clip the gradient norm of an iterable of parameters.
-
-    Args:
-        max_norm (float, optional): norm value to clip to.
-        ord (float, optional): order of the norm. Defaults to 2.
-        mode (str, optional):
-            what to calculate the norm over.
-
-            - "global": calculates and clips the norm of the entire gradient, as if it was a single vector.
-
-            - "param": calculates and clips each param's gradient norm (default).
-
-            - "channel": calculate and clip the norm of gradient of each channel of each param.
-    """
-    def __init__(self, max_norm: float, ord:float=2, mode: typing.Literal["global", "param", "channel"] = "param",):
-        super().__init__({})
-        self.max_norm = max_norm
-        self.ord = ord
-        self.mode: typing.Literal["global", "param", "channel"] = mode
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        _normalize_grad_(
-            ascent,
-            norm_value = self.max_norm,
-            min = self.max_norm,
-            ord = self.ord,
-            mode = self.mode,
-        )
-        return ascent

torchzero/modules/regularization/ortho_grad.py

@@ -1,78 +0,0 @@
-"""
-⟂Grad (read “ortho-grad”) was proposed in https://arxiv.org/abs/2501.04697.
-
-"""
-from collections.abc import Iterable
-
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizerModule, _Targets
-
-
-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.
-
-    Args:
-        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to apply ⟂Grad to.
-        eps (float, optional): epsilon added to the denominator for numerical stability (default: 1e-30)
-
-    reference
-        https://arxiv.org/abs/2501.04697
-    """
-    if not isinstance(params, TensorList): params = TensorList(params)
-    params = params.with_grad()
-    grad = params.grad
-    grad -= (((params*grad).total_sum())/(params*params).total_sum() + eps) * params
-
-class OrthoGrad(OptimizerModule):
-    """⟂Grad - projects gradient of an iterable of parameters to be orthogonal to the weights.
-
-    Args:
-        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to apply ⟂Grad to.
-        eps (float, optional): epsilon added to the denominator for numerical stability (default: 1e-30)
-        renormalize (bool, optional): whether to renormalize gradients back to original norm (default: True).
-        sqrt_scale (bool, optional):
-            uses square root of the scale to make it more impactful, experimental setting and doesn't really work (default: False).
-        add (bool, optional):
-            Experimental option that changes subtraction to addition.
-            I don't think it has any geometric meaning but it drives weights towards zero instead of away from it.
-            and it seems to work well with sqrt_scale = True. It speeds up convergence by a lot compared to using vanilla gradient,
-            but also has INSANE overfitting.
-        target (str, optional):
-            determines what this module updates.
-
-            "ascent" - it updates the ascent (default).
-
-            "grad" - it updates the gradient (and sets `.grad` attributes to updated gradient).
-
-            "closure" - it makes a new closure that sets the updated ascent to the .`grad` attributes.
-
-    reference
-        https://arxiv.org/abs/2501.04697
-    """
-    def __init__(self, eps: float = 1e-30, renormalize=True, sqrt_scale = False, add=False, target: _Targets = 'ascent'):
-        super().__init__({}, target=target)
-        self.eps = eps
-        self.add = add
-        self.renormalize = renormalize
-        self.sqrt_scale = sqrt_scale
-
-    def _update(self, vars, ascent):
-        params = self.get_params()
-
-        if self.renormalize: orig_norm = ascent.norm(2) + self.eps
-        else: orig_norm = 1
-
-        scale = (params*ascent).total_sum() / ((params*params).total_sum() + self.eps)
-        if self.sqrt_scale:
-            scale = scale.abs().sqrt() * scale.sign()
-
-        if self.add: ascent += params * scale
-        else: ascent -= params * scale
-
-        if self.renormalize:
-            ascent *= (orig_norm / ascent.norm(2))
-
-        return ascent
-

torchzero/modules/regularization/weight_decay.py

@@ -1,92 +0,0 @@
-from typing import Literal
-from collections.abc import Iterable
-
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizerModule, _Targets
-
-
-def l2_regularize_(params: Iterable[torch.Tensor], alpha: float = 1e-2):
-    """Adds L2 weight regularization term to the gradients in-place.
-
-    Args:
-        params (Iterable[torch.Tensor]): an iterable of Tensors or a single Tensor.
-        alpha (float, optional): multiplier to the regularizer. Defaults to 1e-2.
-    """
-    p = TensorList(params).with_requires_grad()
-    p.ensure_grad_()
-    p.grad.add_(p, alpha = alpha)
-
-def l1_regularize_(params: Iterable[torch.Tensor], alpha: float = 1e-2):
-    """Adds L1 weight regularization term to the gradients in-place.
-
-    Args:
-        params (Iterable[torch.Tensor]): an iterable of Tensors or a single Tensor.
-        alpha (float, optional): multiplier to the regularizer. Defaults to 1e-2.
-    """
-    p = TensorList(params).with_requires_grad()
-    p.ensure_grad_()
-    p.grad.add_(p.sign(), alpha = alpha)
-
-def weight_decay_penalty(params: Iterable[torch.Tensor], alpha: float = 1e-2, ord:float = 2):
-    """Calculate the weight decay penalty term that can be added to the loss.
-
-    Args:
-        params (Iterable[torch.Tensor]): an iterable of Tensors or a single Tensor.
-        alpha (float): multiplier to the regularizer.
-        ord (int, optional): order of the norm. Defaults to 2.
-    """
-    return TensorList(params).norm(ord) * alpha
-
-def decay_weights_(params: Iterable[torch.Tensor], alpha: float = 1e-2, ord:Literal[1, 2] = 2):
-    """Apply weight decay directly to parameters in-place.
-
-    Args:
-        params (Iterable[torch.Tensor]): an iterable of Tensors or a single Tensor to decay.
-        alpha (float): by how much to decay parameters (default: 1e-2)
-        ord (float, optional):
-            order of the penalty, 1 and 2 are currently supported (L1 and L2 regularization) (default: 2)
-    """
-    params = TensorList(params)
-    if ord == 2: params.mul_(1-alpha)
-    elif ord == 1: params.sub_(params.sign().mul_(alpha))
-    else: raise NotImplementedError(f'order {ord} is not supported')
-
-
-class WeightDecay(OptimizerModule):
-    """Adds weight decay term (L1 or L2 regularization) to the ascent direction.
-
-    Put this at the end to make it decoupled.
-
-    Args:
-        alpha (float, optional): multiplier to the regularizer (default: 1e-2)
-        ord (Literal[1, 2], optional):
-            order of the penalty, 1 and 2 are currently supported (L1 and L2 regularization).
-            Defaults to 2.
-        target (str, optional):
-            determines what this module updates.
-
-            "ascent" - it updates the ascent
-
-            "grad" - it updates the gradient (and sets `.grad` attributes to updated gradient).
-
-            "closure" - it makes a new closure that sets the updated ascent to the .`grad` attributes.
-    """
-    def __init__(self, alpha: float = 1e-2, ord:Literal[1, 2] = 2, target: _Targets = "ascent"):
-        defaults = dict(alpha = alpha)
-        super().__init__(defaults, target = target)
-        self.ord = ord
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-        alpha = self.get_group_key('alpha')
-
-        if any(i != 0 for i in alpha):
-
-            if self.ord == 1: ascent.add_(params.sign() * alpha)
-            elif self.ord == 2: ascent.add_(params * alpha)
-            else: raise NotImplementedError(f'weight descent of order {self.ord} not implemented.')
-
-        return ascent

torchzero/modules/scheduling/__init__.py

@@ -1,2 +0,0 @@
-from .lr_schedulers import LRWarmup
-from .step_size import PolyakStepSize, RandomStepSize

torchzero/modules/scheduling/lr_schedulers.py

@@ -1,131 +0,0 @@
-from collections.abc import Callable
-from functools import partial
-from typing import Any, overload, TYPE_CHECKING
-import random
-
-import torch
-from ...core import OptimizerModule
-
-
-if TYPE_CHECKING:
-    from ...optim import Modular
-
-
-# LR SCHEDULING MOVED TO LR MODULE
-
-# def _set_momentum_hook(optimizer, state, momentum):
-#     for module in optimizer.unrolled_modules:
-#         if 'momentum' in module.defaults:
-#             for g in module.param_groups:
-#                 g['momentum'] = momentum
-#         if 'beta1' in module.defaults:
-#             for g in module.param_groups:
-#                 g['beta1'] = momentum
-
-# def _add_scheduler_hook(opt: "Modular", scheduler_cls, id):
-#     """post-init hook that sets `scheduler_step_fn` to the scheduler step."""
-#     # get LR module
-#     lr_module = opt.get_lr_module()
-
-#     # get current LRScheduler module
-#     scheds = [i for i in opt.unrolled_modules if isinstance(i, LRScheduler)]
-#     scheds = [i for i in scheds if i.id == id]
-#     if len(scheds) != 1:
-#         raise RuntimeError(f"more than 1 module with id {id}: {scheds}")
-
-#     sch_module = scheds[0]
-
-#     # make a scheduler and save the step function
-#     scheduler = scheduler_cls(lr_module)
-#     sch_module.scheduler_step_fn = scheduler.step
-
-
-# class LRScheduler(OptimizerModule):
-#     """Use any pytorch lr scheduler.
-
-#     Important - the lr is applied multiplicatively and multiplies with learning rate of other modules,
-#     so usually base learning rate of the lr scheduler, such as `max_lr` for OneCycleLR, should be set to 1.
-
-#     Args:
-#         lr_scheduler (Callable[[torch.optim.Optimizer], torch.optim.lr_scheduler.LRScheduler  |  Any]):
-#             something like:
-#             .. code:: py
-#             lambda opt: OneCycleLR(opt, max_lr = 1, total_steps = 60000)
-#         update_every (int, optional):
-#             call `step` every n steps, useful for schedulers that only step once per epoch. Defaults to 1.
-#         cycle_momentum (bool, optional):
-#             enables support for cycling momentum with schedulers that support it, such as `OneCycleLR`.
-#             Unlike lr, momentum is not applied multiplicatively, but set to all other modules with
-#             `momentum` or `beta` settings. Has no effect if there are no modules that support momentum. Defaults to False.
-#         init_lr (float, optional):
-#             initial lr, I believe most lr schedulers ignore this. Defaults to 1.
-#         init_momentum (float, optional):
-#             initial init_momentum, I believe most lr schedulers ignore this.
-#             Has no effect if `cycle_momentum` is False or there are no modules that support momentum. Defaults to 0.
-#     """
-#     def __init__(
-#         self,
-#         lr_scheduler: Callable[[torch.optim.Optimizer], torch.optim.lr_scheduler.LRScheduler | Any],
-#         step_every: int = 1,
-#         cycle_momentum: bool = True,
-#     ):
-#         super().__init__({})
-#         scheduler = lr_scheduler(self.dummy_opt)
-#         self.update_every = step_every
-#         self.cycle_momentum = cycle_momentum
-
-#         self.scheduler_step_fn = scheduler.step
-#         self.cur = 0
-#         self.cur_lr = init_lr
-#         self.cur_momentum = init_momentum
-
-#         self.id = random.random()
-
-#     def step(self, vars):
-#         if self.cur % self.update_every == 0:
-#             self.scheduler_step_fn()
-#             self.cur_lr = self.dummy_opt.first_param_group['lr']
-#             self.cur_momentum = self.dummy_opt.first_param_group['momentum']
-
-#         params = self.get_params()
-#         ascent = state.maybe_use_grad_(params)
-#         ascent *= self.cur_lr
-
-#         if self.cycle_momentum:
-#             state.add_post_step_hook(partial(_set_momentum_hook, momentum = self.cur_momentum))
-
-class LRWarmup(OptimizerModule):
-    """Linear learning rate warmup.
-
-    Args:
-        n_steps (int): number of warmup steps.
-        start_lr (float, optional): initial lr. Defaults to 1e-8.
-        end_lr (float, optional): final lr. Defaults to 1.
-        delay_steps (int, optional): number of `start_lr` steps before starting the warmup. Defaults to 0.
-    """
-    def __init__(self, n_steps: int, start_lr: float = 1e-8, end_lr: float = 1, delay_steps: int = 0):
-
-        super().__init__({})
-        self.n_steps = n_steps
-        self.start_lr = start_lr
-        self.end_lr = end_lr
-        self.delay_steps = delay_steps
-
-        self.cur = 0
-
-    def _update(self, vars, ascent):
-        if self.cur < self.delay_steps:
-            if self.start_lr != 1: ascent *= self.start_lr
-
-        elif self.cur >= self.n_steps + self.delay_steps:
-            if self.end_lr != 1: ascent *= self.end_lr
-
-        else:
-            remaining = (self.n_steps - (self.cur-self.delay_steps)) / self.n_steps
-            lr = (self.start_lr * remaining) + self.end_lr * (1 - remaining)
-            ascent *= lr
-
-        self.cur += 1
-        return ascent
-
-