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

torchzero/core/tensorlist_optimizer.py

@@ -1,219 +0,0 @@
-from typing import Literal, Any, overload, TypeVar
-from abc import ABC
-from collections.abc import Callable, Sequence, Iterable, Mapping, MutableSequence
-import numpy as np
-import torch
-import torch.optim.optimizer
-from torch.optim.optimizer import ParamsT
-
-from ..tensorlist import TensorList, NumberList
-from ..utils.torch_tools import totensor, tofloat
-from ..utils.python_tools import _ScalarLoss
-
-_StateInit = Literal['params', 'grad'] | Callable | TensorList
-
-_ClosureType = Callable[..., _ScalarLoss]
-"""
-
-Closure example:
-
-.. code-block:: python
-
-    def closure(backward = True):
-        loss = model(inputs)
-        if backward:
-            optimizer.zero_grad()
-            loss.backward()
-        return loss
-
-This closure will also work with all built in pytorch optimizers including LBFGS, as well as and most custom ones.
-"""
-
-def _maybe_pass_backward(closure: _ClosureType, backward: bool) -> _ScalarLoss:
-    """not passing backward when it is true makes this work with closures with no `backward` argument"""
-    if backward:
-        with torch.enable_grad(): return closure()
-    return closure(False)
-
-CLS = TypeVar('CLS')
-class TensorListOptimizer(torch.optim.Optimizer, ABC):
-    """torch.optim.Optimizer with some additional methods related to TensorList.
-
-    Args:
-        params (ParamsT): iterable of parameters.
-        defaults (_type_): dictionary with default parameters for the optimizer.
-    """
-    def __init__(self, params: ParamsT, defaults):
-        super().__init__(params, defaults)
-        self._params: list[torch.Tensor] = [param for group in self.param_groups for param in group['params']]
-        self.has_complex = any(torch.is_complex(x) for x in self._params)
-        """True if any of the params are complex"""
-
-    def add_param_group(self, param_group: dict[str, Any]) -> None:
-        super().add_param_group(param_group)
-        self._params: list[torch.Tensor] = [param for group in self.param_groups for param in group['params']]
-        self.has_complex = any(torch.is_complex(x) for x in self._params)
-
-    # def get_params[CLS: Any](self, cls: type[CLS] = TensorList) -> CLS:
-    def get_params(self, cls: type[CLS] = TensorList) -> CLS:
-        """returns all params with `requires_grad = True` as a TensorList."""
-        return cls(p for p in self._params if p.requires_grad) # type:ignore
-
-    def ensure_grad_(self):
-        """Replaces None grad attribute with zeroes for all parameters that require grad."""
-        for p in self.get_params():
-            if p.requires_grad and p.grad is None: p.grad = torch.zeros_like(p)
-
-    # def get_state_key[CLS: MutableSequence](self, key: str, init: _StateInit = torch.zeros_like, params=None, cls: type[CLS] = TensorList) -> CLS:
-    def get_state_key(self, key: str, init: _StateInit = torch.zeros_like, params=None, cls: type[CLS] = TensorList) -> CLS:
-        """Returns a tensorlist of all `key` states of all params with `requires_grad = True`.
-
-        Args:
-            key (str): key to create/access.
-            init: Initial value if key doesn't exist. Can be `params`, `grad`, or callable such as `torch.zeros_like`.
-                Defaults to torch.zeros_like.
-            params (optional): optionally pass params if you already created them. Defaults to None.
-            cls (optional): optionally specify any other MutableSequence subclass to use instead of TensorList.
-
-        Returns:
-            TensorList: TensorList with the `key` state. Those tensors are stored in the optimizer, so modify them in-place.
-        """
-        value = cls()
-        if params is None: params = self.get_params()
-        for pi, p in enumerate(params):
-            state = self.state[p]
-            if key not in state:
-                if callable(init): state[key] = init(p)
-                elif isinstance(init, TensorList): state[key] = init[pi].clone()
-                elif init == 'params': state[key] = p.clone().detach()
-                elif init == 'grad': state[key] = p.grad.clone().detach() if p.grad is not None else torch.zeros_like(p)
-                else: raise ValueError(f'unknown init - {init}')
-            value.append(state[key]) # type:ignore
-        return value
-
-    # def get_state_keys[CLS: MutableSequence](
-    def get_state_keys(
-        self,
-        *keys: str,
-        inits: _StateInit | Sequence[_StateInit] = torch.zeros_like,
-        params=None,
-        cls: type[CLS] = TensorList,
-    ) -> list[CLS]:
-        """Returns a TensorList with the `key` states of all `params`. Creates the states if they don't exist."""
-
-        values = [cls() for _ in range(len(keys))]
-        if params is None: params = self.get_params()
-        if callable(inits) or isinstance(inits, str): inits = [inits] * len(keys) # type:ignore
-
-        for pi, p in enumerate(params):
-            state = self.state[p]
-            for i, (key, init) in enumerate(zip(keys, inits)): # type:ignore
-                if key not in state:
-                    if callable(init): state[key] = init(p)
-                    elif isinstance(init, TensorList): state[key] = init[pi].clone()
-                    elif init == 'params': state[key] = p.clone().detach()
-                    elif init == 'grad': state[key] = p.grad.clone().detach() if p.grad is not None else torch.zeros_like(p)
-                    else: raise ValueError(f'unknown init - {init}')
-                values[i].append(state[key]) # type:ignore
-        return values
-
-    def _yield_groups_key(self, key: str):
-        for group in self.param_groups:
-            value = group[key]
-            for p in group['params']:
-                if p.requires_grad: yield value
-
-
-    # def get_group_key[CLS: Any](self, key: str, cls: type[CLS] = NumberList) -> CLS:
-    def get_group_key(self, key: str, cls: type[CLS] = NumberList) -> CLS:
-        """Returns a TensorList with the param_groups `key` setting of each param."""
-        return cls(self._yield_groups_key(key)) # type:ignore
-
-    def get_first_group_key(self, key:str) -> Any:
-        """Returns the param_groups `key` setting of the first param."""
-        return next(iter(self._yield_groups_key(key)))
-
-    # def get_all_group_keys[CLS: Any](self, cls: type[CLS] = NumberList) -> dict[str, CLS]:
-    def get_all_group_keys(self, cls: type[CLS] = NumberList) -> dict[str, CLS]:
-        all_values: dict[str, CLS] = {}
-        for group in self.param_groups:
-
-            n_params = len([p for p in group['params'] if p.requires_grad])
-
-            for key, value in group.items():
-                if key != 'params':
-                    if key not in all_values: all_values[key] = cls(value for _ in range(n_params)) # type:ignore
-                    else: all_values[key].extend([value for _ in range(n_params)]) # type:ignore
-
-        return all_values
-
-    # def get_group_keys[CLS: MutableSequence](self, *keys: str, cls: type[CLS] = NumberList) -> list[CLS]:
-    def get_group_keys(self, *keys: str, cls: type[CLS] = NumberList) -> list[CLS]:
-        """Returns a list with the param_groups `key` setting of each param."""
-
-        all_values: list[CLS] = [cls() for _ in keys]
-        for group in self.param_groups:
-
-            n_params = len([p for p in group['params'] if p.requires_grad])
-
-            for i, key in enumerate(keys):
-                value = group[key]
-                all_values[i].extend([value for _ in range(n_params)]) # type:ignore
-
-        return all_values
-
-    @torch.no_grad
-    def evaluate_loss_at_vec(self, vec, closure=None, params = None, backward=False, ensure_float=False):
-        """_summary_
-
-        Args:
-            vec (_type_): _description_
-            closure (_type_, optional): _description_. Defaults to None.
-            params (_type_, optional): _description_. Defaults to None.
-            backward (bool, optional): _description_. Defaults to False.
-            ensure_float (bool, optional): _description_. Defaults to False.
-
-        Returns:
-            _type_: _description_
-        """
-        vec = totensor(vec)
-        if closure is None: closure = self._closure # type:ignore # pylint:disable=no-member
-        if params is None: params = self.get_params()
-
-        params.from_vec_(vec.to(params[0]))
-        loss = _maybe_pass_backward(closure, backward)
-
-        if ensure_float: return tofloat(loss)
-        return _maybe_pass_backward(closure, backward)
-
-    @overload
-    def evaluate_loss_grad_at_vec(self, vec, closure=None, params = None, to_numpy: Literal[True] = False) -> tuple[float, np.ndarray]: ... # type:ignore
-    @overload
-    def evaluate_loss_grad_at_vec(self, vec, closure=None, params = None, to_numpy: Literal[False] = False) -> tuple[_ScalarLoss, torch.Tensor]: ...
-    @torch.no_grad
-    def evaluate_loss_grad_at_vec(self, vec, closure=None, params = None, to_numpy: Literal[True] | Literal[False]=False):
-        """_summary_
-
-        Args:
-            vec (_type_): _description_
-            closure (_type_, optional): _description_. Defaults to None.
-            params (_type_, optional): _description_. Defaults to None.
-            to_numpy (Literal[True] | Literal[False], optional): _description_. Defaults to False.
-
-        Returns:
-            _type_: _description_
-        """
-        if params is None: params = self.get_params()
-        loss = self.evaluate_loss_at_vec(vec, closure, params, backward = True, ensure_float = to_numpy)
-        grad = params.grad.to_vec()
-
-        if to_numpy: return tofloat(loss), grad.detach().cpu().numpy()
-        return loss, grad
-
-
-    @torch.no_grad
-    def _maybe_evaluate_closure(self, closure, backward=True):
-        loss = None
-        if closure is not None:
-            loss = _maybe_pass_backward(closure, backward)
-        return loss

torchzero/modules/adaptive/__init__.py

@@ -1,4 +0,0 @@
-r"""
-Modules related to adapting the learning rate.
-"""
-from .adaptive import Cautious, UseGradMagnitude, UseGradSign, ScaleLRBySignChange, NegateOnSignChange

torchzero/modules/adaptive/adaptive.py

@@ -1,192 +0,0 @@
-import typing
-import torch
-
-from ...core import OptimizerModule
-
-class Cautious(OptimizerModule):
-    """Negates update for parameters where update and gradient sign is inconsistent.
-    Optionally normalizes the update by the number of parameters that are not masked.
-    This is meant to be used after any momentum-based modules.
-
-    Args:
-        normalize (bool, optional):
-            renormalize update after masking.
-            only has effect when mode is 'zero'. Defaults to False.
-        eps (float, optional): epsilon for normalization. Defaults to 1e-6.
-        mode (str, optional):
-            what to do with updates with inconsistent signs.
-
-            "zero" - set them to zero (as in paper)
-
-            "grad" - set them to the gradient
-
-            "negate" - negate them (same as using update magnitude and gradient sign)
-
-    reference
-        *Cautious Optimizers: Improving Training with One Line of Code.
-        Kaizhao Liang, Lizhang Chen, Bo Liu, Qiang Liu*
-    """
-    def __init__(self, normalize = False, eps=1e-6, mode: typing.Literal['zero', 'grad', 'backtrack'] = 'zero'):
-        super().__init__({})
-        self.eps = eps
-        self.normalize = normalize
-        self.mode: typing.Literal['zero', 'grad', 'backtrack'] = mode
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-        grad = vars.maybe_compute_grad_(params)
-
-        # mask will be > 0 for parameters where both signs are the same
-        mask = (ascent * grad) > 0
-        if self.mode in ('zero', 'grad'):
-            if self.normalize and self.mode == 'zero':
-                fmask = mask.to(ascent[0].dtype)
-                fmask /= fmask.total_mean() + self.eps # type:ignore
-            else:
-                fmask = mask
-
-            ascent *= fmask
-
-            if self.mode == 'grad':
-                ascent += grad * mask.logical_not_()
-
-            return ascent
-
-        # mode = 'backtrack'
-        ascent -= ascent.mul(2).mul_(mask.logical_not_())
-        return ascent
-
-
-class UseGradSign(OptimizerModule):
-    """
-    Uses update magnitude but gradient sign.
-    """
-    def __init__(self):
-        super().__init__({})
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-        grad = vars.maybe_compute_grad_(params)
-
-        return ascent.abs_().mul_(grad.sign())
-
-class UseGradMagnitude(OptimizerModule):
-    """
-    Uses update sign but gradient magnitude.
-    """
-    def __init__(self):
-        super().__init__({})
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-        grad = vars.maybe_compute_grad_(params)
-
-        return ascent.sign_().mul_(grad.abs())
-
-
-class ScaleLRBySignChange(OptimizerModule):
-    """
-    learning rate gets multiplied by `nplus` if ascent/gradient didn't change the sign,
-    or `nminus` if it did.
-
-    This is part of RProp update rule.
-
-    Args:
-        nplus (float): learning rate gets multiplied by `nplus` if ascent/gradient didn't change the sign
-        nminus (float): learning rate gets multiplied by `nminus` if ascent/gradient changed the sign
-        lb (float): lower bound for lr.
-        ub (float): upper bound for lr.
-        alpha (float): initial learning rate.
-
-    """
-    def __init__(self, nplus: float = 1.2, nminus: float = 0.5, lb = 1e-6, ub = 50, alpha=1, use_grad=False):
-        defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub)
-        super().__init__(defaults)
-        self.current_step = 0
-        self.use_grad = use_grad
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-
-        if self.use_grad: cur = vars.maybe_compute_grad_(params)
-        else: cur = ascent
-
-        nplus, nminus, lb, ub = self.get_group_keys('nplus', 'nminus', 'lb', 'ub')
-        prev, lrs = self.get_state_keys('prev_ascent', 'lrs', params=params)
-
-        # initialize on 1st step
-        if self.current_step == 0:
-            lrs.fill_(self.get_group_key('alpha'))
-            ascent.mul_(lrs)
-            prev.copy_(ascent)
-            self.current_step += 1
-            return ascent
-
-        mask = cur * prev
-        sign_changed = mask < 0
-        sign_same = mask > 0
-
-        # multiply magnitudes where sign didn't change
-        lrs.masked_set_(sign_same, lrs * nplus)
-        # multiply magnitudes where sign changed
-        lrs.masked_set_(sign_changed, lrs * nminus)
-        # bounds
-        lrs.clamp_(lb, ub)
-
-        ascent.mul_(lrs)
-        prev.copy_(cur)
-        self.current_step += 1
-        return ascent
-
-
-
-class NegateOnSignChange(OptimizerModule):
-    """Negates or undoes update for parameters where where gradient or update sign changes.
-
-    This is part of RProp update rule.
-
-    Args:
-        normalize (bool, optional): renormalize update after masking. Defaults to False.
-        eps (_type_, optional): epsilon for normalization. Defaults to 1e-6.
-        use_grad (bool, optional): if True, tracks sign change of the gradient,
-            otherwise track sign change of the update. Defaults to True.
-        backtrack (bool, optional): if True, undoes the update when sign changes, otherwise negates it.
-            Defaults to True.
-
-    """
-    # todo: add momentum to negation (to cautious as well and rprop negation as well)
-    def __init__(self, normalize = False, eps=1e-6, use_grad = False, backtrack = True):
-        super().__init__({})
-        self.eps = eps
-        self.normalize = normalize
-        self.use_grad = use_grad
-        self.backtrack = backtrack
-        self.current_step = 0
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-
-        if self.use_grad: cur = vars.maybe_compute_grad_(params)
-        else: cur = ascent
-
-        prev = self.get_state_key('prev')
-
-        # initialize on first step
-        if self.current_step == 0:
-            prev.set_(cur)
-            self.current_step += 1
-            return ascent
-
-        # mask will be > 0 for parameters where both signs are the same
-        mask = (cur * prev) < 0
-        if self.backtrack: ascent.masked_set_(mask, prev)
-        else: ascent.select_set_(mask, 0)
-
-        prev.set_(cur)
-        self.current_step += 1
-        return ascent