torchzero 0.0.1__py3-none-any.whl

torchzero/modules/regularization/noise.py

@@ -0,0 +1,77 @@
+from collections import abc
+from typing import Literal
+
+import torch
+
+from ...core import OptimizerModule
+from ...tensorlist import Distributions, TensorList, _Scalar, _ScalarSequence
+
+
+def add_noise_(
+    grads: abc.Iterable[torch.Tensor],
+    alpha: "_Scalar | _ScalarSequence" = 1e-2,
+    distribution: Distributions = "normal",
+    mode: Literal["absolute", "global", "param", "channel"] = "param",
+):
+    if not isinstance(grads, TensorList): grads = TensorList(grads)
+    if mode == 'absolute':
+        grads += grads.sample_like(alpha, distribution)
+
+    elif mode == 'global':
+        grads += grads.sample_like((grads.total_vector_norm(1)/grads.total_numel() * alpha).detach().cpu().item(), distribution)
+
+    elif mode == 'param':
+        grads += grads.sample_like(grads.abs().mean()*alpha, distribution)
+
+    elif mode == 'channel':
+        grads = grads.unbind_channels()
+        grads += grads.sample_like(grads.abs().mean()*alpha, distribution)
+
+class AddNoise(OptimizerModule):
+    """Add noise to update. By default noise magnitude is relative to the mean of each parameter.
+
+    Args:
+        alpha (float, optional): magnitude of noise. Defaults to 1e-2.
+        distribution (Distributions, optional): distribution of noise. Defaults to 'normal'.
+        mode (str, optional):
+            how to calculate noise magnitude.
+
+            - "absolute": ignores gradient magnitude and always uses `alpha` as magnitude.
+
+            - "global": multiplies `alpha` by mean of the entire gradient, as if it was a single vector.
+
+            - "param": multiplies `alpha` by mean of each individual parameter (default).
+
+            - "channel": multiplies `alpha` by mean of each channel of each parameter.
+        """
+
+    def __init__(
+        self,
+        alpha: float = 1.,
+        distribution: Distributions = "normal",
+        mode: Literal["absolute", "global", "param", "channel"] = "param",
+    ):
+        defaults = dict(alpha = alpha)
+        super().__init__(defaults)
+        self.distribution: Distributions = distribution
+        self.mode: Literal["absolute", "global", "param", "channel"] = mode
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        alpha = self.get_group_key('alpha')
+
+        add_noise_(ascent, alpha, self.distribution, self.mode)
+        return ascent
+
+class Random(OptimizerModule):
+    """uses a random vector as the update. The vector is completely random and isn't checked to be descent direction.
+    This is therefore mainly useful in combination with other modules like Sum, Multiply, etc."""
+    def __init__(self, alpha: float = 1, distribution: Distributions = "normal"):
+        defaults = dict(alpha = alpha)
+        super().__init__(defaults)
+        self.distribution: Distributions = distribution
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        alpha = self.get_group_key('alpha')
+        return ascent.sample_like(alpha, self.distribution)

torchzero/modules/regularization/normalization.py

@@ -0,0 +1,328 @@
+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
+
+@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, state, 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, state, 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, state, 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, state, 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

@@ -0,0 +1,78 @@
+"""
+⟂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, state, 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

@@ -0,0 +1,92 @@
+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, state, 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

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