Auditory Filterbanks 

Notes

This method dispatches to either _forward_sos() or _forward_poly() based on the implementation parameter set during initialization.

The output is real-valued and represents the instantaneous envelope at each frequency channel. The MATLAB convention of multiplying by 2 is applied to compensate for using only the analytic signal’s real part.

Computational Complexity:

SOS: O(nBFT) where n is filter order
Poly: O(BFT(nb + na)) where nb, na are coefficient lengths

For typical parameters (n=4, B=2, F=30, T=16000), expect ~2M operations.

See also

_forward_sos: SOS cascade implementation
_forward_poly: Polynomial expansion implementation

extra_repr()[source]

Extra representation string for module printing.

Returns:: String containing key module parameters: num_channels, fs, n, fc_range, betamul, learnable status.
Return type:: str

T_destination = ~T_destination

add_module(name, module)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (str): name of the child module. The child module can be: accessed from this module using the given name

module (Module): child module to be added to the module.

Parameters:

name (str)
module (Module | None)

Return type:

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Args:: fn (Module -> None): function to be applied to each submodule
Returns:: Module: self

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Parameters:: fn (Callable[[Module], None])
Return type:: Self

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

buffers(recurse=True)

Return an iterator over module buffers.

Args:

recurse (bool): if True, then yields buffers of this module: and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor: module buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Tensor]

call_super_init: bool = False

children()

Return an iterator over immediate children modules.

Return type:: Iterator[Module]

Yields:: Module: a child module

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

cuda(device=None)

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Args:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

double()

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

dump_patches: bool = False

eval()

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Return type:: Self

Returns:: Module: self

float()

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the buffer: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.Tensor: The buffer referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not a buffer

Parameters:: target (str)
Return type:: Tensor

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Return type:: Any

Returns:: object: Any extra state to store in the module’s state_dict

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the Parameter: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.nn.Parameter: The Parameter referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not an nn.Parameter

Parameters:: target (str)
Return type:: Parameter

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

Returns:

torch.nn.Module: The submodule referenced by target

Raises:

AttributeError: If at any point along the path resulting from: the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:: target (str)
Return type:: Module

half()

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

ipu(device=None)

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

load_state_dict(state_dict, strict=True, assign=False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Args:

state_dict (dict): a dict containing parameters and: persistent buffers.
strict (bool, optional): whether to strictly enforce that the keys: in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional): When set to False, the properties of the tensors: in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

NamedTuple with missing_keys and unexpected_keys fields:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Note:

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

Parameters:

state_dict (Mapping[str, Any])
strict (bool)
assign (bool)

modules()

Return an iterator over all modules in the network.

Return type:: Iterator[Module]

Yields:: Module: a module in the network
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Args:

prefix (str): prefix to prepend to all buffer names. recurse (bool, optional): if True, then yields buffers of this module

and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor): Tuple containing the name and buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Return type:: Iterator[tuple[str, Module]]

Yields:: (str, Module): Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Args:: memo: a memo to store the set of modules already added to the result prefix: a prefix that will be added to the name of the module remove_duplicate: whether to remove the duplicated module instances in the result

or not
Yields:: (str, Module): Tuple of name and module
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

Parameters:

memo (set[Module] | None)
prefix (str)
remove_duplicate (bool)

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Args:

prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module

and all submodules. Otherwise, yields only parameters that are direct members of this module.

remove_duplicate (bool, optional): whether to remove the duplicated: parameters in the result. Defaults to True.

Yields:

(str, Parameter): Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Args:

recurse (bool): if True, then yields parameters of this module: and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter: module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Parameter]

register_backward_hook(hook)

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

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

Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
Return type:: RemovableHandle

register_buffer(name, tensor, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (str): name of the buffer. The buffer can be accessed: from this module using the given name
tensor (Tensor or None): buffer to be registered. If None, then operations: that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool): whether the buffer is part of this module’s: state_dict.

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

Parameters:

name (str)
tensor (Tensor | None)
persistent (bool)

Return type:

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Args:

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

before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If True, the hook will be passed the: kwargs given to the forward function. Default: False
always_call (bool): If True the hook will be run regardless of: whether an exception is raised while calling the Module. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...], Any], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any], Any], Any | None])
prepend (bool)
with_kwargs (bool)
always_call (bool)

Return type:

RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Args:

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

all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If true, the hook will be passed the kwargs: given to the forward function. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...]], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None])
prepend (bool)
with_kwargs (bool)

Return type:

RemovableHandle

register_full_backward_hook(hook, prepend=False)

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

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

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Arguments:

hook (Callable): Callable hook that will be invoked before: loading the state dict.

register_module(name, module)

Alias for add_module().

Parameters:

name (str)
module (Module | None)

Return type:

register_parameter(name, param)

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (str): name of the parameter. The parameter can be accessed: from this module using the given name
param (Parameter or None): parameter to be added to the module. If: None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Parameters:

name (str)
param (Parameter | None)

Return type:

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Args:

requires_grad (bool): whether autograd should record operations on: parameters in this module. Default: True.

Returns:

Module: self

Parameters:: requires_grad (bool)
Return type:: Self

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Args:: state (dict): Extra state from the state_dict

Parameters:: state (Any)
Return type:: None

set_submodule(target, module, strict=False)

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

module: The module to set the submodule to. strict: If False, the method will replace an existing submodule

or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError: If the target string is empty or if module is not an instance of nn.Module. AttributeError: If at any point along the path resulting from

the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:

target (str)
module (Module)
strict (bool)

Return type:

share_memory()

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)

Overloads:

self, destination (T_destination), prefix (str), keep_vars (bool) → T_destination
self, prefix (str), keep_vars (bool) → dict[str, Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Args:

destination (dict, optional): If provided, the state of module will: be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional): a prefix added to parameter and buffer: names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional): by default the Tensor s: returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

dict:: a dictionary containing a whole state of the module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)

Overloads:

self, device (DeviceLikeType | None), dtype (dtype | None), non_blocking (bool) → Self
self, dtype (dtype), non_blocking (bool) → Self
self, tensor (Tensor), non_blocking (bool) → Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters: and buffers in this module
dtype (torch.dtype): the desired floating point or complex dtype of: the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired: dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format): the desired memory: format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)

Move the parameters and buffers to the specified device without copying storage.

Args:

device (torch.device): The desired device of the parameters: and buffers in this module.
recurse (bool): Whether parameters and buffers of submodules should: be recursively moved to the specified device.

Returns:

Module: self

Parameters:

device (str | device | int | None)
recurse (bool)

Return type:

Self

train(mode=True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation: mode (False). Default: True.

Returns:

Module: self

Parameters:: mode (bool)
Return type:: Self

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Parameters:: dst_type (dtype | str)
Return type:: Self

xpu(device=None)

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

zero_grad(set_to_none=True)

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: See torch.optim.Optimizer.zero_grad() for details.

Parameters:: set_to_none (bool)
Return type:: None

training: bool

DRNLFilterbank

class torch_amt.DRNLFilterbank(fc, fs, n_channels=50, subject='NH', model='paulick2024', learnable=False, dtype=torch.float64)[source]

Bases: Module

Dual Resonance Non-Linear (DRNL) filterbank for basilar membrane simulation.

Implements the DRNL auditory filterbank model used in Paulick et al. (2024) CASP model. The DRNL consists of two parallel signal paths (linear and nonlinear) that are summed to produce the basilar membrane velocity response.

The architecture models the cochlea’s dual mechanism for sound processing:

Linear path: Provides level-independent frequency selectivity via gammatone bandpass filtering and lowpass smoothing
Nonlinear path: Provides level-dependent compression via broken-stick nonlinearity sandwiched between gammatone filters

Parameters:

fc (Tensor | Tuple[float, float]) –
Center frequencies in Hz. Can be:
- torch.Tensor: Explicit center frequencies of shape (F,)
- tuple (flow, fhigh): Automatically generates ERB-spaced frequencies using erbspacebw()
fs (float) – Sampling rate in Hz.
n_channels (int) – Number of ERB-spaced channels. Only used if fc is tuple. Default: 50.
subject (str) –
Subject type. Default: ‘NH’.
- ’NH’: Normal hearing with intact cochlear compression
- ’HIx’: Hearing impaired without cochlear compression (a=0)
model (str) –
Model parametrization. Default: ‘paulick2024’.
- ’paulick2024’: Current CASP version. Linear: n_gt=2, n_lp=4. Nonlinear: n_gt=2, n_lp=1.
- ’jepsen2008’: Previous version. Both paths: n_gt=3, n_lp=4/3.
learnable (bool) – If True, makes CF-dependent parameters learnable nn.Parameter objects. Default: False.
dtype (dtype) – Data type for computations. Default: torch.float64 (recommended for numerical stability in cascaded IIR filtering).

fc

Center frequencies in Hz, shape (F,).

Type:: torch.Tensor

num_channels

Number of frequency channels F.

Type:: int

fs

Sampling rate in Hz.

Type:: float

subject

Subject type (‘NH’ or ‘HIx’).

Type:: str

model

Model variant (‘paulick2024’ or ‘jepsen2008’).

Type:: str

n_gt_lin, n_gt_nlin

Gammatone filter orders for linear and nonlinear paths.

Type:: int

n_lp_lin, n_lp_nlin

Number of cascaded lowpass filters for linear and nonlinear paths.

Type:: int

CF_lin, CF_nlin

Gammatone center frequencies (Hz), shape (F,). Registered as buffer or parameter.

Type:: torch.Tensor

BW_lin_norm, BW_nlin_norm

Normalized bandwidths (BW/ERB), shape (F,). Registered as buffer or parameter.

Type:: torch.Tensor

g

Linear path gains, shape (F,). Registered as buffer or parameter.

Type:: torch.Tensor

a, b, c

Broken-stick nonlinearity coefficients, shape (F,). Registered as buffer or parameter.

Type:: torch.Tensor

Examples

Basic usage with ERB spacing:

>>> import torch
>>> from torch_amt.common.filterbanks import DRNLFilterbank
>>>
>>> # Create 50-channel DRNL from 250 to 8000 Hz
>>> drnl = DRNLFilterbank((250, 8000), fs=44100, n_channels=50)
>>> print(drnl.num_channels)
50
>>>
>>> # Process 1 second of audio
>>> x = torch.randn(44100)
>>> y = drnl(x)
>>> print(y.shape)
torch.Size([50, 44100])

Batch processing:

>>> x_batch = torch.randn(4, 22050)  # 4 signals, 0.5 seconds
>>> y_batch = drnl(x_batch)
>>> print(y_batch.shape)  # (batch, channels, time)
torch.Size([4, 50, 22050])

Compare NH vs HIx subjects:

>>> drnl_nh = DRNLFilterbank((500, 4000), fs=44100, n_channels=20, subject='NH')
>>> drnl_hi = DRNLFilterbank((500, 4000), fs=44100, n_channels=20, subject='HIx')
>>> print(f"NH compression: a={drnl_nh.a[10]:.1f}")
NH compression: a=10234.5
>>> print(f"HIx compression: a={drnl_hi.a[10]:.1f}")
HIx compression: a=0.0

Notes

Dual-Path Architecture:

The DRNL combines two processing paths:

Linear Path: \(y_{\text{lin}} = \text{LP}(\text{GT}(g \cdot x))\)
- Gain \(g\) controls overall linear path amplitude
- Gammatone bandpass (order n_gt_lin, CF_lin, BW_lin)
- Cascaded 2nd-order Butterworth lowpass (n_lp_lin times, cutoff LP_lin)
Nonlinear Path: \(y_{\text{nlin}} = \text{LP}(\text{GT}(\text{NL}(\text{GT}(x))))\)
- First gammatone extracts frequency channel
- Broken-stick nonlinearity: \(f(x) = \text{sign}(x) \cdot \min(a|x|, b|x|^c)\)
- Second gammatone re-filters after nonlinearity
- Cascaded lowpass smoothing
Summation: \(y_{\text{total}} = y_{\text{lin}} + y_{\text{nlin}}\)

CF-Dependent Parameters:

All filter parameters scale with center frequency CF following empirical fits:

Linear: \(CF_{\text{lin}} = 10^{-0.068+1.017\log_{10}CF}\)
Gain: \(g = 10^{4.204-0.479\log_{10}CF}\)
Nonlinear CF: \(CF_{\text{nlin}} = 10^{-0.053+1.017\log_{10}CF}\)
NH compression: \(a = 10^{1.403+0.819\log_{10}CF}\) for CF ≤ 1000 Hz

For CF > 1000 Hz, parameters freeze at their 1500 Hz values.

Broken-Stick Nonlinearity:

The nonlinearity provides level-dependent compression:

\[f(x) = \text{sign}(x) \cdot \min(a|x|, b|x|^c)\]

Low levels: \(a|x|\) term dominates (approximately linear)
High levels: \(b|x|^c\) term dominates (compressive, c ≈ 0.25)
Transition at: \(|x| = (a/b)^{1/(c-1)}\)

For NH subjects, this creates ~4:1 compression at high levels. For HIx subjects, a=0 removes the linear term, leaving only compression.

Computational Complexity:

Filter coefficient computation: O(F) at initialization
Forward pass: O(BFT(n_gt + n_lp)) where B=batch, F=channels, T=samples
Sequential filtering prevents full GPU vectorization across channels

See also

GammatoneFilterbank: Single-path gammatone filterbank without nonlinearity
erbspacebw: Generate ERB-spaced center frequencies
audfiltbw: Calculate auditory filter bandwidth (ERB)

References

__init__(fc, fs, n_channels=50, subject='NH', model='paulick2024', learnable=False, dtype=torch.float64)[source]

Initialize Dual Resonance Non-Linear filterbank.

Sets up DRNL filterbank with specified center frequencies, computes CF-dependent parameters (gains, bandwidths, nonlinearity coefficients), and precomputes filter coefficients for efficient processing.

Parameters:

fc (Tensor | Tuple[float, float]) –
Center frequencies specification:
- If torch.Tensor: Explicit center frequencies in Hz, shape (n,). User has full control over frequency spacing and count.
- If tuple (flow, fhigh): Auto-generate n_channels ERB-spaced frequencies between flow and fhigh Hz using erbspacebw().
fs (float) – Sampling rate in Hz. Typical values: 44100, 48000, or 32000. Must be at least 2x the highest center frequency to avoid aliasing.
n_channels (int) – Number of frequency channels. Default: 50. Only used when fc is a tuple. Determines ERB spacing resolution: more channels = finer frequency resolution, higher computational cost.
subject (str) –
Simulated subject type. Default: ‘NH’ (Normal Hearing).
- ’NH’: Normal hearing with intact cochlear compression. Nonlinearity parameters (a, b, c) computed from Paulick et al. (2024). At CF <= 1000 Hz: frequency-dependent compression. Above 1000 Hz: frozen at 1500 Hz parameters.
- ’HIx’: Hearing impaired without cochlear compression. Linear nonlinearity (a=0, removes compression term). Models outer hair cell dysfunction.
model (str) –
Model parametrization variant. Default: ‘paulick2024’.
- ’paulick2024’: Current CASP version (Paulick et al. 2024). Linear path: 2 cascaded gammatone + 4 cascaded lowpass. Nonlinear path: 2 cascaded gammatone + 1 lowpass.
- ’jepsen2008’: Previous version (Jepsen et al. 2008). Linear path: 3 cascaded gammatone + 4 cascaded lowpass. Nonlinear path: 3 cascaded gammatone + 3 cascaded lowpass.
learnable (bool) – If True, makes CF-dependent parameters (CF_lin, BW_lin, g, CF_nlin, BW_nlin, a, b, c) into nn.Parameter for gradient-based optimization. Default: False (parameters are buffers, not optimized).
dtype (dtype) – Data type for computations and parameters. Default: torch.float64. Double precision recommended for numerical stability in cascaded IIR filtering (gammatone and Butterworth filters).

Raises:

ValueError – If model is not ‘paulick2024’ or ‘jepsen2008’.
ValueError – If subject is not ‘NH’ or ‘HIx’.

Notes

ERB Spacing:

When fc is a tuple, frequencies are spaced according to the Equivalent Rectangular Bandwidth (ERB) scale:

\[\text{ERB}(f) = 24.7 (4.37 f / 1000 + 1)\]

This spacing matches the frequency resolution of the human auditory system, with finer spacing at low frequencies and coarser spacing at high frequencies.

Parameter Initialization:

During initialization, the following steps occur:

Generate or validate center frequencies (self.fc)
Compute CF-dependent parameters (_compute_parameters): - Linear path: CF_lin, BW_lin, LP_lin_cutoff, g - Nonlinear path: CF_nlin, BW_nlin, LP_nlin_cutoff, a, b, c
Precompute filter coefficients (_compute_filter_coefficients): - Gammatone filters (complex-valued IIR) - Butterworth lowpass filters (real-valued IIR)
Store coefficients as lists for scipy.signal.lfilter compatibility

Computational Cost:

Filter coefficient computation is O(n_channels), performed once at initialization. Forward pass cost is O(n_channels * batch * time), dominated by sequential filtering operations.

Model Variants:

The two model variants differ in filter orders, affecting frequency selectivity and computational cost:

Higher orders = sharper frequency tuning but more cascaded filtering.

forward(x)[source]

Apply DRNL filterbank to input signal.

Processes audio through dual-path (linear + nonlinear) DRNL filterbank, computing basilar membrane velocity response for each frequency channel. Each channel applies independent linear and nonlinear processing, then sums the paths.

Parameters:

x (Tensor) –

Input audio signal. Shape: (batch, samples) or (samples,).

If 1D: Treated as single-channel audio, output is (F, T)
If 2D: Batch processing, output is (B, F, T)

Returns:

Basilar membrane velocity response. Shape: (batch, channels, samples) or (channels, samples).

Channels (F): Number of frequency channels (self.num_channels)
Samples (T): Same length as input

Return type:

Notes

Algorithm Overview:

For each frequency channel and batch, the DRNL applies:

Linear Path:
1. Multiply input by gain \(g\): \(y_{\text{lin}} = g \cdot x\)
2. Gammatone filter (n_gt_lin cascades): bandpass at \(CF_{\text{lin}}\)
3. Lowpass filter (n_lp_lin cascades): cutoff at \(LP_{\text{lin}}\)
Nonlinear Path:
1. Gammatone filter (n_gt_nlin cascades): bandpass at \(CF_{\text{nlin}}\)
2. Broken-stick nonlinearity:
  
  \[\begin{split}y = \\text{sign}(x) \\cdot \\min(a|x|, b|x|^c)\end{split}\]
3. Gammatone filter again (same parameters)
4. Lowpass filter (n_lp_nlin cascades): cutoff at \(LP_{\text{nlin}}\)
Summation: \(y_{\text{total}} = y_{\text{lin}} + y_{\text{nlin}}\)

Broken-Stick Nonlinearity:

The nonlinearity models cochlear compression:

\[\begin{split}f(x) = \\text{sign}(x) \\cdot \\min(a|x|, b|x|^c)\end{split}\]

Linear regime (\(a|x|\) term): Dominates at low levels
Compressive regime (\(b|x|^c\) term): Dominates at high levels, with \(c \\approx 0.25\) providing ~4:1 compression
Transition: Occurs at \(|x| = (a/b)^{1/(c-1)}\)

For normal hearing (NH), this creates level-dependent gain: high gain for quiet sounds, reduced gain for loud sounds. For hearing impaired (HIx), \(a=0\) removes the linear term, leaving only compression.

Computational Complexity:

For input of length T samples, B batches, F channels:

Filtering operations: O(B * F * T * (n_gt + n_lp)) per path
Nonlinearity: O(B * F * T)
Total: O(B * F * T * max_filter_order)
For B=2, F=50, T=44100, n_gt=2, n_lp=4: ~26M operations (~200 ms on CPU)

See also

_compute_parameters: CF-dependent parameter computation
_compute_filter_coefficients: Gammatone and Butterworth coefficient computation
_gammatone_coeffs: Complex gammatone filter design

extra_repr()[source]

Extra representation string for module printing.

Returns:: String containing key module parameters: fs, num_channels, frequency range, learnable status.
Return type:: str

T_destination = ~T_destination

add_module(name, module)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (str): name of the child module. The child module can be: accessed from this module using the given name

module (Module): child module to be added to the module.

Parameters:

name (str)
module (Module | None)

Return type:

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Args:: fn (Module -> None): function to be applied to each submodule
Returns:: Module: self

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Parameters:: fn (Callable[[Module], None])
Return type:: Self

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

buffers(recurse=True)

Return an iterator over module buffers.

Args:

recurse (bool): if True, then yields buffers of this module: and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor: module buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Tensor]

call_super_init: bool = False

children()

Return an iterator over immediate children modules.

Return type:: Iterator[Module]

Yields:: Module: a child module

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

cuda(device=None)

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Args:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

double()

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

dump_patches: bool = False

eval()

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Return type:: Self

Returns:: Module: self

float()

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the buffer: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.Tensor: The buffer referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not a buffer

Parameters:: target (str)
Return type:: Tensor

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Return type:: Any

Returns:: object: Any extra state to store in the module’s state_dict

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the Parameter: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.nn.Parameter: The Parameter referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not an nn.Parameter

Parameters:: target (str)
Return type:: Parameter

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

Returns:

torch.nn.Module: The submodule referenced by target

Raises:

AttributeError: If at any point along the path resulting from: the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:: target (str)
Return type:: Module

half()

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

ipu(device=None)

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

load_state_dict(state_dict, strict=True, assign=False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Args:

state_dict (dict): a dict containing parameters and: persistent buffers.
strict (bool, optional): whether to strictly enforce that the keys: in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional): When set to False, the properties of the tensors: in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

NamedTuple with missing_keys and unexpected_keys fields:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Note:

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

Parameters:

state_dict (Mapping[str, Any])
strict (bool)
assign (bool)

modules()

Return an iterator over all modules in the network.

Return type:: Iterator[Module]

Yields:: Module: a module in the network
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Args:

prefix (str): prefix to prepend to all buffer names. recurse (bool, optional): if True, then yields buffers of this module

and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor): Tuple containing the name and buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Return type:: Iterator[tuple[str, Module]]

Yields:: (str, Module): Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Args:: memo: a memo to store the set of modules already added to the result prefix: a prefix that will be added to the name of the module remove_duplicate: whether to remove the duplicated module instances in the result

or not
Yields:: (str, Module): Tuple of name and module
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

Parameters:

memo (set[Module] | None)
prefix (str)
remove_duplicate (bool)

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Args:

prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module

and all submodules. Otherwise, yields only parameters that are direct members of this module.

remove_duplicate (bool, optional): whether to remove the duplicated: parameters in the result. Defaults to True.

Yields:

(str, Parameter): Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Args:

recurse (bool): if True, then yields parameters of this module: and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter: module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Parameter]

register_backward_hook(hook)

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

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

Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
Return type:: RemovableHandle

register_buffer(name, tensor, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (str): name of the buffer. The buffer can be accessed: from this module using the given name
tensor (Tensor or None): buffer to be registered. If None, then operations: that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool): whether the buffer is part of this module’s: state_dict.

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

Parameters:

name (str)
tensor (Tensor | None)
persistent (bool)

Return type:

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Args:

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

before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If True, the hook will be passed the: kwargs given to the forward function. Default: False
always_call (bool): If True the hook will be run regardless of: whether an exception is raised while calling the Module. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...], Any], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any], Any], Any | None])
prepend (bool)
with_kwargs (bool)
always_call (bool)

Return type:

RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Args:

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

all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If true, the hook will be passed the kwargs: given to the forward function. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...]], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None])
prepend (bool)
with_kwargs (bool)

Return type:

RemovableHandle

register_full_backward_hook(hook, prepend=False)

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

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

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Arguments:

hook (Callable): Callable hook that will be invoked before: loading the state dict.

register_module(name, module)

Alias for add_module().

Parameters:

name (str)
module (Module | None)

Return type:

register_parameter(name, param)

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (str): name of the parameter. The parameter can be accessed: from this module using the given name
param (Parameter or None): parameter to be added to the module. If: None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Parameters:

name (str)
param (Parameter | None)

Return type:

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Args:

requires_grad (bool): whether autograd should record operations on: parameters in this module. Default: True.

Returns:

Module: self

Parameters:: requires_grad (bool)
Return type:: Self

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Args:: state (dict): Extra state from the state_dict

Parameters:: state (Any)
Return type:: None

set_submodule(target, module, strict=False)

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

module: The module to set the submodule to. strict: If False, the method will replace an existing submodule

or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError: If the target string is empty or if module is not an instance of nn.Module. AttributeError: If at any point along the path resulting from

the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:

target (str)
module (Module)
strict (bool)

Return type:

share_memory()

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)

Overloads:

self, destination (T_destination), prefix (str), keep_vars (bool) → T_destination
self, prefix (str), keep_vars (bool) → dict[str, Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Args:

destination (dict, optional): If provided, the state of module will: be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional): a prefix added to parameter and buffer: names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional): by default the Tensor s: returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

dict:: a dictionary containing a whole state of the module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)

Overloads:

self, device (DeviceLikeType | None), dtype (dtype | None), non_blocking (bool) → Self
self, dtype (dtype), non_blocking (bool) → Self
self, tensor (Tensor), non_blocking (bool) → Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters: and buffers in this module
dtype (torch.dtype): the desired floating point or complex dtype of: the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired: dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format): the desired memory: format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)

Move the parameters and buffers to the specified device without copying storage.

Args:

device (torch.device): The desired device of the parameters: and buffers in this module.
recurse (bool): Whether parameters and buffers of submodules should: be recursively moved to the specified device.

Returns:

Module: self

Parameters:

device (str | device | int | None)
recurse (bool)

Return type:

Self

train(mode=True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation: mode (False). Default: True.

Returns:

Module: self

Parameters:: mode (bool)
Return type:: Self

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Parameters:: dst_type (dtype | str)
Return type:: Self

xpu(device=None)

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

zero_grad(set_to_none=True)

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: See torch.optim.Optimizer.zero_grad() for details.

Parameters:: set_to_none (bool)
Return type:: None

training: bool

FastDRNLFilterbank

class torch_amt.FastDRNLFilterbank(*args, ir_length=4096, **kwargs)[source]

Bases: DRNLFilterbank

Fast DRNL filterbank using FFT-based convolution with pre-computed impulse responses.

This is a performance-optimized drop-in replacement for DRNLFilterbank that achieves ~100x speedup while maintaining excellent numerical accuracy (max absolute diff < 3e-4).

Key Optimizations:

Pre-compute impulse responses from all IIR filters at initialization
Use torch.nn.functional.conv1d() with groups for fully parallel convolution
Zero loops in forward pass - all operations vectorized
Pre-flip impulse responses (conv1d does cross-correlation)

Performance Characteristics:

Speedup: ~100x vs original DRNLFilterbank
Throughput: ~0.5x realtime on CPU (vs ~0.005x for original)
Accuracy: max absolute diff < 3e-4, mean diff < 1e-5
Memory: ~5 MB for IR storage (50 channels, ir_length=4096)
Batch scaling: Best with batch size ≤ 4, degrades at batch=8

Training Considerations:

Filter parameters (CF, BW, LP_cutoff) are frozen (not trainable) because impulse responses are pre-computed. Only nonlinearity parameters (a, b, c) remain trainable. This is a design trade-off for performance.

If you need to train filter parameters, use the original DRNLFilterbank.

Parameters:

fc (torch.Tensor or tuple of float) –
Center frequencies. Same as DRNLFilterbank.
- torch.Tensor: Explicit center frequencies of shape \((F,)\)
- tuple (flow, fhigh): Generates ERB-spaced frequencies
fs (float) – Sampling rate in Hz.
n_channels (int, optional) – Number of frequency channels. Only used if fc is tuple. Default: 50.
ir_length (int) – Length of impulse responses in samples. Longer = more accurate but slower and more memory. Default: 4096 (good trade-off).
learnable (bool, optional) – If True, nonlinearity parameters (a, b, c) are learnable. Filter parameters remain frozen. Default: False.
dtype (torch.dtype, optional) – Data type. Default: torch.float32.

ir_length

Length of pre-computed impulse responses in samples.

Type:: int

ir_lin

Pre-computed impulse responses for linear path, shape \((F, L)\) where F is number of channels and L is ir_length.

Type:: torch.Tensor

ir_nlin_1

Pre-computed impulse responses for first nonlinear gammatone (before nonlinearity), shape \((F, L)\).

Type:: torch.Tensor

ir_nlin_2

Pre-computed impulse responses for second nonlinear gammatone (after nonlinearity), shape \((F, L)\).

Type:: torch.Tensor

ir_lp_nlin

Pre-computed impulse responses for cascaded lowpass in nonlinear path, shape \((F, L)\).

Type:: torch.Tensor

num_channels

Number of frequency channels (inherited from parent).

Type:: int

fs

Sampling rate in Hz (inherited from parent).

Type:: float

a

Nonlinearity compression exponent, shape \((F,)\). Learnable if learnable=True.

Type:: torch.Tensor or nn.Parameter

b

Nonlinearity scaling factor, shape \((F,)\). Learnable if learnable=True.

Type:: torch.Tensor or nn.Parameter

c

Nonlinearity offset, shape \((F,)\). Learnable if learnable=True.

Type:: torch.Tensor or nn.Parameter

Shape

-----

- Input

\(B\) = batch size (optional)
\(T\) = time samples

Type:: \((B, T)\) or \((T,)\) where

- Output

\(F\) = number of frequency channels

Type:: \((B, F, T)\) or \((F, T)\) where

Notes

Performance Considerations:

The ~100x speedup is measured relative to the original IIR-based implementation. However, the original is ~200x slower than realtime, so FastDRNLFilterbank is still ~2x slower than realtime on typical CPUs. For real-time applications or very large-scale training, consider:

Using batch size ≤ 4 (batch=8 has poor scaling)
Processing audio in chunks
Using GPU acceleration (if available)
Using a shorter ir_length (e.g., 2048) for speed vs accuracy trade-off

Training Limitations:

Because impulse responses are pre-computed at initialization, filter parameters (center frequencies, bandwidths, lowpass cutoffs) cannot be trained via gradient descent. Only the nonlinearity parameters (a, b, c) remain trainable.

If you need fully trainable filters, use DRNLFilterbank. If you only need to optimize the nonlinearity while keeping filters fixed, FastDRNLFilterbank is the better choice.

Numerical Accuracy:

The FFT-based convolution introduces small numerical differences compared to the original IIR implementation:

Maximum absolute difference: ~3e-4
Mean absolute difference: ~1e-5
Relative error: < 0.01%

These differences are negligible for most applications and are well below numerical precision limits of auditory models.

See also

DRNLFilterbank: Original implementation (slower but fully trainable)
GammatoneFilterbank: Gammatone filterbank (faster but linear only)

Examples

Drop-in replacement for DRNLFilterbank:

>>> import torch
>>> from torch_amt.common.filterbanks import FastDRNLFilterbank, DRNLFilterbank
>>>
>>> # Original (slow)
>>> drnl = DRNLFilterbank((250, 8000), fs=44100, n_channels=50)
>>>
>>> # Fast version (100x speedup)
>>> drnl_fast = FastDRNLFilterbank((250, 8000), fs=44100, n_channels=50)
>>>
>>> x = torch.randn(1, 22050)  # 0.5s @ 44.1kHz
>>> y = drnl_fast(x)  # [1, 50, 22050]
>>> print(f"Output shape: {y.shape}")
Output shape: torch.Size([1, 50, 22050])

With trainable nonlinearity parameters:

>>> drnl = FastDRNLFilterbank((250, 8000), fs=44100, n_channels=50, learnable=True)
>>> optimizer = torch.optim.Adam(drnl.parameters(), lr=1e-3)
>>>
>>> # Only a, b, c will be updated (3*50 = 150 parameters)
>>> print(f"Trainable parameters: {sum(p.numel() for p in drnl.parameters())}")
Trainable parameters: 150
>>>
>>> y = drnl(x)
>>> loss = criterion(y, target)
>>> loss.backward()
>>> optimizer.step()

Adjusting IR length for speed/accuracy trade-off:

>>> # Shorter IR = faster but less accurate
>>> drnl_short = FastDRNLFilterbank((250, 8000), fs=44100, ir_length=2048)
>>>
>>> # Longer IR = slower but more accurate
>>> drnl_long = FastDRNLFilterbank((250, 8000), fs=44100, ir_length=8192)

References

__init__(*args, ir_length=4096, **kwargs)[source]

Initialize Dual Resonance Non-Linear filterbank.

Sets up DRNL filterbank with specified center frequencies, computes CF-dependent parameters (gains, bandwidths, nonlinearity coefficients), and precomputes filter coefficients for efficient processing.

Parameters:

fc (torch.Tensor or tuple of (float, float)) –
Center frequencies specification:
- If torch.Tensor: Explicit center frequencies in Hz, shape (n,). User has full control over frequency spacing and count.
- If tuple (flow, fhigh): Auto-generate n_channels ERB-spaced frequencies between flow and fhigh Hz using erbspacebw().
fs (float) – Sampling rate in Hz. Typical values: 44100, 48000, or 32000. Must be at least 2x the highest center frequency to avoid aliasing.
n_channels (int, optional) – Number of frequency channels. Default: 50. Only used when fc is a tuple. Determines ERB spacing resolution: more channels = finer frequency resolution, higher computational cost.
subject (str, optional) –
Simulated subject type. Default: ‘NH’ (Normal Hearing).
- ’NH’: Normal hearing with intact cochlear compression. Nonlinearity parameters (a, b, c) computed from Paulick et al. (2024). At CF <= 1000 Hz: frequency-dependent compression. Above 1000 Hz: frozen at 1500 Hz parameters.
- ’HIx’: Hearing impaired without cochlear compression. Linear nonlinearity (a=0, removes compression term). Models outer hair cell dysfunction.
model (str, optional) –
Model parametrization variant. Default: ‘paulick2024’.
- ’paulick2024’: Current CASP version (Paulick et al. 2024). Linear path: 2 cascaded gammatone + 4 cascaded lowpass. Nonlinear path: 2 cascaded gammatone + 1 lowpass.
- ’jepsen2008’: Previous version (Jepsen et al. 2008). Linear path: 3 cascaded gammatone + 4 cascaded lowpass. Nonlinear path: 3 cascaded gammatone + 3 cascaded lowpass.
learnable (bool, optional) – If True, makes CF-dependent parameters (CF_lin, BW_lin, g, CF_nlin, BW_nlin, a, b, c) into nn.Parameter for gradient-based optimization. Default: False (parameters are buffers, not optimized).
dtype (torch.dtype, optional) – Data type for computations and parameters. Default: torch.float64. Double precision recommended for numerical stability in cascaded IIR filtering (gammatone and Butterworth filters).

Raises:

ValueError – If model is not ‘paulick2024’ or ‘jepsen2008’.
ValueError – If subject is not ‘NH’ or ‘HIx’.

Notes

ERB Spacing:

When fc is a tuple, frequencies are spaced according to the Equivalent Rectangular Bandwidth (ERB) scale:

\[\text{ERB}(f) = 24.7 (4.37 f / 1000 + 1)\]

This spacing matches the frequency resolution of the human auditory system, with finer spacing at low frequencies and coarser spacing at high frequencies.

Parameter Initialization:

During initialization, the following steps occur:

Generate or validate center frequencies (self.fc)
Compute CF-dependent parameters (_compute_parameters): - Linear path: CF_lin, BW_lin, LP_lin_cutoff, g - Nonlinear path: CF_nlin, BW_nlin, LP_nlin_cutoff, a, b, c
Precompute filter coefficients (_compute_filter_coefficients): - Gammatone filters (complex-valued IIR) - Butterworth lowpass filters (real-valued IIR)
Store coefficients as lists for scipy.signal.lfilter compatibility

Computational Cost:

Filter coefficient computation is O(n_channels), performed once at initialization. Forward pass cost is O(n_channels * batch * time), dominated by sequential filtering operations.

Model Variants:

The two model variants differ in filter orders, affecting frequency selectivity and computational cost:

Higher orders = sharper frequency tuning but more cascaded filtering.

Parameters:: ir_length (int)

forward(x)[source]

Apply fast DRNL filterbank using FFT-based convolution.

Parameters:: x (Tensor) – Input signal, shape (B, T) or (T,).
Returns:: Filtered output, shape (B, F, T) or (F, T).
Return type:: Tensor

T_destination = ~T_destination

add_module(name, module)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (str): name of the child module. The child module can be: accessed from this module using the given name

module (Module): child module to be added to the module.

Parameters:

name (str)
module (Module | None)

Return type:

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Args:: fn (Module -> None): function to be applied to each submodule
Returns:: Module: self

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Parameters:: fn (Callable[[Module], None])
Return type:: Self

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

buffers(recurse=True)

Return an iterator over module buffers.

Args:

recurse (bool): if True, then yields buffers of this module: and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor: module buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Tensor]

call_super_init: bool = False

children()

Return an iterator over immediate children modules.

Return type:: Iterator[Module]

Yields:: Module: a child module

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

cuda(device=None)

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Args:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

double()

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

dump_patches: bool = False

eval()

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Return type:: Self

Returns:: Module: self

extra_repr()

Extra representation string for module printing.

Returns:: String containing key module parameters: fs, num_channels, frequency range, learnable status.
Return type:: str

float()

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the buffer: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.Tensor: The buffer referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not a buffer

Parameters:: target (str)
Return type:: Tensor

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Return type:: Any

Returns:: object: Any extra state to store in the module’s state_dict

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the Parameter: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.nn.Parameter: The Parameter referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not an nn.Parameter

Parameters:: target (str)
Return type:: Parameter

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

Returns:

torch.nn.Module: The submodule referenced by target

Raises:

AttributeError: If at any point along the path resulting from: the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:: target (str)
Return type:: Module

half()

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

ipu(device=None)

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

load_state_dict(state_dict, strict=True, assign=False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Args:

state_dict (dict): a dict containing parameters and: persistent buffers.
strict (bool, optional): whether to strictly enforce that the keys: in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional): When set to False, the properties of the tensors: in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

NamedTuple with missing_keys and unexpected_keys fields:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Note:

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

Parameters:

state_dict (Mapping[str, Any])
strict (bool)
assign (bool)

modules()

Return an iterator over all modules in the network.

Return type:: Iterator[Module]

Yields:: Module: a module in the network
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Args:

prefix (str): prefix to prepend to all buffer names. recurse (bool, optional): if True, then yields buffers of this module

and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor): Tuple containing the name and buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Return type:: Iterator[tuple[str, Module]]

Yields:: (str, Module): Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Args:: memo: a memo to store the set of modules already added to the result prefix: a prefix that will be added to the name of the module remove_duplicate: whether to remove the duplicated module instances in the result

or not
Yields:: (str, Module): Tuple of name and module
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

Parameters:

memo (set[Module] | None)
prefix (str)
remove_duplicate (bool)

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Args:

prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module

and all submodules. Otherwise, yields only parameters that are direct members of this module.

remove_duplicate (bool, optional): whether to remove the duplicated: parameters in the result. Defaults to True.

Yields:

(str, Parameter): Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Args:

recurse (bool): if True, then yields parameters of this module: and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter: module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Parameter]

register_backward_hook(hook)

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

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

Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
Return type:: RemovableHandle

register_buffer(name, tensor, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (str): name of the buffer. The buffer can be accessed: from this module using the given name
tensor (Tensor or None): buffer to be registered. If None, then operations: that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool): whether the buffer is part of this module’s: state_dict.

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

Parameters:

name (str)
tensor (Tensor | None)
persistent (bool)

Return type:

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Args:

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

before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If True, the hook will be passed the: kwargs given to the forward function. Default: False
always_call (bool): If True the hook will be run regardless of: whether an exception is raised while calling the Module. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...], Any], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any], Any], Any | None])
prepend (bool)
with_kwargs (bool)
always_call (bool)

Return type:

RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Args:

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

all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If true, the hook will be passed the kwargs: given to the forward function. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...]], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None])
prepend (bool)
with_kwargs (bool)

Return type:

RemovableHandle

register_full_backward_hook(hook, prepend=False)

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

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

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Arguments:

hook (Callable): Callable hook that will be invoked before: loading the state dict.

register_module(name, module)

Alias for add_module().

Parameters:

name (str)
module (Module | None)

Return type:

register_parameter(name, param)

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (str): name of the parameter. The parameter can be accessed: from this module using the given name
param (Parameter or None): parameter to be added to the module. If: None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Parameters:

name (str)
param (Parameter | None)

Return type:

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Args:

requires_grad (bool): whether autograd should record operations on: parameters in this module. Default: True.

Returns:

Module: self

Parameters:: requires_grad (bool)
Return type:: Self

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Args:: state (dict): Extra state from the state_dict

Parameters:: state (Any)
Return type:: None

set_submodule(target, module, strict=False)

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

module: The module to set the submodule to. strict: If False, the method will replace an existing submodule

or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError: If the target string is empty or if module is not an instance of nn.Module. AttributeError: If at any point along the path resulting from

the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:

target (str)
module (Module)
strict (bool)

Return type:

share_memory()

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)

Overloads:

self, destination (T_destination), prefix (str), keep_vars (bool) → T_destination
self, prefix (str), keep_vars (bool) → dict[str, Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Args:

destination (dict, optional): If provided, the state of module will: be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional): a prefix added to parameter and buffer: names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional): by default the Tensor s: returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

dict:: a dictionary containing a whole state of the module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)

Overloads:

self, device (DeviceLikeType | None), dtype (dtype | None), non_blocking (bool) → Self
self, dtype (dtype), non_blocking (bool) → Self
self, tensor (Tensor), non_blocking (bool) → Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters: and buffers in this module
dtype (torch.dtype): the desired floating point or complex dtype of: the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired: dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format): the desired memory: format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)

Move the parameters and buffers to the specified device without copying storage.

Args:

device (torch.device): The desired device of the parameters: and buffers in this module.
recurse (bool): Whether parameters and buffers of submodules should: be recursively moved to the specified device.

Returns:

Module: self

Parameters:

device (str | device | int | None)
recurse (bool)

Return type:

Self

train(mode=True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation: mode (False). Default: True.

Returns:

Module: self

Parameters:: mode (bool)
Return type:: Self

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Parameters:: dst_type (dtype | str)
Return type:: Self

xpu(device=None)

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

zero_grad(set_to_none=True)

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: See torch.optim.Optimizer.zero_grad() for details.

Parameters:: set_to_none (bool)
Return type:: None

training: bool

ExcitationPattern

class torch_amt.ExcitationPattern(fs=32000, f_min=50.0, f_max=15000.0, erb_step=0.25, learnable=False)[source]

Bases: Module

Excitation pattern with asymmetric spreading for Glasberg & Moore (2002) loudness model.

Applies frequency-domain spreading to simulate the spread of excitation along the basilar membrane. The spreading function models how energy at one frequency activates adjacent auditory filters, with asymmetric and level-dependent characteristics:

Asymmetric: More spreading toward lower frequencies (shallower slope) than higher frequencies (steeper slope)
Level-dependent: Slopes decrease with increasing level (more spreading at high levels)

This is the second stage of the Glasberg & Moore (2002) loudness model, applied after ERB integration (ERBIntegration) to compute the excitation pattern from the excitation in individual ERB bands.

Parameters:

fs (int) – Sampling rate in Hz. Default: 32000. Used for consistency with ERBIntegration but not directly used in spreading computation.
f_min (float) – Minimum center frequency in Hz. Default: 50.0. Defines lower bound of ERB channel range.
f_max (float) – Maximum center frequency in Hz. Default: 15000.0. Defines upper bound of ERB channel range.
erb_step (float) – ERB-rate spacing step. Default: 0.25. Determines ERB channel resolution. Must match ERBIntegration for consistent processing.
learnable (bool) – If True, spreading slope parameters become learnable nn.Parameter objects (4 parameters: upper/lower base slopes, upper/lower level dependencies). Default: False (fixed slopes from Moore & Glasberg 1987).

fs

Sampling rate in Hz.

Type:: int

f_min

Minimum center frequency in Hz.

Type:: float

f_max

Maximum center frequency in Hz.

Type:: float

erb_step

ERB-rate spacing step.

Type:: float

learnable

Whether spreading slopes are learnable.

Type:: bool

fc_erb

ERB channel center frequencies in Hz, shape (n_erb_bands,). Registered as buffer. Matches ERBIntegration channels.

Type:: torch.Tensor

erb_centers

ERB-rate values for each channel, shape (n_erb_bands,). Registered as buffer. Used for distance computation in spreading.

Type:: torch.Tensor

n_erb_bands

Number of ERB channels. Typically 150 for default parameters.

Type:: int

upper_slope_base

Base spreading slope toward higher frequencies at 60 dB SPL, in dB/ERB. Default: 27.0 dB/ERB. Higher values = steeper decay, less spreading.

Type:: torch.Tensor or nn.Parameter

lower_slope_base

Base spreading slope toward lower frequencies at 60 dB SPL, in dB/ERB. Default: 11.0 dB/ERB. Lower than upper slope (asymmetry).

Type:: torch.Tensor or nn.Parameter

upper_slope_per_db

Level dependency of upper slope, in (dB/ERB) per dB SPL. Default: -0.37. Negative = slopes decrease at high levels.

Type:: torch.Tensor or nn.Parameter

lower_slope_per_db

Level dependency of lower slope, in (dB/ERB) per dB SPL. Default: -0.20. Less level-dependent than upper slope.

Type:: torch.Tensor or nn.Parameter

level_ref

Reference level for slope computation, in dB SPL. Fixed at 60.0. Registered as buffer.

Type:: torch.Tensor

Examples

Basic usage in loudness model pipeline:

>>> import torch
>>> from torch_amt.common.filterbanks import (
...     MultiResolutionFFT, ERBIntegration, ExcitationPattern
... )
>>>
>>> # Complete loudness model front-end
>>> mrf = MultiResolutionFFT(fs=32000)
>>> erb_int = ERBIntegration(fs=32000)
>>> exc_pattern = ExcitationPattern(fs=32000)
>>>
>>> # Process audio
>>> audio = torch.randn(2, 32000)  # 2 batches, 1 second
>>> psd, freqs = mrf(audio)  # (2, 32, 1025)
>>> excitation = erb_int(psd, freqs)  # (2, 32, 150) in dB SPL
>>> spread = exc_pattern(excitation)  # (2, 32, 150) with spreading
>>>
>>> print(f"Excitation range: {excitation.min():.1f} - {excitation.max():.1f} dB SPL")
Excitation range: 48.3 - 87.2 dB SPL
>>> print(f"Spread range: {spread.min():.1f} - {spread.max():.1f} dB SPL")
Spread range: 54.1 - 87.3 dB SPL

Inspect spreading slopes at different levels:

>>> exc_pattern = ExcitationPattern()
>>>
>>> for level in [40, 60, 80, 100]:
...     upper, lower = exc_pattern.get_spreading_slopes(level)
...     asymmetry = upper / lower
...     print(f"Level {level} dB: upper={upper:.1f}, lower={lower:.1f}, ratio={asymmetry:.2f}:1")
Level 40 dB: upper=34.4, lower=15.0, ratio=2.29:1
Level 60 dB: upper=27.0, lower=11.0, ratio=2.45:1
Level 80 dB: upper=19.6, lower=10.0, ratio=1.96:1
Level 100 dB: upper=12.2, lower=10.0, ratio=1.22:1

Learnable spreading parameters for optimization:

>>> exc_learnable = ExcitationPattern(learnable=True)
>>> print(f"Learnable parameters: {sum(p.numel() for p in exc_learnable.parameters())}")
Learnable parameters: 4
>>>
>>> # Parameter names
>>> for name, param in exc_learnable.named_parameters():
...     print(f"{name}: {param.item():.2f}")
upper_slope_base: 27.00
lower_slope_base: 11.00
upper_slope_per_db: -0.37
lower_slope_per_db: -0.20
>>>
>>> # Can be optimized
>>> optimizer = torch.optim.Adam(exc_learnable.parameters(), lr=0.01)

Effect of spreading on excitation pattern:

>>> # Single tone at 1000 Hz, 70 dB SPL
>>> excitation_single = torch.zeros(1, 1, 150)
>>> # Find ERB channel closest to 1000 Hz
>>> idx_1000 = torch.argmin(torch.abs(exc_pattern.fc_erb - 1000.0))
>>> excitation_single[0, 0, idx_1000] = 70.0  # 70 dB SPL
>>>
>>> spread_single = exc_pattern(excitation_single)
>>>
>>> # Spreading around 1000 Hz channel
>>> print(f"Original: channel {idx_1000} = {excitation_single[0, 0, idx_1000]:.1f} dB")
Original: channel 50 = 70.0 dB
>>> print(f"Spread: channels {idx_1000-2}:{idx_1000+3}")
>>> print(spread_single[0, 0, idx_1000-2:idx_1000+3])
Spread: channels 48:53
tensor([54.3, 60.8, 70.0, 64.5, 56.2])

Notes

Asymmetric Spreading Function:

For each ERB channel \(i\) with excitation level \(E_i\) (in dB SPL), the spreading to channel \(j\) is computed as:

\[\begin{split}\text{Attenuation}_{ij} = \begin{cases} s_u(E_i) \cdot \Delta \text{ERB} & \text{if } \Delta \text{ERB} \geq 0 \\ -s_l(E_i) \cdot \Delta \text{ERB} & \text{if } \Delta \text{ERB} < 0 \end{cases}\end{split}\]

where \(\Delta \text{ERB} = \text{ERB}_j - \text{ERB}_i\) is the distance between channels (positive = higher frequency, negative = lower frequency).

The contribution from channel \(i\) to channel \(j\) is:

\[C_{ij} = E_i - \text{Attenuation}_{ij}\]

Total excitation at channel \(j\) is the log-sum of all contributions:

\[E_j^{\text{spread}} = 10 \log_{10}\left( \sum_{i} 10^{C_{ij}/10} \right)\]

Level-Dependent Slopes:

The spreading slopes depend on excitation level following Moore & Glasberg (1987):

\[\begin{split}s_u(E) &= s_{u,\text{base}} + k_u (E - E_{\text{ref}}) \\ s_l(E) &= s_{l,\text{base}} + k_l (E - E_{\text{ref}})\end{split}\]

where \(E_{\text{ref}} = 60\) dB SPL is the reference level, \(s_{u,\text{base}} = 27\) dB/ERB and \(s_{l,\text{base}} = 11\) dB/ERB are base slopes, and \(k_u = -0.37\), \(k_l = -0.20\) are level dependencies (slopes decrease at high levels -> more spreading).

Slopes are clamped to minimum 10.0 dB/ERB to prevent unrealistic spreading at very high levels.

Asymmetry Rationale:

The asymmetric spreading (\(s_u > s_l\)) reflects physiological properties of the cochlea:

Upward spreading (toward high frequencies): Steeper decay because high-frequency channels are physically distant from excitation site on basilar membrane.
Downward spreading (toward low frequencies): Shallower decay because traveling wave on basilar membrane spreads more gradually toward apex (low-frequency region).

At 60 dB SPL, asymmetry ratio is \(27/11 \approx 2.45:1\).

Computational Complexity:

The triple-nested loop (batch, frames, channels) with inner loop over all channels gives complexity:

Time: O(batch x n_frames x n_erb_bands²)
For (2, 32, 150): ~1.44M operations
Optimization: Skip contributions with attenuation >50 dB (negligible)
Typical runtime: ~50 ms on CPU, ~5 ms on GPU (for 2 batches, 32 frames)

Spreading Limits:

Contributions with attenuation >50 dB are skipped (threshold check) to avoid:

Numerical instability in log-domain computations
Wasted computation on negligible contributions
At 27 dB/ERB upper slope: >50 dB attenuation at ~1.85 ERB distance
At 11 dB/ERB lower slope: >50 dB attenuation at ~4.55 ERB distance

Log-Domain Summation:

Uses torch.logaddexp(a, b) to compute \(\log(e^a + e^b)\) for dB values without converting to linear domain:

\[\text{logaddexp}(a, b) = \log_{10}(10^{a/10} + 10^{b/10}) \cdot 10\]

This prevents numerical overflow/underflow when summing many contributions.

Device Support:

Supports CPU, CUDA, and MPS. All buffers (fc_erb, erb_centers, level_ref) and parameters automatically moved with .to(device).

See also

ERBIntegration: Computes ERB-band excitation (input for this module).
MultiResolutionFFT: Computes PSD (input for ERBIntegration).
f2erbrate: Convert frequency to ERB-rate scale.
erbrate2f: Convert ERB-rate to frequency.

References

__init__(fs=32000, f_min=50.0, f_max=15000.0, erb_step=0.25, learnable=False)[source]

Initialize excitation pattern spreading module.

Parameters:

fs (int) – Sampling rate in Hz. Default: 32000. Used for consistency with ERBIntegration but not directly used in spreading.
f_min (float) – Minimum ERB center frequency in Hz. Default: 50.0. Must match ERBIntegration for consistent ERB channel alignment.
f_max (float) – Maximum ERB center frequency in Hz. Default: 15000.0. Must match ERBIntegration for consistent ERB channel alignment.
erb_step (float) – ERB-rate spacing step. Default: 0.25. Must match ERBIntegration to ensure same ERB channel count (typically 150 channels).
learnable (bool) – If True, creates learnable spreading slope parameters (4 total: upper_slope_base, lower_slope_base, upper_slope_per_db, lower_slope_per_db). Default: False (fixed slopes from Moore & Glasberg 1987).

Notes

ERB Channel Setup:

Initializes same ERB channels as ERBIntegration:

Convert f_min, f_max to ERB-rate: erb_min, erb_max = f2erbrate(...)
Create uniform grid: erb_centers = arange(erb_min, erb_max, erb_step)
Convert to Hz: fc_erb = erbrate2f(erb_centers)

Both erb_centers (ERB-rate) and fc_erb (Hz) are registered as buffers for device compatibility.

Spreading Slope Parameters:

Default values from Moore & Glasberg (1987):

upper_slope_base = 27.0 dB/ERB (toward high freq, at 60 dB SPL)
lower_slope_base = 11.0 dB/ERB (toward low freq, at 60 dB SPL)
upper_slope_per_db = -0.37 (dB/ERB) per dB (level dependency)
lower_slope_per_db = -0.20 (dB/ERB) per dB (level dependency)

Negative level dependencies mean slopes decrease with increasing level (more spreading at high levels).

Reference Level:

Fixed at 60.0 dB SPL (level_ref), registered as buffer. This is the reference point for level-dependent slope adjustments.

forward(excitation)[source]

Apply excitation pattern spreading to ERB-band excitation.

Parameters:: excitation (Tensor) – Excitation in dB SPL, shape (batch, n_frames, n_erb_bands). Typically output from ERBIntegration.
Returns:: Spread excitation in dB SPL, shape (batch, n_frames, n_erb_bands). Represents excitation pattern after accounting for frequency spreading along the basilar membrane.
Return type:: Tensor

Notes

This is a simple wrapper around _compute_spreading_function() that provides the public API for the module.

Usage in Loudness Model:

The complete Glasberg & Moore (2002) loudness model front-end:

MultiResolutionFFT: audio -> PSD + frequencies
ERBIntegration: PSD -> ERB-band excitation
ExcitationPattern: excitation -> spread excitation (this method)
Specific loudness: spread excitation -> specific loudness (model-dependent)
Total loudness: integrate specific loudness over ERB bands

Examples

>>> import torch
>>> from torch_amt.common.filterbanks import ExcitationPattern
>>>
>>> # Excitation from ERBIntegration (2 batches, 32 frames, 150 ERB bands)
>>> excitation = torch.randn(2, 32, 150) * 10 + 60  # ~60 dB SPL
>>>
>>> exc_pattern = ExcitationPattern()
>>> spread = exc_pattern(excitation)
>>>
>>> print(f\"Input range: {excitation.min():.1f} - {excitation.max():.1f} dB SPL\")
Input range: 37.2 - 82.8 dB SPL
>>> print(f\"Spread range: {spread.min():.1f} - {spread.max():.1f} dB SPL\")
Spread range: 43.5 - 82.9 dB SPL

get_spreading_slopes(level_db)[source]

Get spreading slopes for a specific excitation level.

Computes level-dependent spreading slopes according to Moore & Glasberg (1987).

Parameters:

level_db (float) – Excitation level in dB SPL.

Return type:

Tuple[float, float]

Returns:

upper_slope (float) – Spreading slope toward higher frequencies, in dB/ERB. Higher values = steeper decay, less spreading.
lower_slope (float) – Spreading slope toward lower frequencies, in dB/ERB. Lower values = shallower decay, more spreading.

Notes

Level-Dependent Formula:

\[\begin{split}s_u(E) &= \\max(10, 27 - 0.37(E - 60)) \\\\ s_l(E) &= \\max(10, 11 - 0.20(E - 60))\end{split}\]

where \(E\) is level in dB SPL.

Asymmetry:

At 60 dB SPL: \(s_u/s_l = 27/11 \\approx 2.45:1\)

The upper slope is always steeper than the lower slope, reflecting the physiological asymmetry of basilar membrane vibration patterns.

Level Dependency:

Both slopes decrease with increasing level (negative coefficients), meaning more spreading at high levels. This matches auditory filter tuning measurements showing broader filters at high SPLs.

Examples

>>> from torch_amt.common.filterbanks import ExcitationPattern
>>> exc_pattern = ExcitationPattern()
>>>
>>> # Compare slopes at different levels
>>> for level in [40, 60, 80, 100]:
...     upper, lower = exc_pattern.get_spreading_slopes(level)
...     print(f"{level} dB SPL: upper={upper:.1f}, lower={lower:.1f}, ratio={upper/lower:.2f}:1")
40 dB SPL: upper=34.4, lower=15.0, ratio=2.29:1
60 dB SPL: upper=27.0, lower=11.0, ratio=2.45:1
80 dB SPL: upper=19.6, lower=10.0, ratio=1.96:1
100 dB SPL: upper=12.2, lower=10.0, ratio=1.22:1

extra_repr()[source]

Extra representation string for module printing.

Returns:: String summarizing module parameters: sampling rate, frequency range, ERB spacing, number of channels, and learnable status.
Return type:: str

T_destination = ~T_destination

add_module(name, module)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (str): name of the child module. The child module can be: accessed from this module using the given name

module (Module): child module to be added to the module.

Parameters:

name (str)
module (Module | None)

Return type:

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Args:: fn (Module -> None): function to be applied to each submodule
Returns:: Module: self

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Parameters:: fn (Callable[[Module], None])
Return type:: Self

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

buffers(recurse=True)

Return an iterator over module buffers.

Args:

recurse (bool): if True, then yields buffers of this module: and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor: module buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Tensor]

call_super_init: bool = False

children()

Return an iterator over immediate children modules.

Return type:: Iterator[Module]

Yields:: Module: a child module

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

cuda(device=None)

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Args:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

double()

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

dump_patches: bool = False

eval()

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Return type:: Self

Returns:: Module: self

float()

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the buffer: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.Tensor: The buffer referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not a buffer

Parameters:: target (str)
Return type:: Tensor

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Return type:: Any

Returns:: object: Any extra state to store in the module’s state_dict

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the Parameter: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.nn.Parameter: The Parameter referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not an nn.Parameter

Parameters:: target (str)
Return type:: Parameter

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

Returns:

torch.nn.Module: The submodule referenced by target

Raises:

AttributeError: If at any point along the path resulting from: the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:: target (str)
Return type:: Module

half()

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

ipu(device=None)

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

load_state_dict(state_dict, strict=True, assign=False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Args:

state_dict (dict): a dict containing parameters and: persistent buffers.
strict (bool, optional): whether to strictly enforce that the keys: in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional): When set to False, the properties of the tensors: in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

NamedTuple with missing_keys and unexpected_keys fields:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Note:

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

Parameters:

state_dict (Mapping[str, Any])
strict (bool)
assign (bool)

modules()

Return an iterator over all modules in the network.

Return type:: Iterator[Module]

Yields:: Module: a module in the network
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Args:

prefix (str): prefix to prepend to all buffer names. recurse (bool, optional): if True, then yields buffers of this module

and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor): Tuple containing the name and buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Return type:: Iterator[tuple[str, Module]]

Yields:: (str, Module): Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Args:: memo: a memo to store the set of modules already added to the result prefix: a prefix that will be added to the name of the module remove_duplicate: whether to remove the duplicated module instances in the result

or not
Yields:: (str, Module): Tuple of name and module
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

Parameters:

memo (set[Module] | None)
prefix (str)
remove_duplicate (bool)

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Args:

prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module

and all submodules. Otherwise, yields only parameters that are direct members of this module.

remove_duplicate (bool, optional): whether to remove the duplicated: parameters in the result. Defaults to True.

Yields:

(str, Parameter): Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Args:

recurse (bool): if True, then yields parameters of this module: and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter: module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Parameter]

register_backward_hook(hook)

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

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

Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
Return type:: RemovableHandle

register_buffer(name, tensor, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (str): name of the buffer. The buffer can be accessed: from this module using the given name
tensor (Tensor or None): buffer to be registered. If None, then operations: that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool): whether the buffer is part of this module’s: state_dict.

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

Parameters:

name (str)
tensor (Tensor | None)
persistent (bool)

Return type:

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Args:

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

before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If True, the hook will be passed the: kwargs given to the forward function. Default: False
always_call (bool): If True the hook will be run regardless of: whether an exception is raised while calling the Module. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...], Any], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any], Any], Any | None])
prepend (bool)
with_kwargs (bool)
always_call (bool)

Return type:

RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Args:

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

all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If true, the hook will be passed the kwargs: given to the forward function. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...]], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None])
prepend (bool)
with_kwargs (bool)

Return type:

RemovableHandle

register_full_backward_hook(hook, prepend=False)

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

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

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Arguments:

hook (Callable): Callable hook that will be invoked before: loading the state dict.

register_module(name, module)

Alias for add_module().

Parameters:

name (str)
module (Module | None)

Return type:

register_parameter(name, param)

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (str): name of the parameter. The parameter can be accessed: from this module using the given name
param (Parameter or None): parameter to be added to the module. If: None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Parameters:

name (str)
param (Parameter | None)

Return type:

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Args:

requires_grad (bool): whether autograd should record operations on: parameters in this module. Default: True.

Returns:

Module: self

Parameters:: requires_grad (bool)
Return type:: Self

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Args:: state (dict): Extra state from the state_dict

Parameters:: state (Any)
Return type:: None

set_submodule(target, module, strict=False)

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

module: The module to set the submodule to. strict: If False, the method will replace an existing submodule

or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError: If the target string is empty or if module is not an instance of nn.Module. AttributeError: If at any point along the path resulting from

the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:

target (str)
module (Module)
strict (bool)

Return type:

share_memory()

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)

Overloads:

self, destination (T_destination), prefix (str), keep_vars (bool) → T_destination
self, prefix (str), keep_vars (bool) → dict[str, Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Args:

destination (dict, optional): If provided, the state of module will: be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional): a prefix added to parameter and buffer: names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional): by default the Tensor s: returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

dict:: a dictionary containing a whole state of the module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)

Overloads:

self, device (DeviceLikeType | None), dtype (dtype | None), non_blocking (bool) → Self
self, dtype (dtype), non_blocking (bool) → Self
self, tensor (Tensor), non_blocking (bool) → Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters: and buffers in this module
dtype (torch.dtype): the desired floating point or complex dtype of: the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired: dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format): the desired memory: format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)

Move the parameters and buffers to the specified device without copying storage.

Args:

device (torch.device): The desired device of the parameters: and buffers in this module.
recurse (bool): Whether parameters and buffers of submodules should: be recursively moved to the specified device.

Returns:

Module: self

Parameters:

device (str | device | int | None)
recurse (bool)

Return type:

Self

train(mode=True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation: mode (False). Default: True.

Returns:

Module: self

Parameters:: mode (bool)
Return type:: Self

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Parameters:: dst_type (dtype | str)
Return type:: Self

xpu(device=None)

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

zero_grad(set_to_none=True)

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: See torch.optim.Optimizer.zero_grad() for details.

Parameters:: set_to_none (bool)
Return type:: None

training: bool

Moore2016ExcitationPattern

class torch_amt.Moore2016ExcitationPattern(erb_lower=1.75, erb_upper=39.0, erb_step=0.25, spreading_limit_octaves=4.0, learnable=False, dtype=torch.float32)[source]

Bases: Module

Excitation pattern computation using roex (rounded-exponential) auditory filters.

Converts a sparse spectrum (discrete frequency components with levels) into an excitation pattern on the ERB-rate scale using rounded-exponential (roex) filters with level-dependent lower slopes, following Moore et al. (2016).

The roex filter models the auditory filter shape with asymmetric slopes that depend on the input sound level, capturing the nonlinear frequency selectivity of the human cochlea. Each spectral component spreads to nearby ERB channels according to the roex weighting function.

Parameters:

erb_lower (float) – Lower ERB-rate limit for excitation pattern channels. Default: 1.75 (approximately 47 Hz).
erb_upper (float) – Upper ERB-rate limit for excitation pattern channels. Default: 39.0 (approximately 15 kHz).
erb_step (float) – ERB-rate step size between channels. Default: 0.25 (creates 150 channels from 1.75 to 39.0).
spreading_limit_octaves (float) – Maximum spreading distance in octaves (±). Default: 4.0. Components more than this distance from a channel center frequency are ignored to reduce computational cost.
learnable (bool) – If True, makes level_dep_factor (0.35) a learnable parameter. Default: False.
dtype (dtype) – Data type for computations. Default: torch.float32.

n_channels

Number of ERB channels in excitation pattern (typically 150).

Type:: int

erb_channels

ERB-rate values for each channel, shape (n_channels,).

Type:: torch.Tensor

fc_channels

Center frequencies in Hz for each channel, shape (n_channels,).

Type:: torch.Tensor

p_channels

Base filter slopes p(f) for each channel, shape (n_channels,).

Type:: torch.Tensor

p1000

Reference slope value at 1000 Hz used for level-dependent scaling.

Type:: float

level_dep_factor

Level-dependent slope factor (default 0.35). Registered as buffer or parameter depending on learnable. When learnable, clamped to [0.0, 1.0] for numerical stability.

Type:: torch.Tensor or nn.Parameter

reference_level

Reference level for slope computation (default 51.0 dB SPL). When learnable, clamped to [30.0, 70.0] dB SPL.

Type:: torch.Tensor or nn.Parameter

min_p_lower

Minimum value for p_lower to ensure positivity (default 0.1). When learnable, clamped to [0.01, 1.0].

Type:: torch.Tensor or nn.Parameter

Examples

Basic usage with sparse spectrum:

>>> import torch
>>> from torch_amt.common.filterbanks import Moore2016ExcitationPattern
>>>
>>> # Create excitation pattern module
>>> exc_pattern = Moore2016ExcitationPattern()
>>> print(f"Channels: {exc_pattern.n_channels}")
Channels: 150
>>>
>>> # Sparse spectrum: 3 frequency components
>>> freqs = torch.tensor([[500.0, 1000.0, 2000.0]])  # (batch=1, components=3)
>>> levels = torch.tensor([[60.0, 65.0, 55.0]])  # dB SPL
>>>
>>> # Compute excitation pattern
>>> excitation = exc_pattern(freqs, levels)
>>> print(excitation.shape)
torch.Size([1, 150])
>>> print(f"Peak excitation: {excitation.max():.1f} dB")
Peak excitation: 64.8 dB

Batch processing:

>>> freqs_batch = torch.tensor([
...     [500.0, 1000.0, 2000.0],
...     [750.0, 1500.0, 3000.0]
... ])  # (batch=2, components=3)
>>> levels_batch = torch.tensor([
...     [60.0, 65.0, 55.0],
...     [58.0, 62.0, 57.0]
... ])
>>>
>>> exc_batch = exc_pattern(freqs_batch, levels_batch)
>>> print(exc_batch.shape)
torch.Size([2, 150])

Custom ERB range:

>>> # Narrower frequency range (100 Hz to 8 kHz)
>>> exc_narrow = Moore2016ExcitationPattern(
...     erb_lower=6.0, erb_upper=30.0, erb_step=0.5
... )
>>> print(f"Channels: {exc_narrow.n_channels}")
Channels: 49

Notes

Roex Filter Theory:

The roex (rounded-exponential) filter is defined as:

\[W(p, g) = (1 + p|g|) \cdot e^{-p|g|}\]

where:

\(g = (f - f_c) / f_c\) is the normalized frequency deviation
\(p(f)\) is the filter slope parameter
\(f_c\) is the channel center frequency

The slope parameter is frequency-dependent:

\[p(f) = \frac{4f}{\text{ERB}(f)}\]

where \(\text{ERB}(f) = 24.673(4.368f/1000 + 1)\) is the Equivalent Rectangular Bandwidth in Hz.

Level-Dependent Spreading:

For frequencies below the channel center (\(g < 0\)), the slope is level-dependent:

\[p_l(f, X) = p(f) - 0.35 \cdot \frac{p(f)}{p(1000)} \cdot (X - 51)\]

where \(X\) is the input level in dB SPL. This models the asymmetric broadening of auditory filters at high levels:

At low levels (X < 51 dB): \(p_l > p\) (sharper lower slope)
At high levels (X > 51 dB): \(p_l < p\) (broader lower slope)

The reference level of 51 dB SPL and slope ratio \(p(f)/p(1000)\) ensure consistent scaling across frequencies.

ERB Scale Computation:

Channels are spaced uniformly on the ERB-rate scale from erb_lower to erb_upper with step erb_step. The ERB-rate is converted to frequency using erbrate2f():

\[f = \frac{1}{0.00437}(e^{\text{ERB-rate}/9.2645} - 1)\]

This spacing matches the critical band scale of human hearing, with finer resolution at low frequencies.

Spreading and Computational Efficiency:

To reduce computation, spreading is limited to ±``spreading_limit_octaves`` from each channel. For the default 4.0 octaves, only channels within \(f_c / 16\) to \(16 f_c\) receive contributions from a component at frequency \(f_c\).

Output Units:

The excitation pattern is returned in dB, computed as:

\[E_{\text{dB}}[n] = 10 \log_{10}\left(\sum_i W_i[n] \cdot 10^{L_i/10}\right)\]

where \(L_i\) is the level of component \(i\) and \(W_i[n]\) is the roex weighting from component \(i\) to channel \(n\). Linear power summation is performed before converting to dB.

See also

erbrate2f: Convert ERB-rate to frequency in Hz
f2erb: Convert frequency to ERB bandwidth

References

__init__(erb_lower=1.75, erb_upper=39.0, erb_step=0.25, spreading_limit_octaves=4.0, learnable=False, dtype=torch.float32)[source]

Initialize Moore2016 excitation pattern module.

Sets up ERB-rate channels, computes center frequencies and base filter slopes, and registers buffers/parameters for excitation pattern computation.

Parameters:

erb_lower (float) –
Lower ERB-rate limit. Default: 1.75.

Corresponds to approximately 47 Hz. Values below 1.0 approach DC and may cause numerical issues. Typical range: [1.0, 5.0].
erb_upper (float) –
Upper ERB-rate limit. Default: 39.0.

Corresponds to approximately 15 kHz. Values above 40.0 exceed typical hearing range (20 kHz). Typical range: [35.0, 40.0].
erb_step (float) –
ERB-rate step size between channels. Default: 0.25.

Determines frequency resolution. Smaller steps = finer resolution but higher computational cost:
- 0.25: 150 channels (1.75 to 39.0) - standard
- 0.5: 75 channels - faster, coarser
- 0.1: 373 channels - slower, finer
spreading_limit_octaves (float) –
Maximum spreading distance in octaves (±). Default: 4.0.

Limits excitation spreading to reduce computation. A component at frequency \(f\) only affects channels in the range \([f/2^4, f \cdot 2^4] = [f/16, 16f]\). Smaller values reduce accuracy but increase speed. Typical range: [2.0, 6.0].
learnable (bool) – If True, makes level_dep_factor (0.35) a learnable nn.Parameter for gradient-based optimization. Default: False (fixed buffer).
dtype (dtype) – Data type for computations and parameters. Default: torch.float32. Use torch.float64 for higher precision if needed.

Raises:

ValueError – If erb_lower >= erb_upper (invalid range).
ValueError – If erb_step <= 0 (must be positive).
ValueError – If spreading_limit_octaves <= 0 (must be positive).

Notes

ERB Channel Generation:

Channels are generated uniformly on the ERB-rate scale:

\[\text{ERB-rate}_n = \text{erb\_lower} + n \cdot \text{erb\_step}\]

for \(n = 0, 1, \ldots, N_{\text{channels}}-1\) where \(N_{\text{channels}} = \lceil(\text{erb\_upper} - \text{erb\_lower}) / \text{erb\_step}\rceil + 1\).

Each ERB-rate value is converted to frequency using erbrate2f():

\[f_c = \frac{1}{0.00437}\left(e^{\text{ERB-rate}/9.2645} - 1\right)\]

Filter Slope Computation:

For each channel, the base filter slope is computed as:

\[p(f_c) = \frac{4f_c}{\text{ERB}(f_c)}\]

where the ERB bandwidth is:

\[\text{ERB}(f) = 24.673\left(\frac{4.368f}{1000} + 1\right)\]

This formulation ensures that the roex filter has an equivalent rectangular bandwidth matching the critical band of human hearing.

Reference Slope p(1000 Hz):

The reference slope at 1000 Hz is used for level-dependent scaling:

\[p_{1000} = \frac{4 \cdot 1000}{24.673(4.368 + 1)} \approx 30.20\]

This value normalizes the level-dependent slope adjustment across frequencies.

Computational Cost:

Initialization is O(n_channels), typically ~150 channels requiring ~500 operations (ERB conversion + slope computation). This is negligible compared to forward pass cost.

forward(freqs, levels)[source]

Compute excitation pattern from sparse spectrum (VECTORIZED).

Takes frequency components with their levels and spreads them to ERB-spaced channels using roex filters with level-dependent slopes.

Optimization Note: This implementation is fully vectorized, processing all (batch x components) in parallel instead of nested loops. This provides ~12x speedup compared to the original implementation while maintaining numerical accuracy within 0.001 dB (imperceptible difference).

Parameters:

freqs (Tensor) –
Frequencies of spectral components, shape (batch, n_components), in Hz.

Each value should be positive and within the auditory range (~20 to 20000 Hz). Components with freq < 1e-6 Hz are skipped.
levels (Tensor) –
Levels of spectral components, shape (batch, n_components), in dB SPL.

Typical range: 0 to 100 dB SPL. Components with level < -50 dB are treated as negligible and skipped.

Returns:

excitation – Excitation pattern, shape (batch, n_channels), in dB.

Values typically range from -60 dB (threshold) to ~80 dB (high levels). Output is on the same device and dtype as input freqs.

Return type:

Notes

Spreading Algorithm:

For each spectral component \(i\) at frequency \(f_i\) with level \(L_i\), the contribution to channel \(n\) with center frequency \(f_c[n]\) is computed as:

Calculate normalized deviation: \(g = (f_c[n] - f_i) / f_i\)
Check octave distance: Skip if \(|\log_2(f_c[n]/f_i)| > 4\)
Compute base slope: \(p(f_i) = 4f_i / \text{ERB}(f_i)\)
Apply level-dependent slope for \(g < 0\) (channels below component):

\[p_l = p(f_i) - 0.35 \cdot \frac{p(f_i)}{p(1000)} \cdot (L_i - 51)\]
Calculate roex weight: \(W = (1 + p_{\text{eff}}|g|) e^{-p_{\text{eff}}|g|}\)
Add contribution: \(C[n] = L_i + 10\log_{10}(W)\)

where \(p_{\text{eff}} = p_l\) for \(g < 0\), \(p(f_i)\) for \(g \geq 0\).

Linear Power Summation:

Contributions from all components are summed in linear power:

\[E_{\text{linear}}[n] = \sum_i W_i[n] \cdot 10^{L_i/10}\]

This represents the total excitation power at each channel, accounting for all spectral components within the spreading limit.

dB Conversion:

The final excitation pattern is converted to dB:

\[E_{\text{dB}}[n] = 10\log_{10}(E_{\text{linear}}[n] + \epsilon)\]

where \(\epsilon = 10^{-12}\) prevents log(0) errors. Channels with no contributions have \(E_{\text{linear}} \approx 0\), resulting in \(E_{\text{dB}} \approx -120\) dB (effectively silent).

Vectorization Strategy:

The original implementation used nested loops:

for b in range(batch_size):
    for i in range(n_components):
        contributions = _calculate_input_levels(...)
        excitation[b] += contributions

This new implementation:

Flattens (batch, components) → (batch*components)
Broadcasts all operations over (batch*components, n_channels)
Uses scatter_add to aggregate contributions by batch

Result: ~12x speedup with <0.001 dB difference (imperceptible).

See also

_get_W_vectorized: Vectorized roex filter weighting
_get_p: Base slope computation

T_destination = ~T_destination

add_module(name, module)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (str): name of the child module. The child module can be: accessed from this module using the given name

module (Module): child module to be added to the module.

Parameters:

name (str)
module (Module | None)

Return type:

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Args:: fn (Module -> None): function to be applied to each submodule
Returns:: Module: self

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Parameters:: fn (Callable[[Module], None])
Return type:: Self

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

buffers(recurse=True)

Return an iterator over module buffers.

Args:

recurse (bool): if True, then yields buffers of this module: and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor: module buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Tensor]

call_super_init: bool = False

children()

Return an iterator over immediate children modules.

Return type:: Iterator[Module]

Yields:: Module: a child module

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

cuda(device=None)

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Args:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

double()

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

dump_patches: bool = False

eval()

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Return type:: Self

Returns:: Module: self

float()

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the buffer: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.Tensor: The buffer referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not a buffer

Parameters:: target (str)
Return type:: Tensor

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Return type:: Any

Returns:: object: Any extra state to store in the module’s state_dict

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the Parameter: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.nn.Parameter: The Parameter referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not an nn.Parameter

Parameters:: target (str)
Return type:: Parameter

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

Returns:

torch.nn.Module: The submodule referenced by target

Raises:

AttributeError: If at any point along the path resulting from: the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:: target (str)
Return type:: Module

half()

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

ipu(device=None)

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

load_state_dict(state_dict, strict=True, assign=False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Args:

state_dict (dict): a dict containing parameters and: persistent buffers.
strict (bool, optional): whether to strictly enforce that the keys: in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional): When set to False, the properties of the tensors: in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

NamedTuple with missing_keys and unexpected_keys fields:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Note:

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

Parameters:

state_dict (Mapping[str, Any])
strict (bool)
assign (bool)

modules()

Return an iterator over all modules in the network.

Return type:: Iterator[Module]

Yields:: Module: a module in the network
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Args:

prefix (str): prefix to prepend to all buffer names. recurse (bool, optional): if True, then yields buffers of this module

and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor): Tuple containing the name and buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Return type:: Iterator[tuple[str, Module]]

Yields:: (str, Module): Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Args:: memo: a memo to store the set of modules already added to the result prefix: a prefix that will be added to the name of the module remove_duplicate: whether to remove the duplicated module instances in the result

or not
Yields:: (str, Module): Tuple of name and module
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

Parameters:

memo (set[Module] | None)
prefix (str)
remove_duplicate (bool)

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Args:

prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module

and all submodules. Otherwise, yields only parameters that are direct members of this module.

remove_duplicate (bool, optional): whether to remove the duplicated: parameters in the result. Defaults to True.

Yields:

(str, Parameter): Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Args:

recurse (bool): if True, then yields parameters of this module: and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter: module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Parameter]

register_backward_hook(hook)

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

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

Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
Return type:: RemovableHandle

register_buffer(name, tensor, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (str): name of the buffer. The buffer can be accessed: from this module using the given name
tensor (Tensor or None): buffer to be registered. If None, then operations: that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool): whether the buffer is part of this module’s: state_dict.

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

Parameters:

name (str)
tensor (Tensor | None)
persistent (bool)

Return type:

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Args:

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

before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If True, the hook will be passed the: kwargs given to the forward function. Default: False
always_call (bool): If True the hook will be run regardless of: whether an exception is raised while calling the Module. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...], Any], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any], Any], Any | None])
prepend (bool)
with_kwargs (bool)
always_call (bool)

Return type:

RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Args:

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

all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If true, the hook will be passed the kwargs: given to the forward function. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...]], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None])
prepend (bool)
with_kwargs (bool)

Return type:

RemovableHandle

register_full_backward_hook(hook, prepend=False)

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

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

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Arguments:

hook (Callable): Callable hook that will be invoked before: loading the state dict.

register_module(name, module)

Alias for add_module().

Parameters:

name (str)
module (Module | None)

Return type:

register_parameter(name, param)

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (str): name of the parameter. The parameter can be accessed: from this module using the given name
param (Parameter or None): parameter to be added to the module. If: None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Parameters:

name (str)
param (Parameter | None)

Return type:

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Args:

requires_grad (bool): whether autograd should record operations on: parameters in this module. Default: True.

Returns:

Module: self

Parameters:: requires_grad (bool)
Return type:: Self

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Args:: state (dict): Extra state from the state_dict

Parameters:: state (Any)
Return type:: None

set_submodule(target, module, strict=False)

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

module: The module to set the submodule to. strict: If False, the method will replace an existing submodule

or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError: If the target string is empty or if module is not an instance of nn.Module. AttributeError: If at any point along the path resulting from

the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:

target (str)
module (Module)
strict (bool)

Return type:

share_memory()

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)

Overloads:

self, destination (T_destination), prefix (str), keep_vars (bool) → T_destination
self, prefix (str), keep_vars (bool) → dict[str, Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Args:

destination (dict, optional): If provided, the state of module will: be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional): a prefix added to parameter and buffer: names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional): by default the Tensor s: returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

dict:: a dictionary containing a whole state of the module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)

Overloads:

self, device (DeviceLikeType | None), dtype (dtype | None), non_blocking (bool) → Self
self, dtype (dtype), non_blocking (bool) → Self
self, tensor (Tensor), non_blocking (bool) → Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters: and buffers in this module
dtype (torch.dtype): the desired floating point or complex dtype of: the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired: dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format): the desired memory: format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)

Move the parameters and buffers to the specified device without copying storage.

Args:

device (torch.device): The desired device of the parameters: and buffers in this module.
recurse (bool): Whether parameters and buffers of submodules should: be recursively moved to the specified device.

Returns:

Module: self

Parameters:

device (str | device | int | None)
recurse (bool)

Return type:

Self

train(mode=True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation: mode (False). Default: True.

Returns:

Module: self

Parameters:: mode (bool)
Return type:: Self

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Parameters:: dst_type (dtype | str)
Return type:: Self

xpu(device=None)

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

zero_grad(set_to_none=True)

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: See torch.optim.Optimizer.zero_grad() for details.

Parameters:: set_to_none (bool)
Return type:: None

training: bool

extra_repr()[source]

Extra representation string for module printing.

Returns:: String summarizing module parameters: ERB range, number of channels, spreading limit, and learnable status.
Return type:: str

ERBIntegration

class torch_amt.ERBIntegration(fs=32000, f_min=50.0, f_max=15000.0, erb_step=0.25, learnable=False)[source]

Bases: Module

ERB-scale integration for Glasberg & Moore (2002) loudness model.

Integrates power spectral density (PSD) within Equivalent Rectangular Bandwidth (ERB) spaced frequency bands to compute the excitation pattern. This transforms the FFT-based frequency representation into an auditory-motivated frequency scale that better represents human frequency discrimination.

The module computes 150 ERB-spaced channels from 50 Hz to 15000 Hz with 0.25 ERB-rate spacing, following Glasberg & Moore (2002). Each ERB band integrates PSD energy using rectangular filters centered at the ERB channel frequencies.

Parameters:

fs (int) – Sampling rate in Hz. Default: 32000. Determines Nyquist frequency for PSD integration.
f_min (float) – Minimum center frequency in Hz. Default: 50.0. Lower bound of ERB scale matching typical auditory models.
f_max (float) – Maximum center frequency in Hz. Default: 15000.0. Upper bound avoids aliasing artifacts and matches loudness model requirements.
erb_step (float) – ERB-rate spacing step. Default: 0.25. Determines frequency resolution: smaller values give finer resolution but more channels. 0.25 yields approximately 150 channels matching Glasberg & Moore (2002).
learnable (bool) – If True, integration weights become learnable nn.Parameter objects (one weight per ERB band). Default: False (unit weights).

fs

Sampling rate in Hz.

Type:: int

f_min

Minimum center frequency in Hz.

Type:: float

f_max

Maximum center frequency in Hz.

Type:: float

erb_step

ERB-rate spacing step.

Type:: float

learnable

Whether integration weights are learnable.

Type:: bool

fc_erb

ERB channel center frequencies in Hz, shape (n_erb_bands,). Registered as buffer. Computed from ERB-rate scale.

Type:: torch.Tensor

n_erb_bands

Number of ERB channels. Typically 150 for default parameters.

Type:: int

bandwidth_scale

Global bandwidth scale factor (scalar). Multiplies all ERB bandwidths. When learnable=False, fixed to 1.0. When learnable=True, optimizable. Clamped to [0.1, 10.0] during forward pass for numerical stability. Allows optimization of effective filter bandwidth for better modeling.

Type:: torch.Tensor or nn.Parameter

integration_weights

Per-channel integration weights, shape (n_erb_bands,). When learnable=False, fixed to ones. When learnable=True, optimizable.

Type:: torch.Tensor or nn.Parameter

Examples

Basic usage with PSD from MultiResolutionFFT:

>>> import torch
>>> from torch_amt.common.filterbanks import ERBIntegration, MultiResolutionFFT
>>>
>>> # Generate PSD using multi-resolution FFT
>>> mrf = MultiResolutionFFT(fs=32000)
>>> audio = torch.randn(2, 32000)  # 2 batches, 1 second
>>> psd, freqs = mrf(audio)  # (2, 32, 1025), (1025,)
>>>
>>> # Integrate into ERB bands
>>> erb_int = ERBIntegration(fs=32000)
>>> excitation = erb_int(psd, freqs)  # (2, 32, 150)
>>> print(f"Excitation shape: {excitation.shape}")
Excitation shape: torch.Size([2, 32, 150])
>>> print(f"Excitation range: {excitation.min():.1f} - {excitation.max():.1f} dB SPL")
Excitation range: 87.3 - 132.5 dB SPL

Inspect ERB channel properties:

>>> fc = erb_int.get_erb_frequencies()
>>> bw = erb_int.get_erb_bandwidths()
>>> print(f"ERB channels: {len(fc)}")
ERB channels: 150
>>> print(f"Frequency range: {fc[0]:.1f} - {fc[-1]:.1f} Hz")
Frequency range: 50.0 - 15221.2 Hz
>>> print(f"Bandwidth range: {bw[0]:.1f} - {bw[-1]:.1f} Hz")
Bandwidth range: 30.1 - 1665.1 Hz
>>>
>>> # ERB bandwidth increases with frequency
>>> for i in [0, 50, 100, 149]:
...     print(f"ERB {i}: fc={fc[i]:.1f} Hz, bw={bw[i]:.1f} Hz")
ERB 0: fc=50.0 Hz, bw=30.1 Hz
ERB 50: fc=267.4 Hz, bw=43.4 Hz
ERB 100: fc=1429.9 Hz, bw=146.5 Hz
ERB 149: fc=15221.2 Hz, bw=1665.1 Hz

Learnable integration weights for optimization:

>>> erb_learnable = ERBIntegration(fs=32000, learnable=True)
>>> print(f"Learnable parameters: {sum(p.numel() for p in erb_learnable.parameters())}")
Learnable parameters: 150
>>>
>>> # Weights initialized to ones
>>> print(f"Initial weights: {erb_learnable.integration_weights[:5]}")
Initial weights: tensor([1., 1., 1., 1., 1.], requires_grad=True)
>>>
>>> # Can be optimized
>>> optimizer = torch.optim.Adam(erb_learnable.parameters(), lr=0.01)

Custom ERB configuration:

>>> # Coarser resolution: 0.5 ERB-rate spacing → ~75 channels
>>> erb_coarse = ERBIntegration(fs=32000, erb_step=0.5)
>>> print(f"Coarse ERB bands: {erb_coarse.n_erb_bands}")
Coarse ERB bands: 75
>>>
>>> # Extended low frequency: 20 Hz minimum
>>> erb_extended = ERBIntegration(fs=32000, f_min=20.0, f_max=16000.0)
>>> print(f"Extended ERB bands: {erb_extended.n_erb_bands}")
Extended ERB bands: 155

Notes

ERB Formula (Glasberg & Moore 1990):

The Equivalent Rectangular Bandwidth in Hz is given by:

\[\text{ERB}(f) = 24.673 \left( 4.368 \frac{f}{1000} + 1 \right)\]

where \(f\) is frequency in Hz. This approximates the bandwidth of auditory filters at different center frequencies.

ERB-rate Scale:

The ERB-rate (perceptual frequency scale) is computed as:

\[\text{ERB-rate}(f) = 21.4 \log_{10}(0.00437f + 1)\]

For uniform spacing on the ERB-rate scale (default 0.25), center frequencies are: \(f_c = \text{erbrate2f}(1.75 + 0.25k)\) for \(k = 0, 1, ..., 149\).

Integration Algorithm:

For each ERB band \(i\) with center frequency \(f_{c,i}\) and bandwidth \(\text{ERB}_i\):

Define rectangular filter: \([f_{c,i} - \text{ERB}_i/2, f_{c,i} + \text{ERB}_i/2]\)
Find PSD bins within filter: \(\text{mask} = (f >= f_{\text{low}}) \& (f <= f_{\text{high}})\)
Integrate PSD: \(P_i = \sum_{k \in \text{mask}} \text{PSD}[k] \cdot \Delta f \cdot w_i\)
Convert to dB SPL: \(E_i = 10 \log_{10}(P_i / p_{\text{ref}}^2)\)

where \(\Delta f\) is PSD bin width, \(w_i\) is integration weight, and \(p_{\text{ref}} = 20 \mu\text{Pa}\) is the standard acoustic reference pressure.

dB SPL Conversion:

The output excitation is in dB SPL assuming:

Input PSD is calibrated in Pa²/Hz
Reference pressure \(p_{\text{ref}} = 20 \times 10^{-6}\) Pa (20 μPa)
Formula: \(L_{\text{SPL}} = 10 \log_{10}(P / p_{\text{ref}}^2)\)

For digital audio, calibration typically assumes 1 RMS ≈ some dB SPL (often ~94 dB SPL for full scale, matching typical sound level meters).

Computational Complexity:

For PSD with shape (batch, n_frames, n_freq_bins) and n_erb_bands:

Outer loop: n_erb_bands iterations (typically 150)
Inner operations: O(n_freq_bins) masking + summation per band
Total: O(batch x n_frames x n_erb_bands x n_freq_bins)
For (2, 32, 1025) PSD → 150 bands: ~9.8M operations (~1 ms on GPU)

See also

MultiResolutionFFT: Computes PSD input for this module.
GammatoneFilterbank: Alternative ERB-spaced filterbank in time domain.
f2erbrate: Convert frequency to ERB-rate scale.
erbrate2f: Convert ERB-rate to frequency.
audfiltbw: Auditory filter bandwidth (Glasberg & Moore 1990).

References

__init__(fs=32000, f_min=50.0, f_max=15000.0, erb_step=0.25, learnable=False)[source]

Initialize ERB integration module.

Parameters:

fs (int) – Sampling rate in Hz. Default: 32000.
f_min (float) – Minimum ERB center frequency in Hz. Default: 50.0. Must be positive and less than f_max.
f_max (float) – Maximum ERB center frequency in Hz. Default: 15000.0. Should not exceed Nyquist frequency (fs/2) to avoid aliasing.
erb_step (float) – ERB-rate spacing step. Default: 0.25. Smaller values give finer frequency resolution but more channels (slower). Common values: 0.25 (fine, ~150 channels), 0.5 (coarse, ~75).
learnable (bool) – If True, creates learnable integration weights (one per channel). Default: False (weights fixed to 1.0).

Notes

ERB Channel Computation:

The constructor computes ERB center frequencies as:

Convert f_min, f_max to ERB-rate using f2erbrate()
Create uniform grid: erb_centers = arange(erb_min, erb_max, erb_step)
Convert back to Hz using erbrate2f(): fc_erb = erbrate2f(erb_centers)

This ensures uniform spacing on the perceptual ERB-rate scale, not uniform in Hz (ERB channels become increasingly wider at high frequencies).

Number of Channels:

For default parameters (50-15000 Hz, step=0.25):

erb_min = f2erbrate(50) ≈ 1.75
erb_max = f2erbrate(15000) ≈ 39.0
n_channels = (39.0 - 1.75) / 0.25 = 149-150

forward(psd, freqs)[source]

Integrate power spectral density into ERB-spaced excitation pattern.

For each ERB band, integrates PSD energy within the band’s rectangular frequency range, then converts to dB SPL.

Parameters:

psd (Tensor) – Power spectral density in Pa²/Hz, shape (batch, n_frames, n_freq_bins). Typically output from MultiResolutionFFT.
freqs (Tensor) – Frequency vector in Hz, shape (n_freq_bins,). Must match last dimension of psd.

Returns:

Excitation pattern in dB SPL, shape (batch, n_frames, n_erb_bands). Each value represents integrated energy in one ERB band.

Return type:

Notes

Integration Algorithm:

For each ERB band \(i\) with center frequency \(f_{c,i}\) and bandwidth \(\text{ERB}_i\):

Compute filter edges: \(f_{\text{low}} = f_{c,i} - \text{ERB}_i/2\), \(f_{\text{high}} = f_{c,i} + \text{ERB}_i/2\)
Create frequency mask: \(M = (f >= f_{\text{low}}) \& (f <= f_{\text{high}})\)
Integrate PSD: \(P_i = w_i \sum_{k \in M} \text{PSD}[k] \cdot \Delta f\)
Convert to dB SPL: \(E_i = 10 \log_{10}(P_i / p_{\text{ref}}^2)\)

where \(\Delta f\) is PSD bin width (freqs[1] - freqs[0]), \(w_i\) is integration weight (learnable or 1.0), and \(p_{\text{ref}} = 20 \mu\text{Pa}\).

Rectangular Filters:

This implementation uses simple rectangular frequency masking. More sophisticated models (e.g., roex filters) provide better approximation of auditory filter shapes but are computationally more expensive.

dB SPL Calibration:

Output is in dB SPL (Sound Pressure Level) assuming:

Input PSD is calibrated in Pa²/Hz
Reference pressure: 20 μPa (standard for airborne sound)
Formula: \(L_{\text{SPL}} = 10 \log_{10}(P / p_{\text{ref}}^2)\)

For digital audio, 0 dBFS typically maps to ~94 dB SPL (full scale sine wave), though exact calibration depends on recording setup.

Examples

Basic PSD integration:

>>> import torch
>>> from torch_amt.common.filterbanks import ERBIntegration, MultiResolutionFFT
>>>
>>> # Generate PSD
>>> mrf = MultiResolutionFFT(fs=32000)
>>> audio = torch.randn(4, 32000)  # 4 batches
>>> psd, freqs = mrf(audio)
>>>
>>> # Integrate into ERB bands
>>> erb = ERBIntegration(fs=32000)
>>> excitation = erb(psd, freqs)
>>> print(f"Shape: {psd.shape} -> {excitation.shape}")
Shape: torch.Size([4, 32, 1025]) -> torch.Size([4, 32, 150])
>>> print(f"ERB range: {excitation.min():.1f} - {excitation.max():.1f} dB SPL")
ERB range: 87.3 - 132.5 dB SPL

get_erb_frequencies()[source]

Return ERB channel center frequencies in Hz.

Returns:: ERB center frequencies in Hz, shape (n_erb_bands,). Uniformly spaced on ERB-rate scale, logarithmically spaced in Hz.
Return type:: Tensor

Notes

Center frequencies are computed during initialization as:

\[f_{c,k} = \text{erbrate2f}\left(\text{erb}_{\text{min}} + k \cdot \text{erb\_step}\right)\]

where \(k = 0, 1, ..., n_{\text{erb\_bands}} - 1\).

Examples

>>> from torch_amt.common.filterbanks import ERBIntegration
>>> erb = ERBIntegration(fs=32000)
>>> fc = erb.get_erb_frequencies()
>>> print(f"First 5 ERB centers: {fc[:5].tolist()}")
First 5 ERB centers: [50.0, 52.9, 55.9, 59.1, 62.5]
>>> print(f"Last 5 ERB centers: {fc[-5:].tolist()}")
Last 5 ERB centers: [13587.8, 14096.1, 14628.0, 15184.6, 15221.2]

get_erb_bandwidths()[source]

Return ERB bandwidths in Hz for all channels.

Returns:: ERB bandwidths in Hz, shape (n_erb_bands,). Computed from channel center frequencies using Glasberg & Moore (1990) formula.
Return type:: Tensor

Notes

Bandwidths are computed as:

\[\text{ERB}(f_c) = 24.673 \left( 4.368 \frac{f_c}{1000} + 1 \right)\]

These define the rectangular filter widths used for PSD integration.

Examples

>>> from torch_amt.common.filterbanks import ERBIntegration
>>> erb = ERBIntegration(fs=32000)
>>> bw = erb.get_erb_bandwidths()
>>> fc = erb.get_erb_frequencies()
>>>
>>> # Show bandwidth growth with frequency
>>> for i in [0, 50, 100, 149]:
...     print(f"fc={fc[i]:.1f} Hz: ERB={bw[i]:.1f} Hz (Q={fc[i]/bw[i]:.1f})")
fc=50.0 Hz: ERB=30.1 Hz (Q=1.7)
fc=267.4 Hz: ERB=43.4 Hz (Q=6.2)
fc=1429.9 Hz: ERB=146.5 Hz (Q=9.8)
fc=15221.2 Hz: ERB=1665.1 Hz (Q=9.1)

extra_repr()[source]

Extra representation string for module printing.

Returns:: String summarizing module parameters: sampling rate, frequency range, ERB spacing, number of channels, and learnable status.
Return type:: str

T_destination = ~T_destination

add_module(name, module)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (str): name of the child module. The child module can be: accessed from this module using the given name

module (Module): child module to be added to the module.

Parameters:

name (str)
module (Module | None)

Return type:

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Args:: fn (Module -> None): function to be applied to each submodule
Returns:: Module: self

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Parameters:: fn (Callable[[Module], None])
Return type:: Self

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

buffers(recurse=True)

Return an iterator over module buffers.

Args:

recurse (bool): if True, then yields buffers of this module: and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor: module buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Tensor]

call_super_init: bool = False

children()

Return an iterator over immediate children modules.

Return type:: Iterator[Module]

Yields:: Module: a child module

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

cuda(device=None)

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Args:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

double()

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

dump_patches: bool = False

eval()

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Return type:: Self

Returns:: Module: self

float()

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the buffer: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.Tensor: The buffer referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not a buffer

Parameters:: target (str)
Return type:: Tensor

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Return type:: Any

Returns:: object: Any extra state to store in the module’s state_dict

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the Parameter: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.nn.Parameter: The Parameter referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not an nn.Parameter

Parameters:: target (str)
Return type:: Parameter

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

Returns:

torch.nn.Module: The submodule referenced by target

Raises:

AttributeError: If at any point along the path resulting from: the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:: target (str)
Return type:: Module

half()

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

ipu(device=None)

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

load_state_dict(state_dict, strict=True, assign=False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Args:

state_dict (dict): a dict containing parameters and: persistent buffers.
strict (bool, optional): whether to strictly enforce that the keys: in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional): When set to False, the properties of the tensors: in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

NamedTuple with missing_keys and unexpected_keys fields:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Note:

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

Parameters:

state_dict (Mapping[str, Any])
strict (bool)
assign (bool)

modules()

Return an iterator over all modules in the network.

Return type:: Iterator[Module]

Yields:: Module: a module in the network
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Args:

prefix (str): prefix to prepend to all buffer names. recurse (bool, optional): if True, then yields buffers of this module

and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor): Tuple containing the name and buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Return type:: Iterator[tuple[str, Module]]

Yields:: (str, Module): Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Args:: memo: a memo to store the set of modules already added to the result prefix: a prefix that will be added to the name of the module remove_duplicate: whether to remove the duplicated module instances in the result

or not
Yields:: (str, Module): Tuple of name and module
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

Parameters:

memo (set[Module] | None)
prefix (str)
remove_duplicate (bool)

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Args:

prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module

and all submodules. Otherwise, yields only parameters that are direct members of this module.

remove_duplicate (bool, optional): whether to remove the duplicated: parameters in the result. Defaults to True.

Yields:

(str, Parameter): Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Args:

recurse (bool): if True, then yields parameters of this module: and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter: module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Parameter]

register_backward_hook(hook)

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

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

Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
Return type:: RemovableHandle

register_buffer(name, tensor, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (str): name of the buffer. The buffer can be accessed: from this module using the given name
tensor (Tensor or None): buffer to be registered. If None, then operations: that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool): whether the buffer is part of this module’s: state_dict.

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

Parameters:

name (str)
tensor (Tensor | None)
persistent (bool)

Return type:

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Args:

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

before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If True, the hook will be passed the: kwargs given to the forward function. Default: False
always_call (bool): If True the hook will be run regardless of: whether an exception is raised while calling the Module. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...], Any], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any], Any], Any | None])
prepend (bool)
with_kwargs (bool)
always_call (bool)

Return type:

RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Args:

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

all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If true, the hook will be passed the kwargs: given to the forward function. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...]], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None])
prepend (bool)
with_kwargs (bool)

Return type:

RemovableHandle

register_full_backward_hook(hook, prepend=False)

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

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

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Arguments:

hook (Callable): Callable hook that will be invoked before: loading the state dict.

register_module(name, module)

Alias for add_module().

Parameters:

name (str)
module (Module | None)

Return type:

register_parameter(name, param)

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (str): name of the parameter. The parameter can be accessed: from this module using the given name
param (Parameter or None): parameter to be added to the module. If: None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Parameters:

name (str)
param (Parameter | None)

Return type:

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Args:

requires_grad (bool): whether autograd should record operations on: parameters in this module. Default: True.

Returns:

Module: self

Parameters:: requires_grad (bool)
Return type:: Self

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Args:: state (dict): Extra state from the state_dict

Parameters:: state (Any)
Return type:: None

set_submodule(target, module, strict=False)

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

module: The module to set the submodule to. strict: If False, the method will replace an existing submodule

or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError: If the target string is empty or if module is not an instance of nn.Module. AttributeError: If at any point along the path resulting from

the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:

target (str)
module (Module)
strict (bool)

Return type:

share_memory()

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)

Overloads:

self, destination (T_destination), prefix (str), keep_vars (bool) → T_destination
self, prefix (str), keep_vars (bool) → dict[str, Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Args:

destination (dict, optional): If provided, the state of module will: be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional): a prefix added to parameter and buffer: names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional): by default the Tensor s: returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

dict:: a dictionary containing a whole state of the module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)

Overloads:

self, device (DeviceLikeType | None), dtype (dtype | None), non_blocking (bool) → Self
self, dtype (dtype), non_blocking (bool) → Self
self, tensor (Tensor), non_blocking (bool) → Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters: and buffers in this module
dtype (torch.dtype): the desired floating point or complex dtype of: the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired: dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format): the desired memory: format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)

Move the parameters and buffers to the specified device without copying storage.

Args:

device (torch.device): The desired device of the parameters: and buffers in this module.
recurse (bool): Whether parameters and buffers of submodules should: be recursively moved to the specified device.

Returns:

Module: self

Parameters:

device (str | device | int | None)
recurse (bool)

Return type:

Self

train(mode=True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation: mode (False). Default: True.

Returns:

Module: self

Parameters:: mode (bool)
Return type:: Self

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Parameters:: dst_type (dtype | str)
Return type:: Self

xpu(device=None)

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

zero_grad(set_to_none=True)

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: See torch.optim.Optimizer.zero_grad() for details.

Parameters:: set_to_none (bool)
Return type:: None

training: bool

MultiResolutionFFT

class torch_amt.MultiResolutionFFT(fs=32000, window_lengths=None, hop_fraction=0.5, learnable=False)[source]

Bases: Module

Multi-resolution FFT analysis for Glasberg & Moore (2002) loudness model.

Implements frequency-dependent time-frequency resolution by computing multiple STFTs with different window lengths and selecting the appropriate window for each frequency range. This achieves:

Long windows (better frequency resolution) for low frequencies
Short windows (better temporal resolution) for high frequencies

The window selection follows the original MATLAB implementation in the AMT (Auditory Modeling Toolbox), which uses 6 Hann windows ranging from 2 to 64 ms.

Parameters:

fs (int) – Sampling rate in Hz. Default: 32000 (required by Glasberg & Moore 2002). Note: The model is designed for 32 kHz sampling rate. Other rates may produce incorrect results due to hardcoded frequency thresholds.
window_lengths (list | None) – FFT window lengths in samples. Default: [2048, 1024, 512, 256, 128, 64] corresponding to [64, 32, 16, 8, 4, 2] ms at 32 kHz sampling rate.
hop_fraction (float) – Hop size as fraction of window length. Default: 0.5 (50% overlap). Each window uses its own hop length = window_length * hop_fraction.
learnable (bool) – If True, frequency selection thresholds become learnable nn.Parameter objects. Default: False (fixed thresholds from Glasberg & Moore 2002).

fs

Sampling rate in Hz.

Type:: int

window_lengths

FFT window lengths sorted from longest to shortest.

Type:: list of int

hop_fraction

Hop size fraction for all windows.

Type:: float

learnable

Whether frequency thresholds are learnable.

Type:: bool

freq_thresholds

Frequency boundaries in Hz, shape (5,). Separates 6 frequency ranges: [0, 500], [500, 1000], [1000, 2000], [2000, 4000], [4000, 8000], [8000, Nyquist].

Type:: torch.Tensor or nn.Parameter

windows

Pre-computed Hann windows for each window length.

Type:: dict

hop_lengths

Hop lengths in samples for each window length.

Type:: dict

window_power

Sum of squared window values for PSD normalization.

Type:: dict

Examples

Basic usage with default parameters:

>>> import torch
>>> from torch_amt.common.filterbanks import MultiResolutionFFT
>>>
>>> # Create multi-resolution FFT analyzer
>>> mrf = MultiResolutionFFT(fs=32000)
>>>
>>> # Analyze 1 second of audio (2 batches)
>>> audio = torch.randn(2, 32000)
>>> psd, freqs = mrf(audio)
>>> print(psd.shape)  # (batch=2, n_frames=32, n_freq_bins=1025)
torch.Size([2, 32, 1025])
>>> print(freqs.shape)  # (1025,)
torch.Size([1025])
>>> print(f"Frequency range: {freqs[0]:.1f} - {freqs[-1]:.1f} Hz")
Frequency range: 0.0 - 16000.0 Hz

Inspect window selection for specific frequencies:

>>> freqs_map, window_indices = mrf.get_window_selection_map()
>>> # Check which window is used for 1 kHz
>>> idx_1khz = (torch.abs(freqs_map - 1000)).argmin()
>>> print(f"1 kHz uses window {window_indices[idx_1khz]}")
1 kHz uses window 2
>>> print(f"Window length: {mrf.window_lengths[2]} samples = {mrf.window_lengths[2]/32:.1f} ms")
Window length: 512 samples = 16.0 ms

Learnable frequency thresholds for optimization:

>>> mrf_learnable = MultiResolutionFFT(fs=32000, learnable=True)
>>> print(f"Initial thresholds: {mrf_learnable.freq_thresholds}")
Initial thresholds: tensor([ 500., 1000., 2000., 4000., 8000.])
>>> # Now freq_thresholds can be optimized via backpropagation
>>> optimizer = torch.optim.Adam(mrf_learnable.parameters(), lr=0.01)

Custom window configuration:

>>> # Use only 3 windows: long, medium, short
>>> mrf_custom = MultiResolutionFFT(
...     fs=32000,
...     window_lengths=[1024, 512, 256],  # 32, 16, 8 ms
...     hop_fraction=0.75  # 75% overlap for smoother transitions
... )

Notes

Frequency-Dependent Window Selection:

The default configuration from Glasberg & Moore (2002) uses 6 frequency ranges:

0-500 Hz: 2048 samples (64 ms) → Best for low-frequency resolution
500-1000 Hz: 1024 samples (32 ms)
1000-2000 Hz: 512 samples (16 ms)
2000-4000 Hz: 256 samples (8 ms)
4000-8000 Hz: 128 samples (4 ms)
8000-16000 Hz: 64 samples (2 ms) → Best for transient capture

These ranges balance the time-frequency uncertainty principle to match the temporal and spectral resolution of human auditory perception.

PSD Normalization:

Power spectral density is computed as:

\[\text{PSD}(f) = \frac{|\text{STFT}(f)|^2}{\sum w^2 \cdot f_s}\]

where \(w\) is the window function and \(f_s\) is the sampling rate. This gives units of Pa²/Hz when the input is in Pascals.

Learnable Parameters:

When learnable=True, the 5 frequency thresholds (10 real values total for 2 boundaries per threshold) become learnable. This allows the model to adapt the window selection strategy during training.

Computational Complexity:

FFT operations: O(NWlog(W)) where N is windows, W is window length
Interpolation: O(BTF) where B is batch, T is time, F is frequency bins
Total: dominated by FFT for typical parameters

For 1 second at 32 kHz with 6 windows: ~70M operations.

MATLAB Verification:

This implementation has been verified against the MATLAB reference:

glasberg2002.m: Main loudness model (lines 107-159)
arg_glasberg2002.m: Default parameters
Window lengths: [2, 4, 8, 16, 32, 64] ms confirmed (line 26)
Frequency boundaries: [80, 500, 1250, 2540, 4050] Hz from vLimitingIndices (line 25)
FFT length: 2048 samples (line 24)
Hop size: 1 ms = 32 samples (timeStep, line 27)

Note: The Python implementation uses slightly different thresholds [500, 1000, 2000, 4000, 8000] Hz to better match the published paper, while MATLAB uses [80, 500, 1250, 2540, 4050] Hz for backward compatibility. Both produce perceptually similar results.

See also

Moore2016Spectrum: Extended multi-resolution analysis for binaural loudness
GammatoneFilterbank: Cochlear filterbank for auditory models

References

__init__(fs=32000, window_lengths=None, hop_fraction=0.5, learnable=False)[source]

Initialize multi-resolution FFT analyzer.

Parameters:

fs (int) –
Sampling rate in Hz. Default: 32000 Hz (standard for Glasberg & Moore 2002).

Note: The frequency thresholds are calibrated for 32 kHz sampling. Using other rates may produce suboptimal window-frequency matching.
window_lengths (list | None) –
FFT window lengths in samples, provided in any order (will be sorted internally longest to shortest). Each window length determines both:
- Frequency resolution: Δf = fs / window_length
- Time resolution: Δt = window_length / fs
Default: [2048, 1024, 512, 256, 128, 64] samples, which at 32 kHz correspond to [64, 32, 16, 8, 4, 2] ms windows respectively.

Example custom config: [1024, 512, 256] for 3 windows only.
hop_fraction (float) –
Hop size as fraction of window length, range (0, 1].
- Each window uses: hop_length = int(window_length * hop_fraction)
- Smaller values → more overlap → smoother time evolution
- Larger values → less computation → faster processing
Default: 0.5 (50% overlap). Common alternatives: 0.25 (75% overlap) or 0.75 (25% overlap).
learnable (bool) –
If True, the 5 frequency thresholds [500, 1000, 2000, 4000, 8000] Hz become trainable nn.Parameter objects, allowing gradient-based optimization of the window selection strategy.

Default: False (fixed thresholds from Glasberg & Moore 2002).

Raises:

ValueError – If hop_fraction is not in range (0, 1].
ValueError – If window_lengths contains non-positive integers.

Notes

Frequency-Window Mapping:

The default configuration assigns 6 windows to 6 frequency ranges via thresholds [500, 1000, 2000, 4000, 8000] Hz:

Window 0 (2048 samples, 64 ms): 0-500 Hz → Δf = 15.6 Hz
Window 1 (1024 samples, 32 ms): 500-1000 Hz → Δf = 31.25 Hz
Window 2 (512 samples, 16 ms): 1000-2000 Hz → Δf = 62.5 Hz
Window 3 (256 samples, 8 ms): 2000-4000 Hz → Δf = 125 Hz
Window 4 (128 samples, 4 ms): 4000-8000 Hz → Δf = 250 Hz
Window 5 (64 samples, 2 ms): 8000-16000 Hz → Δf = 500 Hz

This design trades off time-frequency resolution to match human auditory perception characteristics across frequency ranges.

Pre-computed Buffers:

All Hann windows are pre-computed and registered as buffers named window_{length} (e.g., window_2048, window_1024) to ensure:

Reproducibility across devices (CPU, GPU, MPS)
Efficient memory usage (computed once, reused for all STFTs)
Proper gradient flow (buffers move with .to(device))

Window power (sum of squared values) is also pre-computed for PSD normalization.

Learnable Thresholds:

When learnable=True, the model can adapt the window selection during training. This is useful for tasks where the optimal time-frequency trade-off differs from psychoacoustic models (e.g., speech enhancement, music analysis). The thresholds are initialized to [500, 1000, 2000, 4000, 8000] Hz and constrained to remain sorted via projection during optimization.

forward(signal)[source]

Apply multi-resolution FFT analysis to audio signal.

Computes STFTs with all configured window lengths, then selects the appropriate window for each frequency range to produce a single time-frequency representation with frequency-dependent resolution.

Parameters:

signal (Tensor) – Input audio signal, shape (batch, time). Can be any dtype; will be processed as-is by torch.stft.

Return type:

Tuple[Tensor, Tensor]

Returns:

psd (torch.Tensor) – Power spectral density, shape (batch, n_frames, n_freq_bins). Units: Pa²/Hz when input is in Pascals (typical for audio processing).
freqs (torch.Tensor) – Frequency vector in Hz, shape (n_freq_bins,). Linearly spaced from 0 to Nyquist frequency (fs/2).

Notes

Processing Pipeline:

Compute 6 STFTs with different window lengths
For each STFT, compute PSD in its designated frequency range
Interpolate all PSDs to common frequency grid (longest window)
Merge PSDs by frequency range to create final output

Time Resolution:

n_frames depends on the longest window and hop_fraction:

\[n_{\text{frames}} = \left\lceil \frac{T - W_{\text{max}}}{W_{\text{max}} \cdot h} \right\rceil + 1\]

where \(T\) is signal length, \(W_{\text{max}}\) is longest window, \(h\) is hop_fraction.

Computational Cost:

For default configuration (6 windows, 1 sec @ 32 kHz, batch=2):

FFTs: ~60M operations
Interpolation: ~10M operations
Total: ~70M operations (~2 ms on modern GPU)

Examples

>>> import torch
>>> from torch_amt.common.filterbanks import MultiResolutionFFT
>>>
>>> mrf = MultiResolutionFFT(fs=32000)
>>> audio = torch.randn(4, 32000)  # 4 batches, 1 second
>>> psd, freqs = mrf(audio)
>>> print(f"PSD shape: {psd.shape}, Frequency range: {freqs[0]:.0f}-{freqs[-1]:.0f} Hz")
PSD shape: torch.Size([4, 32, 1025]), Frequency range: 0-16000 Hz

get_window_selection_map()[source]

Get frequency-to-window mapping for visualization and analysis.

Returns the window index used for each frequency bin in the reference (longest window) frequency grid. Useful for understanding and visualizing how the multi-resolution analysis distributes different windows across frequencies.

Return type:

Tuple[Tensor, Tensor]

Returns:

freqs (torch.Tensor) – Frequency vector in Hz, shape (n_freq_bins,), where n_freq_bins = longest_window // 2 + 1. Linearly spaced from 0 to Nyquist.
window_indices (torch.Tensor) – Window index for each frequency, shape (n_freq_bins,), dtype torch.long. Values range from 0 to (n_windows - 1), where 0 corresponds to the longest window and (n_windows - 1) to the shortest.

Examples

>>> import torch
>>> import matplotlib.pyplot as plt
>>> from torch_amt.common.filterbanks import MultiResolutionFFT
>>>
>>> mrf = MultiResolutionFFT(fs=32000)
>>> freqs, window_idx = mrf.get_window_selection_map()
>>>
>>> # Visualize window selection
>>> plt.figure(figsize=(10, 4))
>>> plt.scatter(freqs, window_idx, s=1, alpha=0.5)
>>> plt.xlabel('Frequency (Hz)')
>>> plt.ylabel('Window Index')
>>> plt.title('Multi-Resolution Window Selection')
>>> plt.yticks(range(6), [f'{w} smp ({w/32:.0f} ms)'
...                       for w in mrf.window_lengths])
>>> plt.grid(True, alpha=0.3)
>>> plt.show()

Notes

The window index directly maps to self.window_lengths:

Index 0 → window_lengths[0] (longest, e.g., 2048 samples)
Index 1 → window_lengths[1] (e.g., 1024 samples)
…
Index 5 → window_lengths[5] (shortest, e.g., 64 samples)

This method does not require a signal input; it returns the static mapping based on current frequency thresholds.

extra_repr()[source]

Extra representation string for module printing.

Returns:: String containing key module parameters: n_windows, fs, window_range (min/max samples), hop_fraction, learnable status.
Return type:: str

T_destination = ~T_destination

add_module(name, module)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (str): name of the child module. The child module can be: accessed from this module using the given name

module (Module): child module to be added to the module.

Parameters:

name (str)
module (Module | None)

Return type:

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Args:: fn (Module -> None): function to be applied to each submodule
Returns:: Module: self

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Parameters:: fn (Callable[[Module], None])
Return type:: Self

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

buffers(recurse=True)

Return an iterator over module buffers.

Args:

recurse (bool): if True, then yields buffers of this module: and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor: module buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Tensor]

call_super_init: bool = False

children()

Return an iterator over immediate children modules.

Return type:: Iterator[Module]

Yields:: Module: a child module

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

cuda(device=None)

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Args:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

double()

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

dump_patches: bool = False

eval()

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Return type:: Self

Returns:: Module: self

float()

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the buffer: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.Tensor: The buffer referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not a buffer

Parameters:: target (str)
Return type:: Tensor

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Return type:: Any

Returns:: object: Any extra state to store in the module’s state_dict

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the Parameter: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.nn.Parameter: The Parameter referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not an nn.Parameter

Parameters:: target (str)
Return type:: Parameter

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

Returns:

torch.nn.Module: The submodule referenced by target

Raises:

AttributeError: If at any point along the path resulting from: the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:: target (str)
Return type:: Module

half()

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

ipu(device=None)

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

load_state_dict(state_dict, strict=True, assign=False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Args:

state_dict (dict): a dict containing parameters and: persistent buffers.
strict (bool, optional): whether to strictly enforce that the keys: in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional): When set to False, the properties of the tensors: in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

NamedTuple with missing_keys and unexpected_keys fields:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Note:

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

Parameters:

state_dict (Mapping[str, Any])
strict (bool)
assign (bool)

modules()

Return an iterator over all modules in the network.

Return type:: Iterator[Module]

Yields:: Module: a module in the network
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Args:

prefix (str): prefix to prepend to all buffer names. recurse (bool, optional): if True, then yields buffers of this module

and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor): Tuple containing the name and buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Return type:: Iterator[tuple[str, Module]]

Yields:: (str, Module): Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Args:: memo: a memo to store the set of modules already added to the result prefix: a prefix that will be added to the name of the module remove_duplicate: whether to remove the duplicated module instances in the result

or not
Yields:: (str, Module): Tuple of name and module
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

Parameters:

memo (set[Module] | None)
prefix (str)
remove_duplicate (bool)

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Args:

prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module

and all submodules. Otherwise, yields only parameters that are direct members of this module.

remove_duplicate (bool, optional): whether to remove the duplicated: parameters in the result. Defaults to True.

Yields:

(str, Parameter): Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Args:

recurse (bool): if True, then yields parameters of this module: and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter: module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Parameter]

register_backward_hook(hook)

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

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

Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
Return type:: RemovableHandle

register_buffer(name, tensor, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (str): name of the buffer. The buffer can be accessed: from this module using the given name
tensor (Tensor or None): buffer to be registered. If None, then operations: that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool): whether the buffer is part of this module’s: state_dict.

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

Parameters:

name (str)
tensor (Tensor | None)
persistent (bool)

Return type:

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Args:

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

before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If True, the hook will be passed the: kwargs given to the forward function. Default: False
always_call (bool): If True the hook will be run regardless of: whether an exception is raised while calling the Module. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...], Any], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any], Any], Any | None])
prepend (bool)
with_kwargs (bool)
always_call (bool)

Return type:

RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Args:

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

all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If true, the hook will be passed the kwargs: given to the forward function. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...]], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None])
prepend (bool)
with_kwargs (bool)

Return type:

RemovableHandle

register_full_backward_hook(hook, prepend=False)

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

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

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Arguments:

hook (Callable): Callable hook that will be invoked before: loading the state dict.

register_module(name, module)

Alias for add_module().

Parameters:

name (str)
module (Module | None)

Return type:

register_parameter(name, param)

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (str): name of the parameter. The parameter can be accessed: from this module using the given name
param (Parameter or None): parameter to be added to the module. If: None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Parameters:

name (str)
param (Parameter | None)

Return type:

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Args:

requires_grad (bool): whether autograd should record operations on: parameters in this module. Default: True.

Returns:

Module: self

Parameters:: requires_grad (bool)
Return type:: Self

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Args:: state (dict): Extra state from the state_dict

Parameters:: state (Any)
Return type:: None

set_submodule(target, module, strict=False)

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

module: The module to set the submodule to. strict: If False, the method will replace an existing submodule

or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError: If the target string is empty or if module is not an instance of nn.Module. AttributeError: If at any point along the path resulting from

the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:

target (str)
module (Module)
strict (bool)

Return type:

share_memory()

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)

Overloads:

self, destination (T_destination), prefix (str), keep_vars (bool) → T_destination
self, prefix (str), keep_vars (bool) → dict[str, Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Args:

destination (dict, optional): If provided, the state of module will: be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional): a prefix added to parameter and buffer: names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional): by default the Tensor s: returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

dict:: a dictionary containing a whole state of the module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)

Overloads:

self, device (DeviceLikeType | None), dtype (dtype | None), non_blocking (bool) → Self
self, dtype (dtype), non_blocking (bool) → Self
self, tensor (Tensor), non_blocking (bool) → Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters: and buffers in this module
dtype (torch.dtype): the desired floating point or complex dtype of: the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired: dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format): the desired memory: format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)

Move the parameters and buffers to the specified device without copying storage.

Args:

device (torch.device): The desired device of the parameters: and buffers in this module.
recurse (bool): Whether parameters and buffers of submodules should: be recursively moved to the specified device.

Returns:

Module: self

Parameters:

device (str | device | int | None)
recurse (bool)

Return type:

Self

train(mode=True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation: mode (False). Default: True.

Returns:

Module: self

Parameters:: mode (bool)
Return type:: Self

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Parameters:: dst_type (dtype | str)
Return type:: Self

xpu(device=None)

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

zero_grad(set_to_none=True)

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: See torch.optim.Optimizer.zero_grad() for details.

Parameters:: set_to_none (bool)
Return type:: None

training: bool

Moore2016Spectrum

class torch_amt.Moore2016Spectrum(fs=32000, segment_duration=1, db_max=93.98, learnable=False, dtype=torch.float32)[source]

Bases: Module

Multi-resolution spectral analysis for Moore et al. (2016/2018) binaural loudness model.

Implements frequency-dependent time-frequency analysis using 6 Hann windows with lengths from 2 to 64 ms, computing sparse spectral representations with only “relevant” frequency components. This differs fundamentally from MultiResolutionFFT by producing sparse output (frequency/level pairs) rather than dense power spectral densities, enabling efficient loudness computation for time-varying binaural signals.

The model applies 6 overlapping FFT analyses with window lengths matched to the temporal and frequency resolution requirements of auditory perception across the frequency spectrum. Each window contributes only to specific frequency ranges, and only components exceeding perceptual relevance thresholds are retained.

Parameters:

fs (int) – Sampling rate in Hz. Must be 32000 as required by Moore et al. (2016/2018). Other sampling rates will raise ValueError. Default: 32000.
segment_duration (int) – Duration of each analysis segment in milliseconds. Default: 1 (as specified in Moore et al. 2018). This determines the time resolution of the output (1 ms per frame).
db_max (float) – Reference SPL level in dB for full-scale sinusoid. Default: 93.98, which corresponds to dbspl(1) in the AMT MATLAB implementation. This value calibrates intensity units so that 1000.0 corresponds to 0 dB SPL.
learnable (bool) – If True, frequency band limits (freq_limits) and relevance thresholds (threshold_max_minus, threshold_absolute) become learnable nn.Parameter objects for optimization. Default: False (fixed parameters from Moore et al. 2016/2018).
dtype (dtype) – Data type for internal computations and parameters. Default: torch.float32. Use torch.float64 for higher precision if needed.

fs

Sampling rate in Hz (always 32000).

Type:: int

segment_duration

Analysis segment duration in milliseconds.

Type:: int

db_max

Reference SPL level for full-scale sinusoid (dB).

Type:: float

learnable

Whether frequency limits and thresholds are learnable parameters.

Type:: bool

dtype

Data type for computations.

Type:: torch.dtype

window_lengths

FFT window lengths in samples: [2048, 1024, 512, 256, 128, 64], corresponding to [64, 32, 16, 8, 4, 2] ms at 32 kHz.

Type:: list of int

hop_length

Fixed hop size in samples (32 samples = 1 ms at 32 kHz).

Type:: int

freq_limits

Frequency band limits for each window, shape (6, 2). Each row contains [f_low, f_high] in Hz defining which frequency range each window analyzes. Default: [[20, 80], [80, 500], [500, 1250], [1250, 2540], [2540, 4050], [4050, 15000]].

Type:: torch.Tensor or nn.Parameter

threshold_max_minus

Relative threshold in dB below maximum component. Components below max - threshold_max_minus are discarded. Default: 60.0 dB.

Type:: torch.Tensor or nn.Parameter

threshold_absolute

Absolute SPL threshold in dB. Components below this level are discarded regardless of relative level. Default: -30.0 dB SPL.

Type:: torch.Tensor or nn.Parameter

hann_correction

Intensity correction factor for Hann windowing: \(10^{3.32/10} \approx 2.148\). Accounts for power loss due to windowing.

Type:: float

window_{0-5}

Pre-computed zero-padded Hann windows registered as buffers. Each window is centered in a 2048-sample array.

Type:: torch.Tensor

Examples

Basic usage with stereo audio:

>>> import torch
>>> from torch_amt.common.filterbanks import Moore2016Spectrum
>>>
>>> # Create spectrum analyzer
>>> spectrum = Moore2016Spectrum(fs=32000)
>>>
>>> # Analyze 1 second of stereo audio (batch=2)
>>> audio = torch.randn(2, 2, 32000)  # (batch, channels=2, samples)
>>> freqs_l, levels_l, freqs_r, levels_r = spectrum(audio)
>>>
>>> print(f"Left channel: {freqs_l.shape}")  # (2, 936, ~958) - varies!
torch.Size([2, 936, 958])
>>> print(f"Frequency range: {freqs_l[freqs_l > 0].min():.1f} - {freqs_l[freqs_l > 0].max():.1f} Hz")
Frequency range: 31.2 - 14984.4 Hz
>>> print(f"Level range: {levels_l[levels_l != 0].min():.1f} - {levels_l[levels_l != 0].max():.1f} dB SPL")
Level range: 12.3 - 77.8 dB SPL

Understanding sparse output format:

>>> # Check how many relevant components per time frame
>>> n_relevant_per_frame = (freqs_l[0] > 0).sum(dim=-1)  # Non-zero = relevant
>>> print(f"Relevant components per frame: {n_relevant_per_frame[:10]}")
tensor([523, 487, 501, 519, 498, 512, 503, 496, 509, 511])
>>>
>>> # Inspect specific time frame (e.g., frame 100 for batch 0, left ear)
>>> frame_idx = 100
>>> relevant_mask = freqs_l[0, frame_idx] > 0
>>> print(f"Frame {frame_idx} has {relevant_mask.sum()} relevant components")
Frame 100 has 503 relevant components
>>> print(f"Frequencies: {freqs_l[0, frame_idx, relevant_mask][:5]} Hz")
Frequencies: tensor([  31.25,   62.50,   93.75,  125.00,  156.25]) Hz
>>> print(f"Levels: {levels_l[0, frame_idx, relevant_mask][:5]} dB SPL")
Levels: tensor([32.1, 28.5, 35.7, 29.3, 33.8]) dB SPL

Learnable frequency boundaries for optimization:

>>> spectrum_learnable = Moore2016Spectrum(fs=32000, learnable=True)
>>> print(f"Learnable parameters: {sum(p.numel() for p in spectrum_learnable.parameters())}")
Learnable parameters: 14
>>>
>>> # Frequency limits (6 windows x 2 bounds = 12 params)
>>> print(f"Initial freq_limits:\n{spectrum_learnable.freq_limits}")
tensor([[  20.,   80.],
        [  80.,  500.],
        [ 500., 1250.],
        [1250., 2540.],
        [2540., 4050.],
        [4050., 15000.]])
>>>
>>> # Thresholds (2 params)
>>> print(f"Max-minus threshold: {spectrum_learnable.threshold_max_minus} dB")
tensor(60.)
>>> print(f"Absolute threshold: {spectrum_learnable.threshold_absolute} dB SPL")
tensor(-30.)

Custom segment duration for coarser temporal resolution:

>>> # Use 2 ms segments instead of 1 ms (faster, less temporal detail)
>>> spectrum_2ms = Moore2016Spectrum(fs=32000, segment_duration=2)
>>> audio_short = torch.randn(1, 2, 16000)  # 500 ms
>>> freqs_l, levels_l, freqs_r, levels_r = spectrum_2ms(audio_short)
>>> print(f"Frames (2ms resolution): {freqs_l.shape[1]}")
Frames (2ms resolution): 218

Notes

Frequency-Dependent Window Assignment:

Each FFT window analyzes a specific frequency range matched to auditory temporal/frequency resolution tradeoffs:

Sparse Output Format:

Unlike dense spectrograms, this module returns sparse representations: only frequency components exceeding relevance criteria are included. This reduces memory and computational requirements for downstream loudness calculations. Zero frequencies indicate padding (no component).

The number of relevant components varies per time frame and depends on signal characteristics. Frames with simple spectra may have ~100 components, while complex signals may have ~900 components (out of 1024 possible bins).

Relevant Component Criteria:

A frequency component is retained if it satisfies both conditions:

Relative threshold: \(I > I_{\max} / 10^6\) (i.e., within 60 dB of maximum)
Absolute threshold: \(I > 10^{-3}\) (i.e., above -30 dB SPL when calibrated)

where \(I\) is intensity in calibrated linear units. The absolute threshold ensures inaudible components are discarded even if they dominate the spectrum (e.g., DC offset, low-frequency rumble).

Intensity Calibration and Window Correction:

Raw FFT intensities are calibrated to SPL using:

\[I_{\text{calibrated}} = I_{\text{FFT}} \cdot C_{\text{Hann}} \cdot 2^{w} \cdot 10^{d_{\max}/10}\]

where:

\(I_{\text{FFT}} = |X[k]|^2\) is raw FFT power
\(C_{\text{Hann}} = 10^{3.32/10} \approx 2.148\) corrects for Hann window power loss
\(2^{w}\) corrects for window length (w=0 for longest, w=5 for shortest)
\(10^{d_{\max}/10}\) scales so 1000.0 corresponds to 0 dB SPL

The final SPL in dB is: \(L = 10 \log_{10}(I_{\text{calibrated}})\).

Computational Complexity:

For 1 second stereo audio (32000 samples) with default settings:

Time frames: \(\lfloor (32000 - 2048) / 32 \rfloor = 936\)
FFTs per frame: 6 windows x 2 channels = 12
Total FFTs: 936 x 12 = 11,232 FFTs
FFT operations: ~30M complex multiplications
Sparse filtering: ~15M comparisons
Total: ~45M operations (~5 ms on modern GPU)

Comparison with MultiResolutionFFT:

Feature	Moore2016Spectrum	MultiResolutionFFT
Output format	Sparse (freq/level)	Dense (PSD grid)
Hop size	Fixed 1 ms	Fraction of window
Frequency bands	6 fixed ranges	5 adjustable ranges
Typical components	~500/frame	1025 bins/frame
Memory (1s stereo)	~7 MB	~33 MB
Use case	Loudness models	General TF analysis

See also

MultiResolutionFFT: Dense multi-resolution PSD for Glasberg & Moore (2002).
GammatoneFilterbank: Auditory filterbank with ERB spacing.
erbspacebw: Compute ERB-spaced frequency grid for integration.

References

__init__(fs=32000, segment_duration=1, db_max=93.98, learnable=False, dtype=torch.float32)[source]

Initialize Moore2016Spectrum analyzer.

Parameters:

fs (int) – Sampling rate in Hz. Must be 32000 (strictly enforced). Default: 32000.
segment_duration (int) – Analysis segment duration in milliseconds. Determines temporal resolution of output. Default: 1 (1 ms per output frame).
db_max (float) – Reference SPL for full-scale sinusoid in dB. Default: 93.98 (matches AMT dbspl(1) convention).
learnable (bool) – If True, frequency band limits and relevance thresholds become learnable nn.Parameter objects. Default: False (fixed from Moore et al. 2016/2018).
dtype (dtype) – Data type for computations. Default: torch.float32.

Raises:

ValueError – If fs != 32000. The model is specifically designed for 32 kHz sampling rate and frequency band definitions depend on this rate.

Notes

Sampling Rate Requirement:

The 32 kHz sampling rate is strictly enforced because:

Window lengths [2048, 1024, 512, 256, 128, 64] samples correspond to [64, 32, 16, 8, 4, 2] ms only at 32 kHz
Frequency band limits [20-80, 80-500, …] Hz are optimized for 32 kHz Nyquist frequency
Hop size of 32 samples = 1 ms temporal resolution at 32 kHz

Frequency Band Initialization:

The constructor sets up 6 frequency bands matched to window lengths:

[[20, 80], [80, 500], [500, 1250],
 [1250, 2540], [2540, 4050], [4050, 15000]]

These ranges are from Moore et al. (2016) and balance temporal/spectral resolution according to auditory perception characteristics.

Learnable Parameters:

When learnable=True, the following become nn.Parameter objects:

freq_limits: 6x2 tensor of frequency boundaries (12 parameters)
threshold_max_minus: Relative threshold below max (1 parameter)
threshold_absolute: Absolute SPL threshold (1 parameter)
Total: 14 learnable parameters

When learnable=False (default), these are registered as buffers and remain fixed.

forward(audio)[source]

Compute sparse multi-resolution spectrum for binaural audio.

Analyzes stereo audio through 6 overlapping FFT windows, applies frequency band limiting, filters relevant components, and returns sparse frequency/level pairs for each ear separately.

Parameters:

audio (Tensor) –

Binaural audio signal with shape (batch, 2, samples).

Channel 0: Left ear signal
Channel 1: Right ear signal
Samples: Must be at least 2048 (one window length)

Input validation: Raises ValueError if shape is not 3D or if audio.shape[1] != 2 (not stereo).

Return type:

Tuple[Tensor, Tensor, Tensor, Tensor]

Returns:

freqs_left (torch.Tensor) – Relevant frequency components for left ear, shape (batch, n_segments, max_components). Frequencies in Hz. Zero values indicate padding (no actual component).
levels_left (torch.Tensor) – SPL levels in dB for left ear components, shape (batch, n_segments, max_components). Corresponds 1:1 with freqs_left. Zero values indicate padding.
freqs_right (torch.Tensor) – Relevant frequency components for right ear, shape (batch, n_segments, max_components).
levels_right (torch.Tensor) – SPL levels in dB for right ear components, shape (batch, n_segments, max_components).

Raises:

ValueError – If audio is not 3D or does not have exactly 2 channels.

Notes

Processing Pipeline:

Split audio into left and right channels
For each channel independently:
1. Segment signal into 1 ms windows (hop = 32 samples)
2. For each segment, compute 6 FFTs with different window lengths
3. Combine FFTs by frequency range (20-80 Hz from longest window, etc.)
4. Apply relevance filtering (max-60dB and -30dB SPL thresholds)
5. Extract sparse frequency/level pairs
Pad sparse outputs to uniform size across all segments
Return 4 tensors (left/right x freq/level)

Time Resolution:

Number of output frames is determined by:

\[n_{\text{segments}} = \left\lfloor \frac{N_{\text{samples}} - 2048}{32} \right\rfloor\]

where 32 is the hop size (1 ms at 32 kHz). For 1 second audio (32000 samples): \(n_{\text{segments}} = (32000 - 2048) / 32 = 936\) frames.

Computational Cost:

For 1 second stereo audio (batch=1):

Segments: 936
FFTs: 936 x 6 windows x 2 channels = 11,232 FFTs
Sparse filtering: ~936 x 2 x 1024 = 1.9M comparisons
Total: ~30M operations

Examples

>>> import torch
>>> from torch_amt.common.filterbanks import Moore2016Spectrum
>>>
>>> spectrum = Moore2016Spectrum(fs=32000)
>>> audio = torch.randn(2, 2, 32000)  # 2 batches, stereo, 1 second
>>> freqs_l, levels_l, freqs_r, levels_r = spectrum(audio)
>>>
>>> print(f"Output shape: {freqs_l.shape}")  # (2, 936, ~500-900)
Output shape: torch.Size([2, 936, 823])
>>>
>>> # Check sparsity: how many relevant components per frame?
>>> n_relevant = (freqs_l[0] > 0).sum(dim=-1)
>>> print(f"Components per frame: min={n_relevant.min()}, max={n_relevant.max()}, mean={n_relevant.float().mean():.1f}")
Components per frame: min=412, max=823, mean=587.3

See also

_process_channel: Single-channel processing implementation.
_compute_segment_spectrum: Per-segment FFT combination.
_filter_relevant_components: Relevance filtering algorithm.

extra_repr()[source]

Extra representation string for module printing.

Returns:: String containing key module parameters: fs, segment_duration (ms), hop (samples), windows (count), learnable status.
Return type:: str

T_destination = ~T_destination

add_module(name, module)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Args:

name (str): name of the child module. The child module can be: accessed from this module using the given name

module (Module): child module to be added to the module.

Parameters:

name (str)
module (Module | None)

Return type:

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

Typical use includes initializing the parameters of a model (see also torch.nn.init).

Args:: fn (Module -> None): function to be applied to each submodule
Returns:: Module: self

Example:

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) is nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Parameters:: fn (Callable[[Module], None])
Return type:: Self

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

buffers(recurse=True)

Return an iterator over module buffers.

Args:

recurse (bool): if True, then yields buffers of this module: and all submodules. Otherwise, yields only buffers that are direct members of this module.

Yields:

torch.Tensor: module buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Tensor]

call_super_init: bool = False

children()

Return an iterator over immediate children modules.

Return type:: Iterator[Module]

Yields:: Module: a child module

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

Return type:: None

cpu()

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

cuda(device=None)

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Args:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

double()

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

dump_patches: bool = False

eval()

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

Return type:: Self

Returns:: Module: self

float()

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the buffer: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.Tensor: The buffer referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not a buffer

Parameters:: target (str)
Return type:: Tensor

get_extra_state()

Return any extra state to include in the module’s state_dict.

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Return type:: Any

Returns:: object: Any extra state to store in the module’s state_dict

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

Args:

target: The fully-qualified string name of the Parameter: to look for. (See get_submodule for how to specify a fully-qualified string.)

Returns:

torch.nn.Parameter: The Parameter referenced by target

Raises:

AttributeError: If the target string references an invalid: path or resolves to something that is not an nn.Parameter

Parameters:: target (str)
Return type:: Parameter

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

Returns:

torch.nn.Module: The submodule referenced by target

Raises:

AttributeError: If at any point along the path resulting from: the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:: target (str)
Return type:: Module

half()

Casts all floating point parameters and buffers to half datatype.

Note

This method modifies the module in-place.

Returns:: Module: self

Return type:: Self

ipu(device=None)

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on IPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

load_state_dict(state_dict, strict=True, assign=False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Args:

state_dict (dict): a dict containing parameters and: persistent buffers.
strict (bool, optional): whether to strictly enforce that the keys: in state_dict match the keys returned by this module’s state_dict() function. Default: True
assign (bool, optional): When set to False, the properties of the tensors: in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

NamedTuple with missing_keys and unexpected_keys fields:

missing_keys is a list of str containing any keys that are expected
by this module but missing from the provided state_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the provided state_dict.

Note:

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

Parameters:

state_dict (Mapping[str, Any])
strict (bool)
assign (bool)

modules()

Return an iterator over all modules in the network.

Return type:: Iterator[Module]

Yields:: Module: a module in the network
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. So it should be called before constructing the optimizer if the module will live on MTIA while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

named_buffers(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Args:

prefix (str): prefix to prepend to all buffer names. recurse (bool, optional): if True, then yields buffers of this module

and all submodules. Otherwise, yields only buffers that are direct members of this module. Defaults to True.

remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.

Yields:

(str, torch.Tensor): Tuple containing the name and buffer

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Return type:: Iterator[tuple[str, Module]]

Yields:: (str, Module): Tuple containing a name and child module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo=None, prefix='', remove_duplicate=True)

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Args:: memo: a memo to store the set of modules already added to the result prefix: a prefix that will be added to the name of the module remove_duplicate: whether to remove the duplicated module instances in the result

or not
Yields:: (str, Module): Tuple of name and module
Note:: Duplicate modules are returned only once. In the following example, l will be returned only once.

Example:

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

Parameters:

memo (set[Module] | None)
prefix (str)
remove_duplicate (bool)

named_parameters(prefix='', recurse=True, remove_duplicate=True)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Args:

prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module

and all submodules. Otherwise, yields only parameters that are direct members of this module.

remove_duplicate (bool, optional): whether to remove the duplicated: parameters in the result. Defaults to True.

Yields:

(str, Parameter): Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

Parameters:

prefix (str)
recurse (bool)
remove_duplicate (bool)

Return type:

parameters(recurse=True)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Args:

recurse (bool): if True, then yields parameters of this module: and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter: module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

Parameters:: recurse (bool)
Return type:: Iterator[Parameter]

register_backward_hook(hook)

Register a backward hook on the module.

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

Returns:

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

Parameters:: hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
Return type:: RemovableHandle

register_buffer(name, tensor, persistent=True)

Add a buffer to the module.

This is typically used to register a buffer that should not be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

Buffers can be accessed as attributes using given names.

Args:

name (str): name of the buffer. The buffer can be accessed: from this module using the given name
tensor (Tensor or None): buffer to be registered. If None, then operations: that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.
persistent (bool): whether the buffer is part of this module’s: state_dict.

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

Parameters:

name (str)
tensor (Tensor | None)
persistent (bool)

Return type:

register_forward_hook(hook, *, prepend=False, with_kwargs=False, always_call=False)

Register a forward hook on the module.

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature:

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature:

hook(module, args, kwargs, output) -> None or modified output

Args:

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

before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If True, the hook will be passed the: kwargs given to the forward function. Default: False
always_call (bool): If True the hook will be run regardless of: whether an exception is raised while calling the Module. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...], Any], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any], Any], Any | None])
prepend (bool)
with_kwargs (bool)
always_call (bool)

Return type:

RemovableHandle

register_forward_pre_hook(hook, *, prepend=False, with_kwargs=False)

Register a forward pre-hook on the module.

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature:

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature:

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

Args:

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

all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

with_kwargs (bool): If true, the hook will be passed the kwargs: given to the forward function. Default: False

Returns:

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

Parameters:

hook (Callable[[TypeVar(T, bound= Module), tuple[Any, ...]], Any | None] | Callable[[TypeVar(T, bound= Module), tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None])
prepend (bool)
with_kwargs (bool)

Return type:

RemovableHandle

register_full_backward_hook(hook, prepend=False)

Register a backward hook on the module.

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.

If none of the module inputs require gradients, the hook will fire when the gradients are computed with respect to module outputs.

If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_full_backward_pre_hook(hook, prepend=False)

Register a backward pre-hook on the module.

The hook will be called every time the gradients for the module are computed. The hook should have the following signature:

hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.

Warning

Modifying inputs inplace is not allowed when using backward hooks and will raise an error.

Args:

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

all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

Returns:

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

Parameters:

hook (Callable[[Module, tuple[Tensor, ...] | Tensor], None | tuple[Tensor, ...] | Tensor])
prepend (bool)

Return type:

RemovableHandle

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module’s load_state_dict() is called.

It should have the following signature::: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

Returns:

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

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module’s load_state_dict() is called.

It should have the following signature::

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Arguments:

hook (Callable): Callable hook that will be invoked before: loading the state dict.

register_module(name, module)

Alias for add_module().

Parameters:

name (str)
module (Module | None)

Return type:

register_parameter(name, param)

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Args:

name (str): name of the parameter. The parameter can be accessed: from this module using the given name
param (Parameter or None): parameter to be added to the module. If: None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

Parameters:

name (str)
param (Parameter | None)

Return type:

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify the state_dict inplace.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before the state_dict call is made.

requires_grad_(requires_grad=True)

Change if autograd should record operations on parameters in this module.

This method sets the parameters’ requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See Locally disabling gradient computation for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Args:

requires_grad (bool): whether autograd should record operations on: parameters in this module. Default: True.

Returns:

Module: self

Parameters:: requires_grad (bool)
Return type:: Self

set_extra_state(state)

Set extra state contained in the loaded state_dict.

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

Args:: state (dict): Extra state from the state_dict

Parameters:: state (Any)
Return type:: None

set_submodule(target, module, strict=False)

Set the submodule given by target if it exists, otherwise throw an error.

Note

If strict is set to False (default), the method will replace an existing submodule or create a new submodule if the parent module exists. If strict is set to True, the method will only attempt to replace an existing submodule and throw an error if the submodule does not exist.

For example, let’s say you have an nn.Module A that looks like this:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To override the Conv2d with a new submodule Linear, you could call set_submodule("net_b.net_c.conv", nn.Linear(1, 1)) where strict could be True or False

To add a new submodule Conv2d to the existing net_b module, you would call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1)).

In the above if you set strict=True and call set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True), an AttributeError will be raised because net_b does not have a submodule named conv.

Args:

target: The fully-qualified string name of the submodule: to look for. (See above example for how to specify a fully-qualified string.)

module: The module to set the submodule to. strict: If False, the method will replace an existing submodule

or create a new submodule if the parent module exists. If True, the method will only attempt to replace an existing submodule and throw an error if the submodule doesn’t already exist.

Raises:

ValueError: If the target string is empty or if module is not an instance of nn.Module. AttributeError: If at any point along the path resulting from

the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

Parameters:

target (str)
module (Module)
strict (bool)

Return type:

share_memory()

Return type:: Self

state_dict(*args, destination=None, prefix='', keep_vars=False)

Overloads:

self, destination (T_destination), prefix (str), keep_vars (bool) → T_destination
self, prefix (str), keep_vars (bool) → dict[str, Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Args:

destination (dict, optional): If provided, the state of module will: be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.
prefix (str, optional): a prefix added to parameter and buffer: names to compose the keys in state_dict. Default: ''.
keep_vars (bool, optional): by default the Tensor s: returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

dict:: a dictionary containing a whole state of the module

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

to(*args, **kwargs)

Overloads:

self, device (DeviceLikeType | None), dtype (dtype | None), non_blocking (bool) → Self
self, dtype (dtype), non_blocking (bool) → Self
self, tensor (Tensor), non_blocking (bool) → Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters: and buffers in this module
dtype (torch.dtype): the desired floating point or complex dtype of: the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired: dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format): the desired memory: format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*, device, recurse=True)

Move the parameters and buffers to the specified device without copying storage.

Args:

device (torch.device): The desired device of the parameters: and buffers in this module.
recurse (bool): Whether parameters and buffers of submodules should: be recursively moved to the specified device.

Returns:

Module: self

Parameters:

device (str | device | int | None)
recurse (bool)

Return type:

Self

train(mode=True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation: mode (False). Default: True.

Returns:

Module: self

Parameters:: mode (bool)
Return type:: Self

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Parameters:: dst_type (dtype | str)
Return type:: Self

xpu(device=None)

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

Note

This method modifies the module in-place.

Arguments:

device (int, optional): if specified, all parameters will be: copied to that device

Returns:

Module: self

Parameters:: device (int | device | None)
Return type:: Self

zero_grad(set_to_none=True)

Reset gradients of all model parameters.

See similar function under torch.optim.Optimizer for more context.

Args:

set_to_none (bool): instead of setting to zero, set the grads to None.: See torch.optim.Optimizer.zero_grad() for details.

Parameters:: set_to_none (bool)
Return type:: None

training: bool

ERB Scale Utilities

torch_amt.audfiltbw(fc)[source]

Compute equivalent rectangular bandwidth (ERB) of an auditory filter.

Calculates the auditory filter bandwidth based on the center frequency using the relationship defined by Glasberg and Moore (1990). This function computes the bandwidth in Hz, not the ERB-rate scale value.

\[\text{BW}(f_c) = 24.7 + \frac{f_c}{9.265}\]

where \(f_c\) is the center frequency in Hz.

Parameters:: fc (Tensor) – Center frequencies in Hz. Can be scalar or any tensor shape.
Returns:: Auditory filter bandwidths in Hz. Same shape as input.
Return type:: Tensor

Examples

>>> import torch
>>> fc = torch.tensor([100.0, 1000.0, 4000.0])
>>> bw = audfiltbw(fc)
>>> print(bw)
tensor([ 35.4933, 132.6331, 456.4323])

Notes

The bandwidth grows linearly with frequency above ~250 Hz. This relationship reflects the approximately constant-Q behavior of auditory filters at high frequencies and the approximately constant-bandwidth behavior at low frequencies.

References

torch_amt.erb2fc(erb)[source]

Convert ERB-rate scale to frequency in Hz (Natural Logarithm version).

This is the inverse transformation of fc2erb(), converting from the ERB-rate scale (Cams) back to frequency in Hz using the natural logarithm formulation from Glasberg and Moore (1990).

\[f = \frac{1}{0.00437} \left( e^{\frac{\text{ERB-rate}}{9.2645}} - 1 \right)\]

Parameters:: erb (Tensor) – ERB-rate scale values (Cams). Can be scalar or any tensor shape.
Returns:: Frequencies in Hz. Same shape as input.
Return type:: Tensor

Examples

>>> import torch
>>> from torch_amt.common.filterbanks import fc2erb, erb2fc
>>> fc_original = torch.tensor([100.0, 1000.0, 4000.0])
>>> erb_values = fc2erb(fc_original)
>>> fc_reconstructed = erb2fc(erb_values)
>>> print(torch.allclose(fc_original, fc_reconstructed, atol=1e-3))
True

Notes

Numerical Stability: Exponential function can overflow for very large ERB values (>40 Cams ≈ 22 kHz), but this is outside the typical auditory range.

This function uses the natural logarithm formulation. For the base-10 logarithm version used in loudness models, see erbrate2f().

See also

fc2erb: Forward transformation (Hz to ERB-rate).
erbrate2f: Base-10 logarithm version for loudness models.

References

torch_amt.fc2erb(fc)[source]

Convert frequency in Hz to ERB-rate scale (Natural Logarithm version).

Transforms frequency in Hz to the ERB-rate scale (Cams) using the natural logarithm formulation from Glasberg and Moore (1990). The ERB-rate represents the number of equivalent rectangular bandwidths below a given frequency.

\[\text{ERB-rate} = 9.2645 \cdot \ln(1 + f_c \cdot 0.00437)\]

where \(f_c\) is the frequency in Hz.

Parameters:: fc (Tensor) – Frequencies in Hz. Can be scalar or any tensor shape.
Returns:: ERB-rate scale values (Cams). Same shape as input.
Return type:: Tensor

Examples

>>> import torch
>>> fc = torch.tensor([100.0, 1000.0, 4000.0])
>>> erb = fc2erb(fc)
>>> print(erb)
tensor([ 3.3589, 15.5720, 27.0217])
>>> # Verify inverse relationship
>>> fc_back = erb2fc(erb)
>>> print(torch.allclose(fc, fc_back, atol=1e-3))
True

Notes

Numerical Stability: Logarithm is stable for all positive frequencies. For fc=0, result is 0 (continuous at origin).

The ERB-rate scale provides an approximately linear representation of perceived pitch. One unit on the ERB scale corresponds roughly to one critical bandwidth or one auditory filter. The unit is often called “Cams” (after Cambridge).

This function uses the natural logarithm formulation. For the base-10 logarithm version used in loudness models (Moore et al., 1997), see f2erbrate().

See also

erb2fc: Inverse transformation (ERB-rate to Hz).
f2erbrate: Base-10 logarithm version for loudness models.

References

torch_amt.f2erb(f)[source]

Calculate equivalent rectangular bandwidth (ERB) at a given frequency.

Computes the auditory filter bandwidth in Hz according to the formula from Glasberg and Moore (1990) and Moore et al. (1997). This is the same calculation as audfiltbw() but uses the alternative parametrization.

\[\text{ERB}(f) = 24.7 \cdot \left(4.37 \frac{f}{1000} + 1\right)\]

where \(f\) is the frequency in Hz.

Parameters:: f (Tensor) – Frequencies in Hz. Can be scalar or any tensor shape.
Returns:: ERB bandwidths in Hz. Same shape as input.
Return type:: Tensor

Examples

>>> import torch
>>> f = torch.tensor([100.0, 1000.0, 4000.0])
>>> erb_bw = f2erb(f)
>>> print(erb_bw)
tensor([ 35.4939, 132.6390, 456.4560])
>>> # Compare with audfiltbw (should be nearly identical)
>>> from torch_amt.common.filterbanks import audfiltbw
>>> print(torch.allclose(erb_bw, audfiltbw(f), atol=1e-2))
True

Notes

Relation to audfiltbw: This function and audfiltbw() compute the same quantity using slightly different parameterizations. The difference is negligible (<0.1 Hz) across the auditory range.

The ERB bandwidth represents the width in Hz of a rectangular filter that would pass the same total power as the rounded exponential auditory filter. One ERB is also referred to as 1 Cam in the loudness literature.

Note: This function returns bandwidth in Hz. For the ERB-rate scale (number of ERBs from DC), use fc2erb() or f2erbrate().

See also

audfiltbw: Equivalent function with alternative parametrization.
fc2erb: Frequency to ERB-rate scale (cumulative ERBs from DC).

References

torch_amt.f2erbrate(f)[source]

Convert frequency to ERB-rate using Base-10 Logarithm.

This is the base-10 logarithm variant used specifically in loudness models such as Moore et al. (1997), Glasberg and Moore (2002), and Moore et al. (2016). It differs from fc2erb() which uses natural logarithm.

\[E = 21.366 \cdot \log_{10}(4.368 \cdot f_{\text{kHz}} + 1)\]

where \(f_{\text{kHz}}\) is the frequency in kHz.

Parameters:: f (Tensor) – Frequencies in Hz. Can be scalar or any tensor shape.
Returns:: ERB-rate values in Cam units. Same shape as input.
Return type:: Tensor

Examples

>>> import torch
>>> f = torch.tensor([100.0, 1000.0, 4000.0])
>>> erbrate = f2erbrate(f)
>>> print(erbrate)
tensor([ 3.3629, 15.5932, 27.0603])
>>> # Verify inverse
>>> f_back = erbrate2f(erbrate)
>>> print(torch.allclose(f, f_back, atol=1e-2))
True

Notes

Numerical Stability: Stable for all positive frequencies. The log10 function is well-behaved for the argument range encountered in auditory applications (0.44 to 435 for 0-20 kHz).
Difference from fc2erb: This function produces slightly different values (typically <0.05 Cams difference) compared to fc2erb() due to the different logarithm base and constants. The base-10 version is standard in loudness modeling.

The Cam scale (Cambridge ERB-rate scale) represents the cumulative number of equivalent rectangular bandwidths from DC to a given frequency. The human cochlea has approximately 40-43 Cams of frequency resolution from 20 Hz to 20 kHz.

See also

erbrate2f: Inverse transformation (ERB-rate to Hz).
fc2erb: Natural logarithm version.

References

torch_amt.erbrate2f(erbrate)[source]

Convert ERB-rate (Base-10 Logarithm) to frequency in Hz.

This is the inverse transformation of f2erbrate(), converting from the ERB-rate scale (Cams) back to frequency in Hz. Used in loudness models that employ the base-10 logarithm formulation.

\[f = \frac{10^{E/21.366} - 1}{4.368} \cdot 1000\]

where \(E\) is the ERB-rate in Cams and \(f\) is returned in Hz.

Parameters:: erbrate (Tensor) – ERB-rate values in Cam units. Can be scalar or any tensor shape.
Returns:: Frequencies in Hz. Same shape as input.
Return type:: Tensor

Examples

>>> import torch
>>> from torch_amt.common.filterbanks import f2erbrate, erbrate2f
>>> f_original = torch.tensor([100.0, 1000.0, 4000.0])
>>> erbrate = f2erbrate(f_original)
>>> f_reconstructed = erbrate2f(erbrate)
>>> print(torch.allclose(f_original, f_reconstructed, atol=1e-2))
True
>>> # Typical range: 0-43 Cams covers 0-20 kHz
>>> print(erbrate2f(torch.tensor([0.0, 21.5, 43.0])))
tensor([    0.,  3463., 20031.])

Notes

Numerical Stability: The exponential operation (10^x) can produce very large values for high ERB-rates. For erbrate=43 Cams (approximately 20 kHz), the intermediate result is ~10^2, which is well within floating-point range. Overflow only occurs for unrealistic ERB values >100 Cams.

This function is the inverse of f2erbrate() and is used in loudness models to convert back from the perceptual ERB-rate scale to physical frequency. For the natural logarithm version, see erb2fc().

See also

f2erbrate: Forward transformation (Hz to ERB-rate).
erb2fc: Natural logarithm version.

References

torch_amt.erbspacebw(flow, fhigh, bwmul=1.0, basef=None, device=None, dtype=torch.float32)[source]

Generate ERB-spaced center frequencies for auditory filterbanks.

Creates a vector of frequencies spaced equidistantly on the ERB-rate scale, which provides perceptually uniform spacing that matches the frequency resolution of the human auditory system. This is the standard method for selecting center frequencies in auditory models.

\[f_i = \text{erb2fc}(\text{fc2erb}(f_{\text{low}}) + i \cdot \text{bwmul})\]

for \(i = 0, 1, 2, \ldots, N-1\) where \(N\) is determined by the frequency range and spacing density.

Parameters:

flow (float) – Lowest center frequency in Hz. Must be positive.
fhigh (float) – Highest center frequency in Hz. Must be greater than flow.
bwmul (float) – Spacing density in ERB units. Default is 1.0 (one filter per ERB). - bwmul < 1.0: Denser spacing (overlapping filters). - bwmul = 1.0: Standard spacing (adjacent ERBs). - bwmul > 1.0: Sparser spacing (gaps between filters).
basef (float | None) – Reference frequency in Hz. If provided, one filter is placed exactly at this frequency, with others spaced symmetrically above and below. Useful for ensuring coverage of specific frequencies of interest. Default is None (uniform spacing from flow to fhigh).
device (device | None) – Torch device for tensor creation (CPU, CUDA, or MPS). Default is None (uses CPU).
dtype (dtype) – Data type for output tensor. Default is torch.float32.

Returns:

Vector of center frequencies in Hz, monotonically increasing. Shape: [num_channels] where num_channels depends on the frequency range and bwmul.

Return type: