pitchscapes.optimization.WarmAdam

class pitchscapes.optimization.WarmAdam(params, lr=0.001, init_lr=None, lr_beta=0.0, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)[source]

Bases: Optimizer

Modified from Adam class

Public Data Attributes:

`OptimizerPreHook`	alias of `Callable`[[`Self`, `Tuple`[`Any`, ...], `Dict`[`str`, `Any`]], `Tuple`[`Tuple`[`Any`, ...], `Dict`[`str`, `Any`]] \| `None`]
`OptimizerPostHook`	alias of `Callable`[[`Self`, `Tuple`[`Any`, ...], `Dict`[`str`, `Any`]], `None`]
`state`
`param_groups`

Inherited from Optimizer

`OptimizerPreHook`	alias of `Callable`[[`Self`, `Tuple`[`Any`, ...], `Dict`[`str`, `Any`]], `Tuple`[`Tuple`[`Any`, ...], `Dict`[`str`, `Any`]] \| `None`]
`OptimizerPostHook`	alias of `Callable`[[`Self`, `Tuple`[`Any`, ...], `Dict`[`str`, `Any`]], `None`]

Public Methods:

`__init__`(params[, lr, init_lr, lr_beta, ...])
`__setstate__`(state)
`step`([closure])	Performs a single optimization step.

Inherited from Optimizer

`__init__`(params, defaults)
`__getstate__`()
`__setstate__`(state)
`__repr__`()	Return repr(self).
`profile_hook_step`(func)
`register_step_pre_hook`(hook)	Register an optimizer step pre hook which will be called before optimizer step. It should have the following signature::.
`register_step_post_hook`(hook)	Register an optimizer step post hook which will be called after optimizer step. It should have the following signature::.
`register_state_dict_pre_hook`(hook[, prepend])	Register a state dict pre-hook which will be called before `state_dict()` is called. It should have the following signature::.
`register_state_dict_post_hook`(hook[, prepend])	Register a state dict post-hook which will be called after `state_dict()` is called. It should have the following signature::.
`state_dict`()	Returns the state of the optimizer as a `dict`.
`register_load_state_dict_pre_hook`(hook[, ...])	Register a load_state_dict pre-hook which will be called before `load_state_dict()` is called. It should have the following signature::.
`register_load_state_dict_post_hook`(hook[, ...])	Register a load_state_dict post-hook which will be called after `load_state_dict()` is called. It should have the following signature::.
`load_state_dict`(state_dict)	Loads the optimizer state.
`zero_grad`([set_to_none])	Resets the gradients of all optimized `torch.Tensor` s.
`step`([closure])	Performs a single optimization step (parameter update).
`add_param_group`(param_group)	Add a param group to the `Optimizer` s param_groups.

Private Data Attributes:

`_optimizer_step_pre_hooks`
`_optimizer_step_post_hooks`
`_optimizer_state_dict_pre_hooks`
`_optimizer_state_dict_post_hooks`
`_optimizer_load_state_dict_pre_hooks`
`_optimizer_load_state_dict_post_hooks`

Inherited from Optimizer

`_optimizer_step_pre_hooks`
`_optimizer_step_post_hooks`
`_optimizer_state_dict_pre_hooks`
`_optimizer_state_dict_post_hooks`
`_optimizer_load_state_dict_pre_hooks`
`_optimizer_load_state_dict_post_hooks`

Private Methods:

Inherited from Optimizer

`_cuda_graph_capture_health_check`()
`_optimizer_step_code`()	Entry point for torch.profile.profiler.
`_group_tensors_by_device_and_dtype`(...[, ...])	Groups a list of lists of tensors by device and dtype.
`_patch_step_function`()
`_process_value_according_to_param_policy`(...)

OptimizerPostHook: alias of Callable[[Self, Tuple[Any, …], Dict[str, Any]], None]

OptimizerPreHook: alias of Callable[[Self, Tuple[Any, …], Dict[str, Any]], Tuple[Tuple[Any, …], Dict[str, Any]] | None]

__annotations__ = {'OptimizerPostHook': 'TypeAlias', 'OptimizerPreHook': 'TypeAlias', '_optimizer_load_state_dict_post_hooks': '\'OrderedDict[int, Callable[["Optimizer"], None]]\'', '_optimizer_load_state_dict_pre_hooks': '\'OrderedDict[int, Callable[["Optimizer", StateDict], Optional[StateDict]]]\'', '_optimizer_state_dict_post_hooks': '\'OrderedDict[int, Callable[["Optimizer", StateDict], Optional[StateDict]]]\'', '_optimizer_state_dict_pre_hooks': '\'OrderedDict[int, Callable[["Optimizer"], None]]\'', '_optimizer_step_post_hooks': 'Dict[int, OptimizerPostHook]', '_optimizer_step_pre_hooks': 'Dict[int, OptimizerPreHook]', 'param_groups': 'List[Dict[str, Any]]', 'state': 'DefaultDict[torch.Tensor, Any]'}

__dict__ = mappingproxy({'__module__': 'pitchscapes.optimization', '__doc__': '\n Modified from Adam class\n ', '__init__': <function WarmAdam.__init__>, '__setstate__': <function WarmAdam.__setstate__>, 'step': <function WarmAdam.step>, '__annotations__': {'OptimizerPreHook': 'TypeAlias', 'OptimizerPostHook': 'TypeAlias', '_optimizer_step_pre_hooks': 'Dict[int, OptimizerPreHook]', '_optimizer_step_post_hooks': 'Dict[int, OptimizerPostHook]', '_optimizer_state_dict_pre_hooks': '\'OrderedDict[int, Callable[["Optimizer"], None]]\'', '_optimizer_state_dict_post_hooks': '\'OrderedDict[int, Callable[["Optimizer", StateDict], Optional[StateDict]]]\'', '_optimizer_load_state_dict_pre_hooks': '\'OrderedDict[int, Callable[["Optimizer", StateDict], Optional[StateDict]]]\'', '_optimizer_load_state_dict_post_hooks': '\'OrderedDict[int, Callable[["Optimizer"], None]]\'', 'state': 'DefaultDict[torch.Tensor, Any]', 'param_groups': 'List[Dict[str, Any]]'}})

__getstate__()

Return type:: Dict[str, Any]

__init__(params, lr=0.001, init_lr=None, lr_beta=0.0, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)[source]

__module__ = 'pitchscapes.optimization'

__repr__()

Return repr(self).

Return type:: str

__setstate__(state)[source]

__weakref__: list of weak references to the object (if defined)

_cuda_graph_capture_health_check()

Return type:: None

static _group_tensors_by_device_and_dtype(tensorlistlist, with_indices=False)

Groups a list of lists of tensors by device and dtype. Skips this step if we are compiling since this will occur during inductor lowering.

Return type:: Union[Dict[Tuple[None, None], Tuple[List[List[Optional[Tensor]]], List[int]]], Dict[Tuple[device, dtype], Tuple[List[List[Optional[Tensor]]], List[int]]]]

_optimizer_load_state_dict_post_hooks: OrderedDict[int, Callable[["Optimizer"], None]]

_optimizer_load_state_dict_pre_hooks: OrderedDict[int, Callable[["Optimizer", StateDict], Optional[StateDict]]]

_optimizer_state_dict_post_hooks: OrderedDict[int, Callable[["Optimizer", StateDict], Optional[StateDict]]]

_optimizer_state_dict_pre_hooks: OrderedDict[int, Callable[["Optimizer"], None]]

_optimizer_step_code()

Entry point for torch.profile.profiler.

When python tracing is enabled the profiler will hook into this function at the CPython level to inspect the optimizer’s parameters and param groups. It is called it after step() since many optimizers lazily initialize state.

This is a workaround due to lack of a proper step hook on the optimizer, and will be removed if it exists.

Return type:: None

_optimizer_step_post_hooks: Dict[int, OptimizerPostHook]

_optimizer_step_pre_hooks: Dict[int, OptimizerPreHook]

_patch_step_function()

Return type:: None

static _process_value_according_to_param_policy(param, value, param_id, param_groups, key=None)

Return type:: Tensor

add_param_group(param_group)

Add a param group to the Optimizer s param_groups.

This can be useful when fine tuning a pre-trained network as frozen layers can be made trainable and added to the Optimizer as training progresses.

Return type:: None

Args:

param_group (dict): Specifies what Tensors should be optimized along with group: specific optimization options.

load_state_dict(state_dict)

Loads the optimizer state.

Return type:: None

Args:

state_dict (dict): optimizer state. Should be an object returned: from a call to state_dict().

static profile_hook_step(func)

Return type:: Callable[[ParamSpec(_P)], TypeVar(R)]

register_load_state_dict_post_hook(hook, prepend=False)

Register a load_state_dict post-hook which will be called after load_state_dict() is called. It should have the following signature:

hook(optimizer) -> None

The optimizer argument is the optimizer instance being used.

The hook will be called with argument self after calling load_state_dict on self. The registered hook can be used to perform post-processing after load_state_dict has loaded the state_dict.

Return type:: RemovableHandle

Args:

hook (Callable): The user defined hook to be registered. prepend (bool): If True, the provided post hook will be fired before

all the already registered post-hooks on load_state_dict. Otherwise, the provided hook will be fired after all the already registered post-hooks. (default: False)

Returns:

torch.utils.hooks.RemoveableHandle:: a handle that can be used to remove the added hook by calling handle.remove()

register_load_state_dict_pre_hook(hook, prepend=False)

Register a load_state_dict pre-hook which will be called before load_state_dict() is called. It should have the following signature:

hook(optimizer, state_dict) -> state_dict or None

The optimizer argument is the optimizer instance being used and the state_dict argument is a shallow copy of the state_dict the user passed in to load_state_dict. The hook may modify the state_dict inplace or optionally return a new one. If a state_dict is returned, it will be used to be loaded into the optimizer.

The hook will be called with argument self and state_dict before calling load_state_dict on self. The registered hook can be used to perform pre-processing before the load_state_dict call is made.

Return type:: RemovableHandle

Args:

hook (Callable): The user defined hook to be registered. prepend (bool): If True, the provided pre hook will be fired before

all the already registered pre-hooks on load_state_dict. Otherwise, the provided hook will be fired after all the already registered pre-hooks. (default: False)

Returns:

torch.utils.hooks.RemoveableHandle:: a handle that can be used to remove the added hook by calling handle.remove()

register_state_dict_post_hook(hook, prepend=False)

Register a state dict post-hook which will be called after state_dict() is called. It should have the following signature:

hook(optimizer, state_dict) -> state_dict or None

The hook will be called with arguments self and state_dict after generating a state_dict on self. The hook may modify the state_dict inplace or optionally return a new one. The registered hook can be used to perform post-processing on the state_dict before it is returned.

Return type:: RemovableHandle

Args:

hook (Callable): The user defined hook to be registered. prepend (bool): If True, the provided post hook will be fired before

all the already registered post-hooks on state_dict. Otherwise, the provided hook will be fired after all the already registered post-hooks. (default: False)

Returns:

torch.utils.hooks.RemoveableHandle:: a handle that can be used to remove the added hook by calling handle.remove()

register_state_dict_pre_hook(hook, prepend=False)

Register a state dict pre-hook which will be called before state_dict() is called. It should have the following signature:

hook(optimizer) -> None

The optimizer argument is the optimizer instance being used. The hook will be called with argument self before calling state_dict on self. The registered hook can be used to perform pre-processing before the state_dict call is made.

Return type:: RemovableHandle

Args:

hook (Callable): The user defined hook to be registered. prepend (bool): If True, the provided pre hook will be fired before

all the already registered pre-hooks on state_dict. Otherwise, the provided hook will be fired after all the already registered pre-hooks. (default: False)

Returns:

torch.utils.hooks.RemoveableHandle:: a handle that can be used to remove the added hook by calling handle.remove()

register_step_post_hook(hook)

Register an optimizer step post hook which will be called after optimizer step. It should have the following signature:

hook(optimizer, args, kwargs) -> None

The optimizer argument is the optimizer instance being used.

Return type:: RemovableHandle

Args:

hook (Callable): The user defined hook to be registered.

Returns:

torch.utils.hooks.RemovableHandle:: a handle that can be used to remove the added hook by calling handle.remove()

register_step_pre_hook(hook)

Register an optimizer step pre hook which will be called before optimizer step. It should have the following signature:

hook(optimizer, args, kwargs) -> None or modified args and kwargs

The optimizer argument is the optimizer instance being used. If args and kwargs are modified by the pre-hook, then the transformed values are returned as a tuple containing the new_args and new_kwargs.

Return type:: RemovableHandle

Args:

hook (Callable): The user defined hook to be registered.

Returns:

torch.utils.hooks.RemovableHandle:: a handle that can be used to remove the added hook by calling handle.remove()

state_dict()

Returns the state of the optimizer as a dict.

It contains two entries: :rtype: Dict[str, Any]

state: a Dict holding current optimization state. Its content
differs between optimizer classes, but some common characteristics hold. For example, state is saved per parameter, and the parameter itself is NOT saved. state is a Dictionary mapping parameter ids to a Dict with state corresponding to each parameter.
param_groups: a List containing all parameter groups where each
parameter group is a Dict. Each parameter group contains metadata specific to the optimizer, such as learning rate and weight decay, as well as a List of parameter IDs of the parameters in the group.

NOTE: The parameter IDs may look like indices but they are just IDs associating state with param_group. When loading from a state_dict, the optimizer will zip the param_group params (int IDs) and the optimizer param_groups (actual nn.Parameter s) in order to match state WITHOUT additional verification.

A returned state dict might look something like:

{
    'state': {
        0: {'momentum_buffer': tensor(...), ...},
        1: {'momentum_buffer': tensor(...), ...},
        2: {'momentum_buffer': tensor(...), ...},
        3: {'momentum_buffer': tensor(...), ...}
    },
    'param_groups': [
        {
            'lr': 0.01,
            'weight_decay': 0,
            ...
            'params': [0]
        },
        {
            'lr': 0.001,
            'weight_decay': 0.5,
            ...
            'params': [1, 2, 3]
        }
    ]
}

step(closure=None)[source]

Performs a single optimization step.

Arguments:

closure (callable, optional): A closure that reevaluates the model: and returns the loss.

zero_grad(set_to_none=True)

Resets the gradients of all optimized torch.Tensor s.

Return type:: None

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: This will in general have lower memory footprint, and can modestly improve performance. However, it changes certain behaviors. For example: 1. When the user tries to access a gradient and perform manual ops on it, a None attribute or a Tensor full of 0s will behave differently. 2. If the user requests zero_grad(set_to_none=True) followed by a backward pass, .grads are guaranteed to be None for params that did not receive a gradient. 3. torch.optim optimizers have a different behavior if the gradient is 0 or None (in one case it does the step with a gradient of 0 and in the other it skips the step altogether).