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

torchzero/utils/optimizer.py

@@ -0,0 +1,284 @@
+from collections.abc import Callable, Iterable, Mapping, MutableSequence, Sequence, MutableMapping
+from typing import Any, Literal, TypeVar, overload
+
+import torch
+
+from .tensorlist import TensorList
+from .numberlist import NumberList
+from .torch_tools import tofloat, totensor
+
+ListLike = TypeVar('ListLike', bound=MutableSequence)
+
+ParamFilter = Literal["has_grad", "requires_grad", "all"] | Callable[[torch.Tensor], bool]
+def _param_filter(param: torch.Tensor, mode: ParamFilter):
+    if callable(mode): return mode(param)
+    if mode == 'has_grad': return param.grad is not None
+    if mode == 'requires_grad': return param.requires_grad
+    if mode == 'all': return True
+    raise ValueError(f"Unknown mode {mode}")
+
+def get_params(
+    param_groups: Iterable[Mapping[str, Any]],
+    mode: ParamFilter = 'requires_grad',
+    cls: type[ListLike] = TensorList,
+) -> ListLike:
+    return cls(p for g in param_groups for p in g['params'] if _param_filter(p, mode)) # type:ignore[reportCallIssue]
+
+
+@overload
+def get_group_vals(param_groups: Iterable[Mapping[str, Any]], key: str, *,
+                   mode: ParamFilter = 'requires_grad', cls: type[ListLike] = list) -> ListLike: ...
+@overload
+def get_group_vals(param_groups: Iterable[Mapping[str, Any]], key: list[str] | tuple[str,...], *,
+                   mode: ParamFilter = 'requires_grad', cls: type[ListLike] = list) -> list[ListLike]: ...
+@overload
+def get_group_vals(param_groups: Iterable[Mapping[str, Any]], key: str, key2: str, *keys: str,
+                   mode: ParamFilter = 'requires_grad', cls: type[ListLike] = list) -> list[ListLike]: ...
+
+def get_group_vals(param_groups: Iterable[Mapping[str, Any]],
+                   key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
+                   mode: ParamFilter = 'requires_grad', cls: type[ListLike] = list) -> ListLike | list[ListLike]:
+
+    # single key, return single cls
+    if isinstance(key, str) and key2 is None:
+        values = cls()
+        for group in param_groups:
+            num_params = len([p for p in group['params'] if _param_filter(p, mode)])
+            if num_params > 0:
+                group_value = group[key]
+                values.extend(group_value for _ in range(num_params))
+        return values
+
+    # multiple keys
+    k1 = (key,) if isinstance(key, str) else tuple(key)
+    k2 = () if key2 is None else (key2,)
+    keys = k1 + k2 + keys
+
+    values = [cls() for _ in keys]
+    for group in param_groups:
+        num_params = len([p for p in group['params'] if _param_filter(p, mode)])
+        if num_params > 0:
+            for i,key in enumerate(keys):
+                group_value = group[key]
+                values[i].extend(group_value for _ in range(num_params))
+    return values
+
+_InitLiterals = Literal['param', 'grad']
+Init = _InitLiterals | Any | list[_InitLiterals | Any] | tuple[_InitLiterals | Any]
+
+def _make_initial_state_value(param: torch.Tensor, init: Init, i: int | None):
+    if callable(init): return init(param)
+    if isinstance(init, torch.Tensor): return init.detach().clone()
+
+    if isinstance(init, str):
+        if init in ('param','params'): return param.detach().clone()
+        if init in ('grad', 'grads'):
+            if param.grad is None: raise RuntimeError('init is set to "grad, but param.grad is None"')
+            return param.grad.detach().clone()
+
+    if isinstance(init, (list,tuple)):
+        if i is None: raise RuntimeError(f'init is per-parameter ({type(init)}) but parameter index i is None')
+        return _make_initial_state_value(param, init[i], None)
+
+    return init
+
+@overload
+def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], params: Sequence[torch.Tensor],
+                   key: str, *,
+                   must_exist: bool = False, init: Init = torch.zeros_like,
+                   cls: type[ListLike] = list) -> ListLike: ...
+@overload
+def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], params: Sequence[torch.Tensor],
+                   key: list[str] | tuple[str,...], *,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> list[ListLike]: ...
+@overload
+def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], params: Sequence[torch.Tensor],
+                   key: str,  key2: str, *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> list[ListLike]: ...
+
+def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], params: Sequence[torch.Tensor],
+                   key: str | list[str] | tuple[str,...], key2: str | None = None,  *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+                   cls: type[ListLike] = list) -> ListLike | list[ListLike]:
+
+    # single key, return single cls
+    if isinstance(key, str) and key2 is None:
+        values = cls()
+        for i, param in enumerate(params):
+            s = state[param]
+            if key not in s:
+                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                s[key] = _make_initial_state_value(param, init, i)
+            values.append(s[key])
+        return values
+
+    # multiple keys
+    k1 = (key,) if isinstance(key, str) else tuple(key)
+    k2 = () if key2 is None else (key2,)
+    keys = k1 + k2 + keys
+
+    values = [cls() for _ in keys]
+    for i, param in enumerate(params):
+        s = state[param]
+        for k_i, key in enumerate(keys):
+            if key not in s:
+                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                k_init = init[k_i] if isinstance(init, (list,tuple)) else init
+                s[key] = _make_initial_state_value(param, k_init, i)
+            values[k_i].append(s[key])
+
+    return values
+
+
+
+def loss_at_params(closure, params: Iterable[torch.Tensor],
+                   new_params: Sequence[torch.Tensor] | Any, backward: bool, restore=False):
+    params = TensorList(params)
+
+    old_params = params.clone() if restore else None
+
+    if isinstance(new_params, Sequence) and isinstance(new_params[0], torch.Tensor):
+        # when not restoring, copy new_params to params to avoid unexpected bugs due to shared storage
+        # when restoring params will be set back to old_params so its fine
+        if restore: params.set_(new_params)
+        else: params.copy_(new_params) # type:ignore
+
+    else:
+        new_params = totensor(new_params)
+        params.from_vec_(new_params)
+
+    if backward: loss = closure()
+    else: loss = closure(False)
+
+    if restore:
+        assert old_params is not None
+        params.set_(old_params)
+
+    return tofloat(loss)
+
+def loss_grad_at_params(closure, params: Iterable[torch.Tensor], new_params: Sequence[torch.Tensor], restore=False):
+    params = TensorList(params)
+    old_params = params.clone() if restore else None
+    loss = loss_at_params(closure, params, new_params, backward=True, restore=False)
+    grad = params.ensure_grad_().grad
+
+    if restore:
+        assert old_params is not None
+        params.set_(old_params)
+
+    return loss, grad
+
+def grad_at_params(closure, params: Iterable[torch.Tensor], new_params: Sequence[torch.Tensor], restore=False):
+    return loss_grad_at_params(closure=closure,params=params,new_params=new_params,restore=restore)[1]
+
+def loss_grad_vec_at_params(closure, params: Iterable[torch.Tensor], new_params: Any, restore=False):
+    params = TensorList(params)
+    old_params = params.clone() if restore else None
+    loss = loss_at_params(closure, params, new_params, backward=True, restore=False)
+    grad = params.ensure_grad_().grad.to_vec()
+
+    if restore:
+        assert old_params is not None
+        params.set_(old_params)
+
+    return loss, grad
+
+def grad_vec_at_params(closure, params: Iterable[torch.Tensor], new_params: Any, restore=False):
+    return loss_grad_vec_at_params(closure=closure,params=params,new_params=new_params,restore=restore)[1]
+
+
+
+class Optimizer(torch.optim.Optimizer):
+    """subclass of torch.optim.Optimizer with some helper methods for fast experimentation, it's not used anywhere in torchzero.
+
+    Args:
+        params (iterable): an iterable of :class:`torch.Tensor` s or
+            :class:`dict` s. Specifies what Tensors should be optimized.
+        defaults (dict | None): a dict containing default values of optimization
+            options (used when a parameter group doesn't specify them).
+    """
+    def __init__(self, params, defaults: dict[str, Any] | None = None, **_defaults):
+        if defaults is None: defaults = {}
+        defaults.update(_defaults)
+
+        super().__init__(params, defaults)
+        self.global_state = self.state[self.param_groups[0]['params'][0]]
+        """state of 1st parameter, can be used as global state which is how L-BFGS uses it in pytorch, and there is some kind of good reason to do it like that"""
+
+    def get_params(self, mode: ParamFilter = 'requires_grad', cls: type[ListLike] = TensorList) -> ListLike:
+        return get_params(self.param_groups, mode, cls)
+
+    @overload
+    def group_vals(self, key: str, *,
+                   mode: ParamFilter = 'requires_grad', cls: type[ListLike] = NumberList) -> ListLike: ...
+    @overload
+    def group_vals(self, key: list[str] | tuple[str,...], *,
+                   mode: ParamFilter = 'requires_grad', cls: type[ListLike] = NumberList) -> list[ListLike]: ...
+    @overload
+    def group_vals(self, key: str, key2: str, *keys: str,
+                   mode: ParamFilter = 'requires_grad', cls: type[ListLike] = NumberList) -> list[ListLike]: ...
+
+    def group_vals(self, key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
+                   mode: ParamFilter = 'requires_grad', cls: type[ListLike] = NumberList) -> ListLike | list[ListLike]:
+        return get_group_vals(self.param_groups, key, key2, *keys, mode = mode, cls = cls) # pyright:ignore[reportArgumentType]
+
+
+    @overload
+    def state_vals(self, key: str, *,
+                   init: Init = torch.zeros_like,
+                   mode: ParamFilter | list[torch.Tensor] | tuple[torch.Tensor, ...] = 'requires_grad',
+                   cls: type[ListLike] = TensorList) -> ListLike: ...
+    @overload
+    def state_vals(self, key: list[str] | tuple[str,...], *,
+                   init: Init | Sequence[Init] = torch.zeros_like,
+                   mode: ParamFilter | list[torch.Tensor] | tuple[torch.Tensor, ...] = 'requires_grad',
+                   cls: type[ListLike] = TensorList) -> list[ListLike]: ...
+    @overload
+    def state_vals(self, key: str, key2: str, *keys: str,
+                   init: Init | Sequence[Init] = torch.zeros_like,
+                   mode: ParamFilter | list[torch.Tensor] | tuple[torch.Tensor, ...] = 'requires_grad',
+                   cls: type[ListLike] = TensorList) -> list[ListLike]: ...
+
+    def state_vals(self, key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
+                   init: Init | Sequence[Init] = torch.zeros_like,
+                   mode: ParamFilter | list[torch.Tensor] | tuple[torch.Tensor, ...] = 'requires_grad',
+                   cls: type[ListLike] = TensorList) -> ListLike | list[ListLike]:
+
+        if isinstance(mode, (list,tuple)): params = mode
+        else: params = self.get_params(mode)
+
+        return get_state_vals(self.state, params, key, key2, *keys, init = init, cls = cls) # type:ignore[reportArgumentType]
+
+    def loss_at_params(self, closure, params: Sequence[torch.Tensor] | Any, backward: bool, restore=False):
+        return loss_at_params(closure=closure,params=self.get_params(),new_params=params,backward=backward,restore=restore)
+
+    def loss_grad_at_params(self, closure, params: Sequence[torch.Tensor] | Any, restore=False):
+        return loss_grad_at_params(closure=closure,params=self.get_params(),new_params=params,restore=restore)
+
+    def grad_at_params(self, closure, new_params: Sequence[torch.Tensor], restore=False):
+        return self.loss_grad_at_params(closure=closure,params=new_params,restore=restore)[1]
+
+    def loss_grad_vec_at_params(self, closure, params: Any, restore=False):
+        return loss_grad_vec_at_params(closure=closure,params=self.get_params(),new_params=params,restore=restore)
+
+    def grad_vec_at_params(self, closure, params: Any, restore=False):
+        return self.loss_grad_vec_at_params(closure=closure,params=params,restore=restore)[1]
+
+
+def zero_grad_(params: Iterable[torch.Tensor], set_to_none):
+    if set_to_none:
+        for p in params:
+            p.grad = None
+
+    else:
+        grads = [p.grad for p in params if p.grad is not None]
+        for grad in grads:
+            # taken from torch.optim.Optimizer.zero_grad
+            if grad.grad_fn is not None:
+                grad.detach_()
+            else:
+                grad.requires_grad_(False)
+
+        torch._foreach_zero_(grads)

torchzero/utils/optuna_tools.py

@@ -0,0 +1,40 @@
+import optuna
+
+from ..core import Chain, Module
+
+from ..modules import (
+    EMA,
+    NAG,
+    Cautious,
+    ClipNorm,
+    ClipNormGrowth,
+    ClipValue,
+    ClipValueGrowth,
+    Debias,
+    Normalize,
+)
+
+
+def get_momentum(trial: optuna.Trial, prefix: str, conditional: bool=True) -> list[Module]:
+    cond = trial.suggest_categorical(f'{prefix}_use_momentum', [True,False]) if conditional else True
+    if cond:
+        beta = trial.suggest_float(f'{prefix}_beta', -1, 2)
+        dampening = trial.suggest_float(f'{prefix}_dampening', -1, 2)
+        lerp = trial.suggest_categorical(f'{prefix}_use_lerp', [True, False])
+        nag = trial.suggest_categorical(f'{prefix}_use_NAG', [True, False])
+        debiased = trial.suggest_categorical(f'{prefix}_debiased', [True, False])
+        if nag:
+            m = NAG(beta, dampening, lerp)
+            if debiased: m = Chain(m, Debias(beta1=beta))
+        else:
+            m = EMA(beta, dampening, debiased=debiased, lerp=lerp)
+        return [m]
+    return []
+
+def get_clip_value(trial: optuna.Trial, prefix: str, conditional: bool=True) -> list[Module]:
+    cond = trial.suggest_categorical(f'{prefix}_use_clip_value', [True,False]) if conditional else True
+    if cond:
+        return [ClipValue(value = trial.suggest_float(f'{prefix}_clip_value', 0, 10))]
+    return []
+
+

torchzero/utils/params.py

@@ -0,0 +1,149 @@
+from typing import Any
+from collections.abc import Sequence, Iterable, Mapping
+import warnings
+import torch, numpy as np
+
+
+
+Params = Iterable[torch.Tensor | tuple[str, torch.Tensor] | Mapping[str, Any]]
+
+def _validate_params_are_unique_(params: Sequence[torch.Tensor]):
+    # this is from pytorch add_param_group
+    if len(params) != len(set(params)):
+        warnings.warn(
+            "optimizer contains a parameter group with duplicate parameters; "
+            "in future, this will cause an error; "
+            "see github.com/pytorch/pytorch/issues/40967 for more information",
+            stacklevel=3,
+        )
+
+def _validate_param_is_differentiable_(tensor: torch.Tensor | Any):
+    """Checks that param is torch.Tensor and isn't a leaf parameter unless differentiable is True, otherwise this raises, this is taken from torch.optim.Optimizer."""
+    if not (tensor.is_leaf or tensor.retains_grad):
+        raise ValueError("can't optimize a non-leaf Tensor")
+
+def _validate_at_least_one_param_requires_grad_(params: Iterable[torch.Tensor]):
+    params = list(params)
+    if not any(p.requires_grad for p in params):
+        warnings.warn(
+            "Parameter group contains no parameters which require gradients. "
+            "Note for gradient-free optimizers, they still only optimize parameters with requires_grad=True, "
+            "so if needed, use `with torch.no_grad():` context instead.", stacklevel=3)
+
+
+
+def _copy_param_groups(param_groups: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """copies param_groups, doesn't copy the tensors."""
+    new_param_group = []
+
+    for g in param_groups:
+        assert isinstance(g, dict)
+        g_copy = g.copy()
+
+        for k in ('params', 'updates', 'grads'):
+            if k in g_copy:
+                assert isinstance(g_copy[k], list)
+                g_copy[k] = g_copy[k].copy()
+
+        new_param_group.append(g_copy)
+
+    return new_param_group
+
+def _process_param_group_(param_group: dict[str, Any]) -> dict[str, Any]:
+    """makes sure `param_group["params"]` is a list of tensors, and sets `param_group["param_names"]` if params are named."""
+    if 'params' not in param_group: raise KeyError("Param group doesn't have a `params` key.")
+
+    if isinstance(param_group['params'], torch.Tensor): param_group['params'] = [param_group['params']]
+
+    tensors: list[torch.Tensor] = []
+    names: list[str] | None = []
+
+    for p in param_group['params']:
+        if isinstance(p, torch.Tensor):
+            tensors.append(p)
+
+        elif isinstance(p, tuple):
+            if len(p) != 2:
+                raise ValueError(f'named_parameters must be a tuple of (name, tensor), got length {len(p)} tuple')
+            if (not isinstance(p[0], str)) or (not isinstance(p[1], torch.Tensor)):
+                raise ValueError(f'named_parameters must be a tuple of (name, tensor), got {[type(a) for a in p]}')
+            names.append(p[0])
+            tensors.append(p[1])
+
+        else:
+            raise ValueError(f'Parameters must be tensors or tuples (name, tensor), got parameter of type {type(p)}')
+
+    if len(tensors) == 0: warnings.warn('got an empty parameter group')
+
+    param_group['params'] = tensors
+
+    if len(names) != 0:
+        if len(names) != len(tensors):
+            raise ValueError(f"Number of parameters {len(tensors)} doesn't match number of names {len(names)}")
+        param_group['param_names'] = names
+
+    return param_group
+
+def _make_param_groups(params: Params, differentiable: bool) -> list[dict[str, Any]]:
+    params = list(params)
+
+    param_groups: list[dict[str, Any]] = [dict(p) for p in params if isinstance(p, Mapping)]
+    tensors = [p for p in params if isinstance(p, torch.Tensor)]
+    named_tensors = [p for p in params if isinstance(p, tuple)]
+
+    if len(tensors) != 0: param_groups.append({"params": tensors})
+    if len(named_tensors) != 0: param_groups.append({"params": named_tensors})
+
+    # process param_groups
+    for g in param_groups:
+        _process_param_group_(g)
+
+    # validate
+    all_params = [p for g in param_groups for p in g['params']]
+    _validate_params_are_unique_(all_params)
+    _validate_at_least_one_param_requires_grad_(all_params)
+    if not differentiable:
+        for p in all_params: _validate_param_is_differentiable_(p)
+
+    return param_groups
+
+def _add_defaults_to_param_groups_(param_groups: list[dict[str, Any]], defaults: dict[str, Any]) -> list[dict[str, Any]]:
+    for group in param_groups:
+        for k, v in defaults.items():
+            if k not in group:
+                group[k] = v
+    return param_groups
+
+def _add_updates_grads_to_param_groups_(param_groups: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    for group in param_groups:
+        if 'updates' in group: raise ValueError('updates in group')
+        group['updates'] = [None for _ in group['params']]
+
+        if 'grads' in group: raise ValueError('grads in group')
+        group['grads'] = [None for _ in group['grads']]
+
+    return param_groups
+
+
+def _set_update_and_grad_(
+    param_groups: list[dict[str, Any]],
+    updates: list[torch.Tensor] | None,
+    grads: list[torch.Tensor] | None,
+) -> list[dict[str, Any]]:
+    if updates is None and grads is None: return param_groups
+
+    updates_iter = iter(updates) if updates is not None else None
+    grads_iter = iter(grads) if grads is not None else None
+
+    for group in param_groups:
+        group_params = group['params']
+        group_updates = group['updates']
+        group_grads = group['grads']
+
+        for i, param in enumerate(group_params):
+            if not param.requires_grad: continue
+            if updates_iter is not None: group_updates[i] = next(updates_iter)
+            if grads_iter is not None: group_grads[i] = next(grads_iter)
+
+    return param_groups
+

torchzero/utils/python_tools.py

@@ -1,25 +1,40 @@
-import functools
-import operator
-from typing import Any, TypeVar
-from collections.abc import Iterable
-
-import torch
-
-def _flatten_no_check(iterable: Iterable) -> list[Any]:
-    """Flatten an iterable of iterables, returns a flattened list. Note that if `iterable` is not Iterable, this will return `[iterable]`."""
-    if isinstance(iterable, Iterable):
-        return [a for i in iterable for a in _flatten_no_check(i)]
-    return [iterable]
-
-def flatten(iterable: Iterable) -> list[Any]:
-    """Flatten an iterable of iterables, returns a flattened list. If `iterable` is not iterable, raises a TypeError."""
-    if isinstance(iterable, Iterable): return [a for i in iterable for a in _flatten_no_check(i)]
-    raise TypeError(f'passed object is not an iterable, {type(iterable) = }')
-
-X = TypeVar("X")
-# def reduce_dim[X](x:Iterable[Iterable[X]]) -> list[X]: # pylint:disable=E0602
-def reduce_dim(x:Iterable[Iterable[X]]) -> list[X]: # pylint:disable=E0602
-    """Reduces one level of nesting. Takes an iterable of iterables of X, and returns an iterable of X."""
-    return functools.reduce(operator.iconcat, x, [])
-
-_ScalarLoss = int | float | bool | torch.Tensor
+import functools
+import operator
+from typing import Any, TypeVar
+from collections.abc import Iterable, Callable
+from collections import UserDict
+
+
+def _flatten_no_check(iterable: Iterable) -> list[Any]:
+    """Flatten an iterable of iterables, returns a flattened list. Note that if `iterable` is not Iterable, this will return `[iterable]`."""
+    if isinstance(iterable, Iterable) and not isinstance(iterable, str):
+        return [a for i in iterable for a in _flatten_no_check(i)]
+    return [iterable]
+
+def flatten(iterable: Iterable) -> list[Any]:
+    """Flatten an iterable of iterables, returns a flattened list. If `iterable` is not iterable, raises a TypeError."""
+    if isinstance(iterable, Iterable): return [a for i in iterable for a in _flatten_no_check(i)]
+    raise TypeError(f'passed object is not an iterable, {type(iterable) = }')
+
+X = TypeVar("X")
+# def reduce_dim[X](x:Iterable[Iterable[X]]) -> list[X]: # pylint:disable=E0602
+def reduce_dim(x:Iterable[Iterable[X]]) -> list[X]: # pylint:disable=E0602
+    """Reduces one level of nesting. Takes an iterable of iterables of X, and returns an iterable of X."""
+    return functools.reduce(operator.iconcat, x, [])
+
+def generic_eq(x: int | float | Iterable[int | float], y: int | float | Iterable[int | float]) -> bool:
+    """generic equals function that supports scalars and lists of numbers"""
+    if isinstance(x, (int,float)):
+        if isinstance(y, (int,float)): return x==y
+        return all(i==x for i in y)
+    if isinstance(y, (int,float)):
+        return all(i==y for i in x)
+    return all(i==j for i,j in zip(x,y))
+
+def zipmap(self, fn: Callable, other: Any | list | tuple, *args, **kwargs):
+    """If `other` is list/tuple, applies `fn` to self zipped with `other`.
+    Otherwise applies `fn` to this sequence and `other`.
+    Returns a new sequence with return values of the callable."""
+    if isinstance(other, (list, tuple)): return self.__class__(fn(i, j, *args, **kwargs) for i, j in zip(self, other))
+    return self.__class__(fn(i, other, *args, **kwargs) for i in self)
+