Loudness Processing

Compression, specific loudness transformation, and temporal integration components.

Compression

BrokenStickCompression

class torch_amt.BrokenStickCompression(knee_db=30.0, exponent=0.3, dboffset=100.0, num_channels=None, learnable=False)

Bases: Module

Broken-stick compression for auditory nerve fiber dynamics.

Implements a piecewise power-law compression where signals below a knee point pass through unchanged (linear), while signals above the knee are compressed using a power-law function. This creates a “broken stick” transfer function that simulates the nonlinear input-output characteristics of inner hair cells and auditory nerve fibers.

Algorithm Overview

The compression is applied element-wise:

\[\begin{split}y(t) = \\begin{cases} x(t) & \\text{if } |x(t)| \\leq \\text{knee} \\\\ \\text{sign}(x(t)) \\cdot |x(t)|^n \\cdot \\text{knee}^{1-n} & \\text{if } |x(t)| > \\text{knee} \\end{cases}\end{split}\]

where: - \(n\) is the compression exponent (typically 0.3) - \(\\text{knee}\) is the threshold in linear units - The formula ensures continuity at the knee point

Knee point calculation:

\[\begin{split}\\text{knee}_{\\text{linear}} = 10^{(\\text{knee}_{\\text{dB}} - \\text{dboffset})/20}\end{split}\]

param knee_db:

Knee point in dB relative to dboffset. Default: 30 dB.

type knee_db:

float

param exponent:

Compression exponent n. Default: 0.3.

• n < 1: compression above knee

• n = 1: linear (no compression)

• n > 1: expansion above knee

type exponent:

float | Tensor

param dboffset:

Reference level in dB SPL for full scale. Default: 100 dB. AMT convention uses 94 dB, King2019 uses 100 dB.

type dboffset:

float

param num_channels:

Number of frequency channels for per-channel exponents. Default: None. If provided, exponent can vary across channels.

type num_channels:

int | None

param learnable:

If True, both exponent and knee_db become learnable nn.Parameters. Default: False.

type learnable:

bool

knee

Linear knee point value, shape ().

Type:

torch.Tensor

exponent

Compression exponent, shape () or (num_channels,).

Type:

torch.Tensor

knee_db

Knee point in dB (stored for reference).

Type:

float

dboffset

Reference level in dB SPL.

Type:

float

Shape

-----

- Input

• \(T\) = time samples

• \(C\) = channels (optional)

Type:

\((T, C)\) or \((T,)\) where

- Output

Type:

Same shape as input

Notes

Compression behavior:

• Below knee: No compression (linear passthrough)

• Above knee: Power-law compression with exponent n

• For King et al. (2019): knee_db=30, exponent=0.3, dboffset=100

Per-channel exponents:

Exponent can be a scalar (same for all channels) or a tensor with one value per channel, allowing frequency-dependent compression characteristics.

See also

PowerCompression

Full power-law compression without knee

References

[1]

A. King, L. Varnet, and C. Lorenzi, “Accounting for masking of frequency modulation by amplitude modulation with the modulation filter-bank concept,” J. Acoust. Soc. Am., vol. 145, no. 4, pp. 2277-2293, 2019.

Examples

>>> # Standard King2019 compression
>>> comp = BrokenStickCompression(knee_db=30, exponent=0.3, dboffset=100)
>>> signal = torch.randn(1000, 31)
>>> compressed = comp(signal)
>>> print(compressed.shape)
torch.Size([1000, 31])

>>> # Per-channel compression (different exponent per channel)
>>> exponents = torch.linspace(0.2, 0.4, 31)  # Vary from 0.2 to 0.4
>>> comp_perchan = BrokenStickCompression(exponent=exponents, num_channels=31)
>>> compressed = comp_perchan(signal)

__init__(knee_db=30.0, exponent=0.3, dboffset=100.0, num_channels=None, learnable=False)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• knee_db (float)

• exponent (float | Tensor)

• dboffset (float)

• num_channels (int | None)

• learnable (bool)

forward(x)

Apply broken-stick compression.

Parameters:

x (Tensor) – Input signal of shape (time, channels) or (time,).

Returns:

Compressed signal with same shape as input.

Return type:

Tensor

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

Return type:

str

get_parameters()

Get compression parameters.

Returns:

Dictionary with knee_db, exponent (mean), dboffset, learnable

Return type:

dict

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

PowerCompression

class torch_amt.PowerCompression(knee_db=30.0, exponent=0.3, dboffset=100.0, num_channels=None, learnable=False)

Bases: Module

Full power-law compression for auditory processing.

Applies a power-law compression/expansion to the entire signal without a linear region. Unlike BrokenStickCompression, this affects all signal levels and can cause expansion below the knee when exponent < 1.

Algorithm Overview

The compression is applied globally:

\[\begin{split}y(t) = \\text{sign}(x(t)) \\cdot \\left|\\frac{x(t)}{\\text{knee}}\\right|^n \\cdot \\text{knee}\end{split}\]

where: - \(n\) is the power-law exponent - \(\\text{knee}\) is the reference level in linear units

Compression characteristics:

• For n < 1: Compresses signals above knee, expands signals below knee

• For n = 1: Linear (no effect)

• For n > 1: Expands signals above knee, compresses signals below knee

param knee_db:

Reference level in dB relative to dboffset. Default: 30 dB.

type knee_db:

float

param exponent:

Power-law exponent n. Default: 0.3.

• n < 1: compression above knee, expansion below

• n = 1: linear (identity transform)

• n > 1: expansion above knee, compression below

type exponent:

float | Tensor

param dboffset:

Reference level in dB SPL for full scale. Default: 100 dB SPL.

type dboffset:

float

param num_channels:

Number of frequency channels for per-channel exponents. Default: None.

type num_channels:

int | None

param learnable:

If True, both exponent and knee_db become learnable nn.Parameters. Default: False.

type learnable:

bool

knee

Linear reference level, shape ().

Type:

torch.Tensor

exponent

Power-law exponent, shape () or (num_channels,).

Type:

torch.Tensor

knee_db

Reference level in dB.

Type:

float

dboffset

dB SPL reference.

Type:

float

Shape

-----

- Input

• \(T\) = time samples

• \(C\) = channels (optional)

Type:

\((T, C)\) or \((T,)\) where

- Output

Type:

Same shape as input

Notes

Difference from BrokenStickCompression:

• PowerCompression: Affects all signal levels, no linear region

• BrokenStickCompression: Linear below knee, compressed above knee

WARNING: With typical exponent values (n < 1), this compression can amplify low-level signals (expansion below knee), which may not be physiologically accurate for auditory modeling.

Usage in King et al. (2019):

This compression type is specific to the PEMO model. The Dau et al. (1997) model does not use compression at this stage.

See also

BrokenStickCompression

Piecewise compression with linear region below knee

References

[1]

A. King, L. Varnet, and C. Lorenzi, “Accounting for masking of frequency modulation by amplitude modulation with the modulation filter-bank concept,” J. Acoust. Soc. Am., vol. 145, no. 4, pp. 2277-2293, 2019.

Examples

>>> # Standard power-law compression
>>> comp = PowerCompression(knee_db=30, exponent=0.3, dboffset=100)
>>> signal = torch.randn(1000, 31)
>>> compressed = comp(signal)
>>> print(compressed.shape)
torch.Size([1000, 31])

__init__(knee_db=30.0, exponent=0.3, dboffset=100.0, num_channels=None, learnable=False)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• knee_db (float)

• exponent (float | Tensor)

• dboffset (float)

• num_channels (int | None)

• learnable (bool)

forward(x)

Apply power-law compression.

Parameters:

x (Tensor) – Input signal of shape (time, channels) or (time,).

Returns:

Compressed signal with same shape as input.

Return type:

Tensor

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

Return type:

str

get_parameters()

Get compression parameters.

Returns:

Dictionary with knee_db, exponent (mean), dboffset, learnable

Return type:

dict

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

Specific Loudness

SpecificLoudness

class torch_amt.SpecificLoudness(fs=32000, f_min=50.0, f_max=15000.0, erb_step=0.25, learnable=False)

Bases: Module

Specific loudness transformation for Glasberg & Moore (2002) loudness model.

Transforms excitation pattern from ERB filterbank to specific loudness using a three-regime model that accounts for absolute threshold, linear region, and compressive region. Based on Moore & Glasberg (1997) model with ISO 226 threshold curves.

Algorithm Overview

The transformation uses a piecewise function with three regimes:

Regime 1: Sub-threshold (E ≤ E_Thrq)

\[N(f,t) = 0\]

Regime 2: Linear region (E_Thrq < E ≤ E_Thrq + E_0)

\[\begin{split}N(f,t) = C \\cdot (E - E_{Thrq})\end{split}\]

Regime 3: Compressive region (E > E_Thrq + E_0)

\[\begin{split}N(f,t) = C \\cdot E_0^{1-\\alpha} \\cdot (E - E_{Thrq})^{\\alpha}\end{split}\]

where:

• \(E(f,t)\) is the excitation level in dB SPL at ERB frequency \(f\)

• \(E_{Thrq}(f)\) is the absolute threshold in quiet (ISO 226), frequency-dependent

• \(C = 0.047\) is the gain constant (Moore & Glasberg 1997)

• \(\\alpha = 0.2\) is the compression exponent

• \(E_0 = 10\) dB is the transition point from linear to compressive regime

• \(N(f,t)\) is the specific loudness in sone/ERB

Absolute Threshold Computation (ISO 226 approximation):

\[\begin{split}E_{Thrq}(f) = 3.64 \\left(\\frac{f}{1000}\\right)^{-0.8} - 6.5 \\exp\\left[-0.6\\left(\\frac{f}{1000} - 3.3\\right)^2\\right] + 10^{-3} \\left(\\frac{f}{1000}\\right)^4\end{split}\]

param fs:

Sampling rate in Hz. Default: 32000

type fs:

int

param f_min:

Minimum frequency for ERB filterbank in Hz. Default: 50.0

type f_min:

float

param f_max:

Maximum frequency for ERB filterbank in Hz. Default: 15000.0

type f_max:

float

param erb_step:

ERB frequency step for filterbank spacing. Default: 0.25 ERB

type erb_step:

float

param learnable:

If True, C, α, E_0, and threshold adjustments become learnable parameters. Default: False

type learnable:

bool

fc_erb

ERB filterbank center frequencies in Hz, shape (n_erb_bands,)

Type:

torch.Tensor

n_erb_bands

Number of ERB frequency channels

Type:

int

ethrq_base

Base absolute threshold in quiet (ISO 226), shape (n_erb_bands,)

Type:

torch.Tensor

ethrq_adjustment

Additive adjustment to threshold in dB, shape (n_erb_bands,)

Type:

torch.Tensor or nn.Parameter

C

Gain constant, scalar. Fixed at 0.047 (Moore & Glasberg 1997)

Type:

torch.Tensor or nn.Parameter

alpha

Compression exponent, scalar. Fixed at 0.2

Type:

torch.Tensor or nn.Parameter

E0_offset

Transition point from linear to compressive regime in dB above threshold. Fixed at 10.0 dB

Type:

torch.Tensor or nn.Parameter

Input Shape

-----------

excitation

Excitation pattern in dB SPL, shape (batch, n_frames, n_erb_bands)

Type:

torch.Tensor

Output Shape

------------

specific_loudness

Specific loudness in sone/ERB, shape (batch, n_frames, n_erb_bands)

Type:

torch.Tensor

Notes

• The three-regime model captures the transition from complete masking (sub-threshold), through a linear loudness growth region, to the compressive loudness region

• The absolute threshold \(E_{Thrq}\) is frequency-dependent and follows ISO 226, with minimum threshold (~4 dB SPL) around 2-5 kHz

• The linear region (Regime 2) extends from threshold to ~10 dB above threshold

• The compressive region (Regime 3) with \(\\alpha=0.2\) implements the well-known power-law loudness growth (approximately doubling loudness per 10 dB)

• When learnable=True, the model can adapt the threshold, gain, and compression characteristics through backpropagation

• This implementation is compatible with glasberg2002 model from AMT MATLAB

See also

Moore2016SpecificLoudness

Specific loudness for Moore2016 model (ANSI S3.4-2007)

LoudnessIntegration

Temporal integration for Glasberg2002 model

auditoryfilterbank

ERB gammatone filterbank preprocessing

Examples

Basic usage with default Glasberg2002 parameters:

>>> import torch
>>> from torch_amt.common.loudness import SpecificLoudness
>>>
>>> # Create module
>>> spec_loud = SpecificLoudness(fs=32000, f_min=50, f_max=15000, erb_step=0.25)
>>> print(f"Number of ERB bands: {spec_loud.n_erb_bands}")
Number of ERB bands: 150
>>>
>>> # Simulate excitation pattern (e.g., from ERB filterbank)
>>> batch, n_frames, n_erb = 2, 100, 150
>>> excitation_db = torch.randn(batch, n_frames, n_erb) * 20 + 60  # ~60 dB SPL mean
>>>
>>> # Transform to specific loudness
>>> N = spec_loud(excitation_db)
>>> print(f"Specific loudness shape: {N.shape}, range: [{N.min():.3f}, {N.max():.3f}] sone/ERB")
Specific loudness shape: torch.Size([2, 100, 150]), range: [0.000, 15.234] sone/ERB

Check absolute threshold in quiet:

>>> threshold = spec_loud.get_threshold()
>>> print(f"Threshold at 1 kHz: {threshold[spec_loud.fc_erb.argmin((spec_loud.fc_erb - 1000).abs())]:.2f} dB SPL")
Threshold at 1 kHz: 4.23 dB SPL
>>> print(f"Min threshold: {threshold.min():.2f} dB SPL at {spec_loud.fc_erb[threshold.argmin()]:.0f} Hz")
Min threshold: 3.85 dB SPL at 3500 Hz

Learnable parameters for model adaptation:

>>> spec_loud_learn = SpecificLoudness(fs=32000, learnable=True)
>>> params = spec_loud_learn.get_parameters()
>>> print(f"C={params['C']:.4f}, alpha={params['alpha']:.3f}, E0={params['E0_offset']:.1f} dB")
C=0.0470, alpha=0.200, E0=10.0 dB
>>>
>>> # Can now train these parameters with backpropagation
>>> optimizer = torch.optim.Adam(spec_loud_learn.parameters(), lr=1e-3)

References

[1]

Glasberg, B. R., & Moore, B. C. (2002). A model of loudness applicable to time-varying sounds. Journal of the Audio Engineering Society, 50(5), 331-342.

[2]

Moore, B. C. J., Glasberg, B. R., & Baer, T. (1997). A Model for the Prediction of Thresholds, Loudness, and Partial Loudness. J. Audio Eng. Soc, 45(4), 224-240.

[3]

ISO 226:2003. Acoustics – Normal equal-loudness-level contours.

__init__(fs=32000, f_min=50.0, f_max=15000.0, erb_step=0.25, learnable=False)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• fs (int)

• f_min (float)

• f_max (float)

• erb_step (float)

• learnable (bool)

forward(excitation)

Transform excitation to specific loudness.

Parameters:

excitation (Tensor) – Excitation in dB SPL, shape (batch, n_frames, n_erb_bands)

Returns:

Specific loudness in sone/ERB, shape (batch, n_frames, n_erb_bands)

Return type:

Tensor

get_threshold()

Get absolute threshold in quiet for all ERB channels.

Returns:

Threshold in dB SPL, shape (n_erb_bands,)

Return type:

Tensor

get_parameters()

Get model parameters.

Returns:

Dictionary with C, alpha, E0_offset, learnable

Return type:

dict

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

Moore2016SpecificLoudness

class torch_amt.Moore2016SpecificLoudness(learnable=False, dtype=torch.float32)

Bases: Module

Specific loudness transformation for Moore et al. (2016) binaural loudness model.

Implements the ANSI S3.4-2007 specific loudness transformation with three loudness regimes (sub-threshold, standard, high-level) and frequency-dependent parameters derived from lookup tables. Uses binaural constant C = 0.0631 (Moore & Glasberg 2007).

Algorithm Overview

The specific loudness is computed using three regimes based on excitation level:

Regime 1: Sub-threshold (E < E_Thrq)

\[\begin{split}N_2(f) = C \\cdot \\left(\\frac{2E}{E + E_{Thrq}}\\right)^{1.5} \\cdot \\left[(G(f)\\cdot E + A(f))^{\\alpha(f)} - A(f)^{\\alpha(f)}\\right]\end{split}\]

Regime 2: Standard above-threshold (E_Thrq ≤ E < 10^{10})

\[\begin{split}N_1(f) = C \\cdot \\left[(G(f)\\cdot E + A(f))^{\\alpha(f)} - A(f)^{\\alpha(f)}\\right]\end{split}\]

Regime 3: Very high level (E ≥ 10^{10})

\[\begin{split}N_3(f) = C \\cdot \\left(\\frac{E}{1.0707}\\right)^{0.2}\end{split}\]

where:

• \(E(f)\) is the excitation (linear scale, not dB) at ERB frequency \(f\)

• \(E_{Thrq}(f)\) is the absolute threshold in quiet (ISO 226), frequency-dependent

• \(C = 0.0631\) is the binaural constant (Moore & Glasberg 2007)

• \(G(f)\) is the low-level gain parameter from lookup table (150 values)

• \(\\alpha(f)\) is the compression exponent from lookup table, range: [0.2, 0.267]

• \(A(f)\) is the additive constant from lookup table (88 unique values)

• \(N(f)\) is the specific loudness in sone/ERB

Frequency-Dependent Parameters:

The model uses three lookup tables (G, α, A) derived from ANSI S3.4-2007:

1. G(f): Low-level gain, computed piecewise:

   \[\begin{split}G(f) = \\begin{cases} 10^{(\\text{ERB}_c - 13)/15} & \\text{ERB}_c < 13 \\\\ 10^{(\\text{ERB}_c - 13)/7.5} & \\text{ERB}_c \\geq 13 \\end{cases}\end{split}\]

   where \(\\text{ERB}_c = 21.366 \\cdot \\log_{10}(f/228.7 + 1)\)

2. α(f): Compression exponent, interpolated from 6-point lookup table based on G

   • Low frequencies (low G): α ≈ 0.267 (stronger compression)

   • High frequencies (high G): α ≈ 0.200 (weaker compression)

3. A(f): Additive constant, interpolated from 88-point lookup table based on G

   • Range: [10^{4.72} to 10^{8.85}] (linear scale)

   • Ensures smooth transition between loudness regimes

Absolute Threshold (ISO 226 approximation):

\[\begin{split}E_{Thrq}(f) = 3.64 \\left(\\frac{f}{1000}\\right)^{-0.8} - 6.5 \\exp\\left[-0.6\\left(\\frac{f}{1000} - 3.3\\right)^2\\right] + 10^{-3} \\left(\\frac{f}{1000}\\right)^4\end{split}\]

param learnable:

If True, C and lookup table parameters (G, Alpha, A) become learnable. Default: False

type learnable:

bool

param dtype:

Data type for computations. Default: torch.float32

type dtype:

dtype

erb_scale

ERB scale from 1.75 to 39 in 0.25 steps, shape (150,)

Type:

torch.Tensor

fc

Center frequencies in Hz for each ERB channel, shape (150,)

Type:

torch.Tensor

G

Low-level gain parameter for each channel, shape (150,)

Type:

torch.Tensor or nn.Parameter

Alpha

Compression exponent for each channel (frequency-dependent), shape (150,)

Type:

torch.Tensor or nn.Parameter

A

Additive constant for each channel, shape (150,)

Type:

torch.Tensor or nn.Parameter

threshold_db

Absolute threshold in dB SPL for each channel, shape (150,)

Type:

torch.Tensor

C

Binaural loudness constant (0.0631)

Type:

float or nn.Parameter

Input Shape

-----------

excitation_db

Excitation pattern in dB SPL, shape (batch, 150) or (150,) Typically from Moore2016ExcitationPattern output

Type:

torch.Tensor

Output Shape

------------

specific_loudness

Specific loudness in sone/ERB, same shape as input

Type:

torch.Tensor

Notes

• This module operates on single time frames with 150 ERB channels

• Input shape is (batch, 150), different from Glasberg2002 SpecificLoudness which processes time series (batch, n_frames, n_erb)

• The 150 channels correspond to ERB scale 1.75 to 39 in 0.25 ERB steps

• Lookup tables G, Alpha, A are derived from ANSI S3.4-2007 standard

• The three-regime model provides smooth transitions: * Sub-threshold: Gradual onset with threshold-dependent weighting * Standard: Main loudness growth with frequency-dependent compression * High-level: Simplified power-law to prevent overflow

• Binaural constant C = 0.0631 accounts for binaural summation (approximately √2 loudness increase for identical binaural signals)

See also

SpecificLoudness

Specific loudness for Glasberg2002 model

Moore2016BinauralLoudness

Complete Moore2016 binaural loudness pipeline

Moore2016AGC

Automatic gain control for Moore2016 model

Examples

Basic usage with Moore2016 model:

>>> import torch
>>> from torch_amt.common.loudness import Moore2016SpecificLoudness
>>>
>>> # Create module (150 ERB channels fixed)
>>> spec_loud = Moore2016SpecificLoudness()
>>> params = spec_loud.get_parameters()
>>> print(f"Channels: {params['n_channels']}, C={params['C']:.4f}")
Channels: 150, C=0.0631
>>> print(f"Alpha range: [{params['Alpha_min']:.3f}, {params['Alpha_max']:.3f}]")
Alpha range: [0.200, 0.267]

Process single-frame excitation pattern:

>>> # Simulate excitation from Moore2016ExcitationPattern
>>> batch = 4
>>> excitation_db = torch.randn(batch, 150) * 15 + 60  # ~60 dB SPL mean
>>> N = spec_loud(excitation_db)
>>> print(f"Specific loudness shape: {N.shape}, range: [{N.min():.3f}, {N.max():.3f}] sone/ERB")
Specific loudness shape: torch.Size([4, 150]), range: [0.000, 18.456] sone/ERB

1D input (single excitation pattern):

>>> excitation_1d = torch.randn(150) * 10 + 55
>>> N_1d = spec_loud(excitation_1d)
>>> print(f"Output shape: {N_1d.shape}")
Output shape: torch.Size([150])

Check frequency-dependent parameters:

>>> print(f"Low-freq (100 Hz): G={spec_loud.G[0]:.3f}, Alpha={spec_loud.Alpha[0]:.3f}")
Low-freq (100 Hz): G=0.234, Alpha=0.267
>>> print(f"High-freq (10 kHz): G={spec_loud.G[-10]:.3f}, Alpha={spec_loud.Alpha[-10]:.3f}")
High-freq (10 kHz): G=8.456, Alpha=0.203

Learnable parameters for model adaptation:

>>> spec_loud_learn = Moore2016SpecificLoudness(learnable=True)
>>> # Can train C, G, Alpha, A with backpropagation
>>> optimizer = torch.optim.Adam(spec_loud_learn.parameters(), lr=1e-4)

References

[1]

Moore, B. C. J., Glasberg, B. R., & Schlittenlacher, J. (2016). A model of binaural loudness perception based on the outputs of an auditory periphery model. Acta Acustica united with Acustica, 102(5), 824-837.

[2]

Moore, B. C. J., & Glasberg, B. R. (2007). Modeling binaural loudness. J. Acoust. Soc. Am., 121(3), 1604-1612.

[3]

ANSI S3.4-2007. Procedure for the Computation of Loudness of Steady Sounds. American National Standards Institute.

[4]

ISO 226:2003. Acoustics – Normal equal-loudness-level contours.

__init__(learnable=False, dtype=torch.float32)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• learnable (bool)

• dtype (dtype)

forward(excitation_db)

Transform excitation pattern to specific loudness.

Parameters:

excitation_db (Tensor) – Excitation pattern in dB SPL. Shape: (batch, 150) or (150,)

Returns:

specific_loudness – Specific loudness in sone/ERB. Same shape as input.

Return type:

Tensor

get_parameters()

Get model parameters.

Returns:

params – Dictionary containing C, G range, Alpha range, A range

Return type:

dict

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

Binaural Processing

SpatialSmoothing

class torch_amt.SpatialSmoothing(kernel_width=18.0, sigma=0.08, learnable=False, dtype=torch.float32)

Bases: Module

Gaussian smoothing over ERB frequency channels for spatial integration.

Applies a Gaussian kernel to smooth specific loudness across frequency channels, implementing the spatial integration mechanism of the auditory system. This models the spread of excitation across the auditory nerve and represents the limited frequency selectivity of loudness integration.

Algorithm Overview

The spatial smoothing operation is a 1D convolution with a Gaussian kernel:

\[\begin{split}N_{smooth}(f) = \\sum_{g} W(g) \\cdot N(f + g)\end{split}\]

where the Gaussian kernel is defined as:

\[\begin{split}W(g) = \\frac{\\exp\\left[-(\\sigma \\cdot g)^2\\right]}{\\sum_g \\exp\\left[-(\\sigma \\cdot g)^2\\right]}\end{split}\]

• \(N(f)\) is the specific loudness at ERB frequency \(f\) (sone/ERB)

• \(g\) is the distance in ERB steps, ranging from \(-w\) to \(+w\)

• \(w\) = kernel_width (default: 18.0 ERB)

• \(\\sigma\) = 0.08 is the Gaussian standard deviation parameter

• \(W(g)\) is normalized to sum to 1 (energy-preserving)

• ERB step size = 0.25 ERB, so kernel spans \(2w/0.25 + 1 = 145\) channels

The kernel is applied via reflection padding at boundaries to avoid edge artifacts.

param kernel_width:

Half-width of the kernel in ERB units. Default: 18.0 Spans ±18 ERB, corresponding to ±72 channels (18 / 0.25 = 72)

type kernel_width:

float

param sigma:

Standard deviation parameter for Gaussian kernel. Default: 0.08

type sigma:

float

param learnable:

If True, both sigma and kernel_width become learnable parameters. Default: False

type learnable:

bool

param dtype:

Data type for computations. Default: torch.float32

type dtype:

dtype

gaussian_kernel

Normalized Gaussian kernel weights, shape (kernel_size,) Pre-computed when learnable=False, computed dynamically when learnable=True

Type:

torch.Tensor or None

sigma

Gaussian standard deviation parameter, scalar

Type:

torch.Tensor or nn.Parameter

g

ERB distance array from -kernel_width to +kernel_width in 0.25 steps Only used when learnable=True

Type:

torch.Tensor

kernel_width

Half-width of the kernel in ERB units

Type:

float

Input Shape

-----------

specific_loudness

Specific loudness in sone/ERB, shape (batch, 150) or (150,) Typically from Moore2016SpecificLoudness output

Type:

torch.Tensor

Output Shape

------------

smoothed_loudness

Spatially smoothed specific loudness, same shape as input

Type:

torch.Tensor

Notes

• The Gaussian kernel provides smooth spatial integration across frequency

• Kernel width of 18 ERB corresponds to 72 channels (18 / 0.25 = 72)

• Total kernel size: 2 * 72 + 1 = 145 channels

• Reflection padding at boundaries preserves energy and avoids edge artifacts

• The kernel is normalized to sum to 1, ensuring energy conservation

• Small sigma (0.08) means the Gaussian drops off rapidly, providing localized smoothing

• At g = ±18 ERB (kernel edges), W ≈ exp(-(0.08*18)²) ≈ 0.10 (10% of center weight)

See also

Moore2016SpecificLoudness

Computes specific loudness before smoothing

BinauralInhibition

Applies binaural inhibition after smoothing

Moore2016BinauralLoudness

Complete pipeline including this module

Examples

Basic usage with default Moore2016 parameters:

>>> import torch
>>> from torch_amt.common.loudness import SpatialSmoothing
>>>
>>> # Create module
>>> smoothing = SpatialSmoothing(kernel_width=18.0, sigma=0.08)
>>> params = smoothing.get_parameters()
>>> print(f"Kernel size: {params['kernel_size']}, sigma: {params['sigma']:.3f}")
Kernel size: 145, sigma: 0.080

Apply smoothing to specific loudness:

>>> # Simulate specific loudness (from Moore2016SpecificLoudness)
>>> batch = 2
>>> N_specific = torch.randn(batch, 150).abs() * 5  # Non-negative loudness
>>> N_smooth = smoothing(N_specific)
>>> print(f"Input shape: {N_specific.shape}, Output shape: {N_smooth.shape}")
Input shape: torch.Size([2, 150]), Output shape: torch.Size([2, 150])
>>> print(f"Before: [{N_specific[0, 70:75].tolist()}]")
Before: [[2.34, 5.67, 1.23, 8.90, 3.45]]
>>> print(f"After: [{N_smooth[0, 70:75].tolist()}]")
After: [[3.52, 4.31, 4.76, 4.98, 4.52]]

Energy conservation check:

>>> energy_before = N_specific.sum()
>>> energy_after = N_smooth.sum()
>>> print(f"Energy before: {energy_before:.2f}, after: {energy_after:.2f}")
Energy before: 1500.34, after: 1500.34
>>> print(f"Relative difference: {abs(energy_before - energy_after) / energy_before * 100:.4f}%")
Relative difference: 0.0001%

1D input (single loudness pattern):

>>> N_1d = torch.randn(150).abs() * 3
>>> N_smooth_1d = smoothing(N_1d)
>>> print(f"1D output shape: {N_smooth_1d.shape}")
1D output shape: torch.Size([150])

Learnable sigma for model adaptation:

>>> smoothing_learn = SpatialSmoothing(kernel_width=18.0, sigma=0.08, learnable=True)
>>> # Kernel computed dynamically in forward() based on learned sigma
>>> optimizer = torch.optim.Adam(smoothing_learn.parameters(), lr=1e-3)

References

[1]

Moore, B. C. J., Glasberg, B. R., & Schlittenlacher, J. (2016). A model of binaural loudness perception based on the outputs of an auditory periphery model. Acta Acustica united with Acustica, 102(5), 824-837.

[2]

Moore, B. C. J., & Glasberg, B. R. (2007). Modeling binaural loudness. J. Acoust. Soc. Am., 121(3), 1604-1612.

__init__(kernel_width=18.0, sigma=0.08, learnable=False, dtype=torch.float32)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• kernel_width (float)

• sigma (float)

• learnable (bool)

• dtype (dtype)

forward(specific_loudness)

Apply Gaussian spatial smoothing.

Parameters:

specific_loudness (Tensor) – Specific loudness in sone/ERB. Shape: (batch, 150) or (150,)

Returns:

smoothed_loudness – Spatially smoothed specific loudness. Same shape as input.

Return type:

Tensor

get_parameters()

Get smoothing parameters.

Returns:

Dictionary with kernel_width, sigma, kernel_size, learnable

Return type:

dict

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

BinauralInhibition

class torch_amt.BinauralInhibition(p=1.5978, learnable=False, dtype=torch.float32)

Bases: Module

Cross-ear binaural inhibition using hyperbolic secant (sech) function.

Implements the binaural inhibition mechanism where specific loudness in one ear is suppressed based on the loudness ratio between ears. The inhibition is symmetric and models the competitive interaction between the two ears in loudness perception.

Algorithm Overview

For left and right specific loudness patterns \(N_L(f)\) and \(N_R(f)\), the inhibition factors are computed as:

Left ear inhibition factor:

\[\begin{split}I_L(f) = \\frac{2}{1 + \\text{sech}\\left(\\frac{N_R(f)}{N_L(f)}\\right)^p}\end{split}\]

Right ear inhibition factor:

\[\begin{split}I_R(f) = \\frac{2}{1 + \\text{sech}\\left(\\frac{N_L(f)}{N_R(f)}\\right)^p}\end{split}\]

where the hyperbolic secant function is defined via:

\[\begin{split}\\text{sech}(r) = \\frac{1}{\\cosh(\\ln(r))} = \\frac{2}{e^{\\ln(r)} + e^{-\\ln(r)}} = \\frac{2}{r + 1/r}\end{split}\]

• \(N_L(f)\), \(N_R(f)\) are left and right smoothed specific loudness (sone/ERB)

• \(r = N_R/N_L\) (or \(N_L/N_R\)) is the loudness ratio

• \(p = 1.5978\) is the inhibition exponent (empirically determined)

• \(I_L(f)\), \(I_R(f)\) are inhibition factors ranging from ≈1 to 2

Inhibition factor properties:

• When \(N_L = N_R\) (diotic/equal): \(\\text{sech}(1) = 1\), so \(I = 2/(1+1^p) = 1\) → Maximum inhibition (no binaural advantage)

• When \(N_L \\gg N_R\) or \(N_L \\ll N_R\) (dichotic/unequal): \(\\text{sech}(r) \\to 0\), so \(I \\to 2\) → Minimum inhibition (full binaural advantage)

• The exponent \(p = 1.5978\) controls the steepness of the inhibition function

param p:

Exponent parameter for sech function. Default: 1.5978 Controls the steepness of the inhibition transition

type p:

float

param learnable:

If True, p becomes a learnable parameter. Default: False

type learnable:

bool

param dtype:

Data type for computations. Default: torch.float32

type dtype:

dtype

p

Sech exponent parameter, scalar

Type:

torch.Tensor or nn.Parameter

Input Shape

-----------

left_smoothed

Smoothed left specific loudness, shape (batch, 150) or (150,) Typically from SpatialSmoothing output

Type:

torch.Tensor

right_smoothed

Smoothed right specific loudness, shape (batch, 150) or (150,)

Type:

torch.Tensor

Output Shape

------------

inhib_left

Left inhibition factors, same shape as input

Type:

torch.Tensor

inhib_right

Right inhibition factors, same shape as input

Type:

torch.Tensor

Notes

• Inhibition factors range from ≈1 (strong inhibition, diotic) to 2 (no inhibition, dichotic)

• The inhibition models the competitive interaction between ears in binaural loudness

• Small epsilon (1e-13) is added to prevent division by zero

• The sech function via cosh(ln(x)) is numerically stable for positive ratios

• After computing inhibition factors, they divide the specific loudness: \(N_{inhib,L}(f) = N_L(f) / I_L(f)\)

• Total binaural loudness is then the sum across frequency and ears

See also

SpatialSmoothing

Gaussian smoothing applied before inhibition

Moore2016BinauralLoudness

Complete pipeline including this module

Examples

Basic usage with default Moore2016 parameters:

>>> import torch
>>> from torch_amt.common.loudness import BinauralInhibition
>>>
>>> # Create module
>>> inhibition = BinauralInhibition(p=1.5978)
>>> params = inhibition.get_parameters()
>>> print(f"Inhibition exponent p: {params['p']:.4f}")
Inhibition exponent p: 1.5978

Compute inhibition for smoothed specific loudness:

>>> # Simulate smoothed specific loudness (from SpatialSmoothing)
>>> batch = 2
>>> N_left = torch.randn(batch, 150).abs() * 5
>>> N_right = torch.randn(batch, 150).abs() * 5
>>> I_left, I_right = inhibition(N_left, N_right)
>>> print(f"Inhibition shapes: {I_left.shape}, {I_right.shape}")
Inhibition shapes: torch.Size([2, 150]), torch.Size([2, 150])
>>> print(f"Left inhibition range: [{I_left.min():.3f}, {I_left.max():.3f}]")
Left inhibition range: [1.023, 1.987]
>>> print(f"Right inhibition range: [{I_right.min():.3f}, {I_right.max():.3f}]")
Right inhibition range: [1.015, 1.993]

Diotic case (equal left/right):

>>> N_diotic = torch.ones(1, 150) * 10.0  # Equal loudness
>>> I_L_diotic, I_R_diotic = inhibition(N_diotic, N_diotic)
>>> print(f"Diotic inhibition (should be ~1): {I_L_diotic[0, 0]:.4f}, {I_R_diotic[0, 0]:.4f}")
Diotic inhibition (should be ~1): 1.0000, 1.0000

Dichotic case (unequal left/right):

>>> N_L_dichotic = torch.ones(1, 150) * 100.0  # Much louder in left
>>> N_R_dichotic = torch.ones(1, 150) * 1.0    # Quiet in right
>>> I_L_dichotic, I_R_dichotic = inhibition(N_L_dichotic, N_R_dichotic)
>>> print(f"Dichotic inhibition (should approach 2): Left={I_L_dichotic[0, 0]:.4f}, Right={I_R_dichotic[0, 0]:.4f}")
Dichotic inhibition (should approach 2): Left=1.9856, Right=1.9856

1D input (single loudness patterns):

>>> N_L_1d = torch.randn(150).abs() * 3
>>> N_R_1d = torch.randn(150).abs() * 3
>>> I_L_1d, I_R_1d = inhibition(N_L_1d, N_R_1d)
>>> print(f"1D output shapes: {I_L_1d.shape}, {I_R_1d.shape}")
1D output shapes: torch.Size([150]), torch.Size([150])

Apply inhibition to specific loudness:

>>> N_L_inhib = N_left / I_left  # Inhibited left loudness
>>> N_R_inhib = N_right / I_right  # Inhibited right loudness
>>> print(f"After inhibition: Left sum={N_L_inhib.sum():.2f}, Right sum={N_R_inhib.sum():.2f}")
After inhibition: Left sum=523.45, Right sum=487.32

Learnable exponent for model adaptation:

>>> inhibition_learn = BinauralInhibition(p=1.5978, learnable=True)
>>> optimizer = torch.optim.Adam(inhibition_learn.parameters(), lr=1e-3)

References

[1]

Moore, B. C. J., Glasberg, B. R., & Schlittenlacher, J. (2016). A model of binaural loudness perception based on the outputs of an auditory periphery model. Acta Acustica united with Acustica, 102(5), 824-837.

[2]

Moore, B. C. J., & Glasberg, B. R. (2007). Modeling binaural loudness. J. Acoust. Soc. Am., 121(3), 1604-1612.

__init__(p=1.5978, learnable=False, dtype=torch.float32)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• p (float)

• learnable (bool)

• dtype (dtype)

forward(left_smoothed, right_smoothed)

Compute binaural inhibition factors.

Parameters:

• left_smoothed (Tensor) – Smoothed left specific loudness. Shape: (batch, 150) or (150,)

• right_smoothed (Tensor) – Smoothed right specific loudness. Shape: (batch, 150) or (150,)

Return type:

Tuple[Tensor, Tensor]

Returns:

• inhib_left (torch.Tensor) – Left inhibition factors. Same shape as input.

• inhib_right (torch.Tensor) – Right inhibition factors. Same shape as input.

get_parameters()

Get inhibition parameters.

Returns:

Dictionary with p parameter

Return type:

dict

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

Moore2016BinauralLoudness

class torch_amt.Moore2016BinauralLoudness(kernel_width=18.0, sigma=0.08, p=1.5978, learnable=False, dtype=torch.float32)

Bases: Module

Complete binaural loudness computation for Moore et al. (2016) model.

Combines spatial smoothing, binaural inhibition, and loudness integration to compute total binaural loudness from left and right specific loudness patterns. Implements the complete loudness stage following Moore, Glasberg & Schlittenlacher (2016) with ANSI S3.4-2007 normalization.

Algorithm Overview

The binaural loudness computation follows a 5-step pipeline:

Step 1: Spatial Smoothing (Gaussian convolution)

\[ \begin{align}\begin{aligned}\begin{split}N_{smooth,L}(f) = \\sum_g W(g) \\cdot N_L(f+g)\end{split}\\\begin{split}N_{smooth,R}(f) = \\sum_g W(g) \\cdot N_R(f+g)\end{split}\end{aligned}\end{align} \]

where \(W(g) = \\exp[-(\\sigma g)^2]\) normalized, \(\\sigma=0.08\), \(g \\in [-18, +18]\) ERB

Step 2: Binaural Inhibition (sech function)

\[ \begin{align}\begin{aligned}\begin{split}I_L(f) = \\frac{2}{1 + \\text{sech}\\left(\\frac{N_{smooth,R}(f)}{N_{smooth,L}(f)}\\right)^p}\end{split}\\\begin{split}I_R(f) = \\frac{2}{1 + \\text{sech}\\left(\\frac{N_{smooth,L}(f)}{N_{smooth,R}(f)}\\right)^p}\end{split}\end{aligned}\end{align} \]

where \(p = 1.5978\), \(\\text{sech}(r) = 1/\\cosh(\\ln(r))\)

Step 3: Apply Inhibition

\[ \begin{align}\begin{aligned}\begin{split}N_{inhib,L}(f) = \\frac{N_L(f)}{I_L(f)}\end{split}\\\begin{split}N_{inhib,R}(f) = \\frac{N_R(f)}{I_R(f)}\end{split}\end{aligned}\end{align} \]

Step 4: Frequency Integration (sum across ERB channels)

\[ \begin{align}\begin{aligned}\begin{split}L_L = \\frac{1}{4} \\sum_f N_{inhib,L}(f)\end{split}\\\begin{split}L_R = \\frac{1}{4} \\sum_f N_{inhib,R}(f)\end{split}\end{aligned}\end{align} \]

Division by 4 follows ANSI S3.4-2007 normalization convention

Step 5: Total Binaural Loudness

\[L_{total} = L_L + L_R\]

param kernel_width:

Spatial smoothing kernel half-width in ERB units. Default: 18.0

type kernel_width:

float

param sigma:

Gaussian standard deviation for spatial smoothing. Default: 0.08

type sigma:

float

param p:

Inhibition exponent parameter for sech function. Default: 1.5978

type p:

float

param learnable:

If True, smoothing and inhibition parameters become learnable. Default: False

type learnable:

bool

param dtype:

Data type for computations. Default: torch.float32

type dtype:

dtype

spatial_smoothing

Gaussian spatial smoothing module with kernel_width and sigma

Type:

SpatialSmoothing

inhibition

Cross-ear inhibition module with exponent p

Type:

BinauralInhibition

Input Shape

-----------

specific_loud_left

Left ear specific loudness in sone/ERB, shape (batch, 150) or (150,) Typically from Moore2016SpecificLoudness output

Type:

torch.Tensor

specific_loud_right

Right ear specific loudness in sone/ERB, shape (batch, 150) or (150,)

Type:

torch.Tensor

Output Shape

------------

loudness

Total binaural loudness in sone, shape (batch,) or scalar

Type:

torch.Tensor

loudness_left

Left ear loudness contribution in sone, shape (batch,) or scalar

Type:

torch.Tensor

loudness_right

Right ear loudness contribution in sone, shape (batch,) or scalar

Type:

torch.Tensor

Notes

• The pipeline models the complete binaural loudness perception mechanism

• Spatial smoothing (Step 1) models the spread of excitation across frequency

• Inhibition (Steps 2-3) models the competitive interaction between ears: * Diotic signals (L=R): Maximum inhibition, I≈1, minimal binaural advantage * Dichotic signals (L≠R): Minimal inhibition, I→2, maximum binaural advantage

• Division by 4 (Step 4) is an ANSI S3.4-2007 calibration constant

• Total loudness is the sum (not average) of left and right contributions

• For monaural stimuli (one ear silent), inhibition is minimal and loudness ≈ monaural loudness

• For identical binaural stimuli, loudness ≈ monaural loudness (inhibition cancels binaural summation)

• For uncorrelated binaural stimuli, loudness ≈ 2x monaural loudness

See also

SpatialSmoothing

Gaussian smoothing module (Step 1)

BinauralInhibition

Binaural inhibition module (Steps 2-3)

Moore2016SpecificLoudness

Computes specific loudness before this stage

Moore2016AGC

Automatic gain control for Moore2016 model

Moore2016TemporalIntegration

Temporal integration after this stage

Examples

Basic usage with default Moore2016 parameters:

>>> import torch
>>> from torch_amt.common.loudness import Moore2016BinauralLoudness
>>>
>>> # Create module
>>> binaural = Moore2016BinauralLoudness()
>>> params = binaural.get_parameters()
>>> print(f"Spatial: {params['spatial_smoothing']}")
Spatial: {'kernel_width': 18.0, 'sigma': 0.08, 'kernel_size': 145}
>>> print(f"Inhibition: {params['inhibition']}")
Inhibition: {'p': 1.5978}

Compute binaural loudness from specific loudness:

>>> # Simulate specific loudness (from Moore2016SpecificLoudness)
>>> batch = 3
>>> N_left = torch.randn(batch, 150).abs() * 8
>>> N_right = torch.randn(batch, 150).abs() * 8
>>> L_total, L_left, L_right = binaural(N_left, N_right)
>>> print(f"Total loudness shape: {L_total.shape}, values: {L_total.tolist()}")
Total loudness shape: torch.Size([3]), values: [18.45, 21.32, 19.67]
>>> print(f"Left contribution: {L_left.tolist()}")
Left contribution: [9.23, 10.56, 9.81]
>>> print(f"Right contribution: {L_right.tolist()}")
Right contribution: [9.22, 10.76, 9.86]

Diotic case (identical left/right):

>>> N_diotic = torch.ones(1, 150) * 10.0  # 10 sone/ERB in all channels
>>> L_diotic, L_L_diotic, L_R_diotic = binaural(N_diotic, N_diotic)
>>> print(f"Diotic: Total={L_diotic.item():.2f}, Left={L_L_diotic.item():.2f}, Right={L_R_diotic.item():.2f}")
Diotic: Total=750.00, Left=375.00, Right=375.00
>>> print(f"Inhibition effect: Total loudness ≈ monaural loudness (150*10/4 = 375 per ear)")

Monaural case (one ear silent):

>>> N_monaural_L = torch.ones(1, 150) * 10.0
>>> N_monaural_R = torch.zeros(1, 150)  # Right ear silent
>>> L_monaural, L_L_mon, L_R_mon = binaural(N_monaural_L, N_monaural_R)
>>> print(f"Monaural: Total={L_monaural.item():.2f}, Left={L_L_mon.item():.2f}, Right={L_R_mon.item():.2f}")
Monaural: Total=745.23, Left=745.23, Right=0.00
>>> print(f"Minimal inhibition: Left loudness ≈ 2x (150*10/4) due to I→2")

1D input (single loudness patterns):

>>> N_L_1d = torch.randn(150).abs() * 5
>>> N_R_1d = torch.randn(150).abs() * 5
>>> L_1d, L_L_1d, L_R_1d = binaural(N_L_1d, N_R_1d)
>>> print(f"1D output shapes: {L_1d.shape}, scalars: {L_1d.item():.2f} sone")
1D output shapes: torch.Size([]), scalars: 12.34 sone

Learnable parameters for model adaptation:

>>> binaural_learn = Moore2016BinauralLoudness(learnable=True)
>>> # Can train spatial_smoothing.sigma and inhibition.p
>>> optimizer = torch.optim.Adam(binaural_learn.parameters(), lr=1e-3)

Integration into Moore2016 pipeline:

>>> # Typical usage in Moore2016 model:
>>> # 1. Gammatone filterbank → excitation
>>> # 2. Moore2016SpecificLoudness(excitation_L/R) → N_L, N_R
>>> # 3. Moore2016BinauralLoudness(N_L, N_R) → L_total, L_left, L_right
>>> # 4. Moore2016TemporalIntegration(L_total) → STL, LTL

References

[1]

Moore, B. C. J., Glasberg, B. R., & Schlittenlacher, J. (2016). A model of binaural loudness perception based on the outputs of an auditory periphery model. Acta Acustica united with Acustica, 102(5), 824-837.

[2]

Moore, B. C. J., & Glasberg, B. R. (2007). Modeling binaural loudness. J. Acoust. Soc. Am., 121(3), 1604-1612.

[3]

ANSI S3.4-2007. Procedure for the Computation of Loudness of Steady Sounds. American National Standards Institute.

__init__(kernel_width=18.0, sigma=0.08, p=1.5978, learnable=False, dtype=torch.float32)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• kernel_width (float)

• sigma (float)

• p (float)

• learnable (bool)

• dtype (dtype)

forward(specific_loud_left, specific_loud_right)

Compute binaural loudness with spatial smoothing and inhibition.

Parameters:

• specific_loud_left (Tensor) – Left specific loudness in sone/ERB. Shape: (batch, 150) or (150,)

• specific_loud_right (Tensor) – Right specific loudness in sone/ERB. Shape: (batch, 150) or (150,)

Return type:

Tuple[Tensor, Tensor, Tensor]

Returns:

• loudness (torch.Tensor) – Total binaural loudness in sone. Shape: (batch,) or scalar

• loudness_left (torch.Tensor) – Left loudness contribution in sone. Shape: (batch,) or scalar

• loudness_right (torch.Tensor) – Right loudness contribution in sone. Shape: (batch,) or scalar

get_parameters()

Get all binaural processing parameters.

Returns:

Dictionary with spatial_smoothing and inhibition parameters

Return type:

dict

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

Temporal Integration

LoudnessIntegration

class torch_amt.LoudnessIntegration(fs=32000, learnable=False)

Bases: Module

Loudness integration for Glasberg & Moore (2002) model.

Implements two-stage loudness integration: 1. Spatial integration: Sum specific loudness across ERB frequency channels

to obtain Short-Term Loudness (STL)

2. Temporal integration: Apply asymmetric attack/release IIR filter to obtain Long-Term Loudness (LTL)

Algorithm Overview

Stage 1: Spatial Integration

Sum specific loudness across ERB frequency channels:

\[\begin{split}\\text{STL}(t) = \\sum_{f} N(f,t)\end{split}\]

where:

• \(N(f,t)\) is the specific loudness at ERB frequency \(f\) and time \(t\) (sone/ERB)

• \(\\text{STL}(t)\) is the Short-Term Loudness (sone)

• Sum is computed over all ERB channels (e.g., 150 channels for 50-15000 Hz @ 0.25 ERB step)

Stage 2: Temporal Integration (Asymmetric IIR Filter)

Apply first-order IIR lowpass filter with asymmetric time constants:

\[\begin{split}\\text{LTL}[n] = (1 - \\alpha[n]) \\cdot \\text{STL}[n] + \\alpha[n] \\cdot \\text{LTL}[n-1]\end{split}\]

where the coefficient \(\\alpha[n]\) depends on whether the signal is increasing or decreasing:

\[\begin{split}\\alpha[n] = \\begin{cases} \\exp(-\\Delta t / \\tau_{attack}) & \\text{if } \\text{STL}[n] > \\text{LTL}[n-1] \\text{ (increasing)} \\\\ \\exp(-\\Delta t / \\tau_{release}) & \\text{if } \\text{STL}[n] \\leq \\text{LTL}[n-1] \\text{ (decreasing)} \\end{cases}\end{split}\]

• \(\\tau_{attack} = 0.05\) s (50 ms, fast response to increases)

• \(\\tau_{release} = 0.20\) s (200 ms, slow response to decreases)

• \(\\Delta t = 1 / f_{frame}\) is the frame period (inverse of frame rate)

• \(\\text{LTL}[n]\) is the Long-Term Loudness (sone)

This asymmetric integration models the auditory system’s fast adaptation to loudness increases and slow adaptation to decreases, preventing abrupt loudness drops in time-varying signals.

type fs:

int, optional

param fs:

Sampling rate in Hz. Default: 32000 Used to estimate frame rate if not provided to forward()

type fs:

int, optional

type learnable:

bool, optional

param learnable:

If True, time constants τ_attack and τ_release become learnable parameters. Default: False

type learnable:

bool, optional

fs

Sampling rate in Hz

Type:

int

tau_attack

Attack time constant in seconds, scalar. Fixed at 0.05s (50 ms)

Type:

torch.Tensor or nn.Parameter

tau_release

Release time constant in seconds, scalar. Fixed at 0.20s (200 ms)

Type:

torch.Tensor or nn.Parameter

ltl_state

Current LTL state for temporal integration, shape (batch,) Initialized to zeros, updated during forward pass

Type:

torch.Tensor or None

Input Shape

-----------

specific_loudness

Specific loudness in sone/ERB, shape (batch, n_frames, n_erb_bands) Typically from SpecificLoudness output

Type:

torch.Tensor

Output Shape

------------

ltl

Long-Term Loudness in sone, shape (batch, n_frames)

Type:

torch.Tensor

stl

Short-Term Loudness in sone, shape (batch, n_frames)

Type:

torch.Tensor (optional, if return_stl=True)

Notes

• The two-stage integration separates spatial (frequency) and temporal (time) processing

• Spatial integration (Stage 1) is a simple sum across ERB channels

• Temporal integration (Stage 2) uses an asymmetric IIR filter: * Fast attack (50 ms): Tracks sudden loudness increases quickly * Slow release (200 ms): Prevents rapid loudness drops, models auditory persistence

• Frame rate is typically ~62.5 Hz (fs/512 for hop_length=512 at 32kHz sampling)

• The LTL state is maintained across forward() calls for streaming operation

• Use reset_state() to clear the temporal integration state between non-continuous signals

• Time constants can be made learnable for model adaptation (learnable=True)

See also

SpecificLoudness

Computes specific loudness before integration

Moore2016TemporalIntegration

Temporal integration for Moore2016 model (different time constants)

Examples

Basic usage with Glasberg2002 model:

>>> import torch
>>> from torch_amt.common.loudness import LoudnessIntegration
>>>
>>> # Create module
>>> integration = LoudnessIntegration(fs=32000)
>>> tau_attack, tau_release = integration.get_time_constants()
>>> print(f"Time constants: attack={tau_attack*1000:.0f}ms, release={tau_release*1000:.0f}ms")
Time constants: attack=50ms, release=200ms

Integrate specific loudness to loudness:

>>> # Simulate specific loudness (from SpecificLoudness)
>>> batch, n_frames, n_erb = 2, 200, 150
>>> N_specific = torch.randn(batch, n_frames, n_erb).abs() * 5  # sone/ERB
>>>
>>> # Compute LTL (default output)
>>> ltl = integration(N_specific)
>>> print(f"LTL shape: {ltl.shape}, range: [{ltl.min():.2f}, {ltl.max():.2f}] sone")
LTL shape: torch.Size([2, 200]), range: [0.00, 450.23] sone

Return both STL and LTL:

>>> ltl, stl = integration(N_specific, return_stl=True)
>>> print(f"STL shape: {stl.shape}, range: [{stl.min():.2f}, {stl.max():.2f}] sone")
STL shape: torch.Size([2, 200]), range: [0.00, 523.45] sone
>>> print(f"LTL smooths STL: STL_max={stl.max():.2f}, LTL_max={ltl.max():.2f}")
LTL smooths STL: STL_max=523.45, LTL_max=450.23

Asymmetric temporal integration (attack vs release):

>>> # Create impulse: sudden increase then decrease
>>> impulse = torch.zeros(1, 100, 150)
>>> impulse[:, 20:30, :] = 10.0  # 10-frame pulse
>>>
>>> integration.reset_state()  # Clear state
>>> ltl_impulse, stl_impulse = integration(impulse, return_stl=True)
>>>
>>> # STL shows immediate jump at frame 20
>>> print(f"STL at frame 19-21: {stl_impulse[0, 19:22].tolist()}")
STL at frame 19-21: [0.0, 1500.0, 1500.0]
>>>
>>> # LTL rises fast (attack) but falls slowly (release)
>>> print(f"LTL at frame 19-21: {ltl_impulse[0, 19:22]:.2f} (fast rise)")
LTL at frame 19-21: [0.0, 350.45, 750.23] (fast rise)
>>> print(f"LTL at frame 29-32: {ltl_impulse[0, 29:33]:.2f} (slow fall)")
LTL at frame 29-32: [1450.12, 1398.34, 1350.67, 1305.23] (slow fall)

Reset state between non-continuous signals:

>>> # First signal
>>> N1 = torch.randn(1, 50, 150).abs() * 5
>>> ltl1 = integration(N1)
>>> print(f"LTL1 final state: {integration.ltl_state[0]:.2f}")
LTL1 final state: 235.67
>>>
>>> # Reset before processing second signal
>>> integration.reset_state()
>>> N2 = torch.randn(1, 50, 150).abs() * 5
>>> ltl2 = integration(N2)
>>> print(f"LTL2 starts from: {ltl2[0, 0]:.2f} (should be small)")
LTL2 starts from: 8.45 (should be small)

Learnable time constants for model adaptation:

>>> integration_learn = LoudnessIntegration(fs=32000, learnable=True)
>>> optimizer = torch.optim.Adam(integration_learn.parameters(), lr=1e-3)
>>> # Can train τ_attack and τ_release with backpropagation

References

[1]

Glasberg, B. R., & Moore, B. C. (2002). A model of loudness applicable to time-varying sounds. Journal of the Audio Engineering Society, 50(5), 331-342.

[2]

Moore, B. C. J., Glasberg, B. R., & Baer, T. (1997). A Model for the Prediction of Thresholds, Loudness, and Partial Loudness. J. Audio Eng. Soc, 45(4), 224-240.

__init__(fs=32000, learnable=False)

Initialize loudness integration.

Parameters:

• fs (int, optional) – Sampling rate in Hz. Default: 32000

• learnable (bool, optional) – Whether time constants are learnable parameters. Default: False

reset_state()

Reset temporal integration state.

forward(specific_loudness, return_stl=False)

Integrate specific loudness to obtain loudness.

Parameters:

• specific_loudness (torch.Tensor) – Specific loudness, shape (batch, time, n_erb_bands), in sone/ERB

• return_stl (bool, optional) – Whether to return short-term loudness (STL) in addition to LTL. Default: False

Returns:

If return_stl=False: Long-term loudness, shape (batch, time), in sone. If return_stl=True: Tuple of (ltl, stl) both with shape (batch, time), in sone.

Return type:

torch.Tensor or Tuple[torch.Tensor, torch.Tensor]

get_time_constants()

Get temporal integration time constants.

Returns:

(tau_attack, tau_release) in seconds

Return type:

tuple

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

Moore2016AGC

class torch_amt.Moore2016AGC(attack_alpha=0.045, release_alpha=0.033, learnable=False, dtype=torch.float32)

Bases: Module

Automatic Gain Control (AGC) for temporal smoothing in Moore et al. (2016) model.

Implements asymmetric first-order IIR filtering with different attack and release coefficients, following the AGC mechanism in Moore et al. (2016) and ANSI S3.4-2007. This module provides temporal smoothing that models the auditory system’s integration time with different responses to signal increases (attack) and decreases (release).

Algorithm Overview

The AGC filter is a frame-by-frame first-order IIR lowpass filter with asymmetric time constants:

\[\begin{split}y[n] = \\begin{cases} \\alpha_{attack} \\cdot x[n] + (1 - \\alpha_{attack}) \\cdot y[n-1] & \\text{if } x[n] > y[n-1] \\text{ (increasing)} \\\\ \\alpha_{release} \\cdot x[n] + (1 - \\alpha_{release}) \\cdot y[n-1] & \\text{if } x[n] \\leq y[n-1] \\text{ (decreasing)} \\end{cases}\end{split}\]

where:

• \(x[n]\) is the input signal at frame \(n\)

• \(y[n]\) is the filtered output at frame \(n\)

• \(\\alpha_{attack}\) controls attack speed (higher = faster response to increases)

• \(\\alpha_{release}\) controls release speed (higher = faster response to decreases)

• Coefficients are in range [0, 1]: 0 = no filtering (instant), 1 = infinite hold

Two-Stage Usage in Moore2016:

1. Short-Term AGC: Instantaneous → Short-Term Specific Loudness

   • Operates on (n_frames, 150) specific loudness patterns

   • \(\\alpha_{attack} = 0.045\) (fast attack, ~22 frames to reach 63%)

   • \(\\alpha_{release} = 0.033\) (fast release, ~30 frames to reach 63%)

   • Smooths specific loudness independently per ERB channel

2. Long-Term AGC: Short-Term → Long-Term Loudness

   • Operates on (n_frames,) scalar loudness values

   • \(\\alpha_{attack} = 0.01\) (slow attack, ~100 frames to reach 63%)

   • \(\\alpha_{release} = 0.00133\) (very slow release, ~752 frames to reach 63%)

   • Models long-term loudness adaptation

The effective time constant \(\\tau\) relates to \(\\alpha\) via:

\[\begin{split}\\tau = -\\frac{\\Delta t}{\\ln(1 - \\alpha)}\end{split}\]

where \(\\Delta t\) is the frame period.

param attack_alpha:

Attack coefficient (0 to 1). Default: 0.045 (short-term AGC) Higher values = faster response to signal increases

type attack_alpha:

float

param release_alpha:

Release coefficient (0 to 1). Default: 0.033 (short-term AGC) Higher values = faster response to signal decreases

type release_alpha:

float

param learnable:

If True, α_attack and α_release become learnable parameters. Default: False

type learnable:

bool

param dtype:

Data type for computations. Default: torch.float32

type dtype:

dtype

attack_alpha

Attack coefficient, scalar

Type:

torch.Tensor or nn.Parameter

release_alpha

Release coefficient, scalar

Type:

torch.Tensor or nn.Parameter

Input Shape

-----------

x

Input signal, shape (n_frames, n_channels) or (n_frames,) Can be specific loudness (150 channels) or scalar loudness

Type:

torch.Tensor

state

Initial filter state, shape (n_channels,) or scalar If None, starts from zeros

Type:

torch.Tensor, optional

Output Shape

------------

output

Filtered signal, same shape as input

Type:

torch.Tensor

Notes

• The AGC is applied frame-by-frame to maintain causality

• Attack/release selection is element-wise (per channel or per batch element)

• State is updated at each frame and carried to the next

• Alpha coefficients interpretation: * α = 0: No filtering (output = input, instantaneous) * α = 0.5: Medium smoothing (half-weight to previous state) * α = 1: Infinite hold (output = previous state, no update)

• Short-term AGC (α_attack=0.045, α_release=0.033): * Fast attack: ~22 frames to 63% (1-e^-1), ~100 frames to 99% * Fast release: ~30 frames to 63%, ~138 frames to 99%

• Long-term AGC (α_attack=0.01, α_release=0.00133): * Slow attack: ~100 frames to 63%, ~460 frames to 99% * Very slow release: ~752 frames to 63%, ~3460 frames to 99%

• For frame rate ≈ 62.5 fps (hop=512 @ 32kHz), 100 frames ≈ 1.6 seconds

See also

Moore2016TemporalIntegration

Complete two-stage AGC pipeline

LoudnessIntegration

Glasberg2002 temporal integration (different approach)

Examples

Short-term AGC for specific loudness smoothing:

>>> import torch
>>> from torch_amt.common.loudness import Moore2016AGC
>>>
>>> # Create short-term AGC
>>> stl_agc = Moore2016AGC(attack_alpha=0.045, release_alpha=0.033)
>>> params = stl_agc.get_parameters()
>>> print(f"Short-term AGC: attack={params['attack_alpha']:.3f}, release={params['release_alpha']:.3f}")
Short-term AGC: attack=0.045, release=0.033
>>>
>>> # Apply to instantaneous specific loudness (150 ERB channels)
>>> n_frames, n_channels = 100, 150
>>> inst_spec_loud = torch.randn(n_frames, n_channels).abs() * 8
>>> st_spec_loud = stl_agc(inst_spec_loud)
>>> print(f"Input shape: {inst_spec_loud.shape}, Output shape: {st_spec_loud.shape}")
Input shape: torch.Size([100, 150]), Output shape: torch.Size([100, 150])

Long-term AGC for scalar loudness smoothing:

>>> # Create long-term AGC
>>> ltl_agc = Moore2016AGC(attack_alpha=0.01, release_alpha=0.00133)
>>> params_lt = ltl_agc.get_parameters()
>>> print(f"Long-term AGC: attack={params_lt['attack_alpha']:.5f}, release={params_lt['release_alpha']:.5f}")
Long-term AGC: attack=0.01000, release=0.00133
>>>
>>> # Apply to short-term loudness (scalar time series)
>>> st_loud = torch.randn(100).abs() * 50
>>> lt_loud = ltl_agc(st_loud)
>>> print(f"Input shape: {st_loud.shape}, Output shape: {lt_loud.shape}")
Input shape: torch.Size([100]), Output shape: torch.Size([100])

Asymmetric response to impulse:

>>> # Create impulse: sudden increase then return to baseline
>>> impulse = torch.zeros(200)
>>> impulse[50:60] = 100.0  # 10-frame pulse at high level
>>>
>>> agc = Moore2016AGC(attack_alpha=0.045, release_alpha=0.033)
>>> filtered = agc(impulse)
>>>
>>> # Fast attack (frame 50-60)
>>> print(f"Attack: frame 50={filtered[50]:.2f}, frame 55={filtered[55]:.2f}, frame 59={filtered[59]:.2f}")
Attack: frame 50=4.50, frame 55=23.45, frame 59=67.89
>>>
>>> # Slow release (frame 60-100)
>>> print(f"Release: frame 60={filtered[60]:.2f}, frame 70={filtered[70]:.2f}, frame 80={filtered[80]:.2f}")
Release: frame 60=67.45, frame 70=54.32, frame 80=43.21

With initial state (streaming):

>>> # First chunk
>>> chunk1 = torch.randn(50, 150).abs() * 5
>>> out1 = stl_agc(chunk1)
>>> final_state = out1[-1, :]  # Last frame as state
>>>
>>> # Second chunk with initial state
>>> chunk2 = torch.randn(50, 150).abs() * 5
>>> out2 = stl_agc(chunk2, state=final_state)
>>> print(f"Continuous processing: out2 starts from state={final_state.mean():.2f}")
Continuous processing: out2 starts from state=3.45

Learnable coefficients for model adaptation:

>>> agc_learn = Moore2016AGC(attack_alpha=0.045, release_alpha=0.033, learnable=True)
>>> optimizer = torch.optim.Adam(agc_learn.parameters(), lr=1e-3)
>>> # Can train α_attack and α_release with backpropagation

References

[1]

Moore, B. C. J., Glasberg, B. R., & Schlittenlacher, J. (2016). A model of binaural loudness perception based on the outputs of an auditory periphery model. Acta Acustica united with Acustica, 102(5), 824-837.

[2]

ANSI S3.4-2007. Procedure for the Computation of Loudness of Steady Sounds. American National Standards Institute.

__init__(attack_alpha=0.045, release_alpha=0.033, learnable=False, dtype=torch.float32)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• attack_alpha (float)

• release_alpha (float)

• learnable (bool)

• dtype (dtype)

forward(x, state=None)

Apply AGC filtering frame-by-frame.

Parameters:

• x (Tensor) – Input signal. Shape: (n_frames, n_channels) or (n_frames,)

• state (Tensor) – Initial state. If None, starts from zeros. Shape: (n_channels,) or scalar (matching x channels)

Returns:

output – Filtered signal. Same shape as input.

Return type:

Tensor

get_parameters()

Get AGC parameters.

Returns:

Dictionary with attack_alpha and release_alpha

Return type:

dict

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

Moore2016TemporalIntegration

class torch_amt.Moore2016TemporalIntegration(stl_attack=0.045, stl_release=0.033, ltl_attack=0.01, ltl_release=0.00133, learnable=False, dtype=torch.float32)

Bases: Module

Complete two-stage temporal integration for Moore et al. (2016) binaural loudness model.

Implements the complete temporal processing pipeline with two AGC stages followed by frequency integration, transforming instantaneous specific loudness into long-term loudness. This models the auditory system’s temporal integration at multiple time scales.

Algorithm Overview

The temporal integration consists of three sequential stages:

Stage 1: Short-Term AGC (per-channel specific loudness smoothing)

Apply asymmetric IIR filter to each of 150 ERB channels independently:

\[\begin{split}N_{ST}[n,f] = \\begin{cases} \\alpha_{ST,attack} \\cdot N_{inst}[n,f] + (1-\\alpha_{ST,attack}) \\cdot N_{ST}[n-1,f] & \\text{if increasing} \\\\ \\alpha_{ST,release} \\cdot N_{inst}[n,f] + (1-\\alpha_{ST,release}) \\cdot N_{ST}[n-1,f] & \\text{if decreasing} \\end{cases}\end{split}\]

• \(N_{inst}[n,f]\) is instantaneous specific loudness at frame \(n\), ERB channel \(f\)

• \(N_{ST}[n,f]\) is short-term specific loudness (sone/ERB)

• \(\\alpha_{ST,attack} = 0.045\) (fast attack, ~22 frames to 63%)

• \(\\alpha_{ST,release} = 0.033\) (fast release, ~30 frames to 63%)

Stage 2: Frequency Integration (sum across ERB channels)

\[\begin{split}L_{ST}[n] = \\frac{1}{4} \\sum_{f=1}^{150} N_{ST}[n,f]\end{split}\]

• \(L_{ST}[n]\) is short-term loudness (sone)

• Division by 4 follows ANSI S3.4-2007 normalization

Stage 3: Long-Term AGC (scalar loudness smoothing)

Apply asymmetric IIR filter to scalar loudness:

\[\begin{split}L_{LT}[n] = \\begin{cases} \\alpha_{LT,attack} \\cdot L_{ST}[n] + (1-\\alpha_{LT,attack}) \\cdot L_{LT}[n-1] & \\text{if increasing} \\\\ \\alpha_{LT,release} \\cdot L_{ST}[n] + (1-\\alpha_{LT,release}) \\cdot L_{LT}[n-1] & \\text{if decreasing} \\end{cases}\end{split}\]

• \(L_{LT}[n]\) is long-term loudness (sone)

• \(\\alpha_{LT,attack} = 0.01\) (slow attack, ~100 frames to 63%)

• \(\\alpha_{LT,release} = 0.00133\) (very slow release, ~752 frames to 63%)

The two-stage AGC provides multi-scale temporal integration: - Short-term (Stage 1): Fast adaptation (~0.3-0.5s at 62.5 fps) - Long-term (Stage 3): Slow adaptation (~1.6-12s at 62.5 fps)

param stl_attack:

Short-term attack coefficient (Stage 1). Default: 0.045

type stl_attack:

float

param stl_release:

Short-term release coefficient (Stage 1). Default: 0.033

type stl_release:

float

param ltl_attack:

Long-term attack coefficient (Stage 3). Default: 0.01

type ltl_attack:

float

param ltl_release:

Long-term release coefficient (Stage 3). Default: 0.00133

type ltl_release:

float

param learnable:

If True, all four α coefficients become learnable parameters. Default: False

type learnable:

bool

param dtype:

Data type for computations. Default: torch.float32

type dtype:

dtype

stl_agc

Short-term AGC module for Stage 1 (150-channel filtering)

Type:

Moore2016AGC

ltl_agc

Long-term AGC module for Stage 3 (scalar filtering)

Type:

Moore2016AGC

Input Shape

-----------

inst_spec_loud

Instantaneous specific loudness in sone/ERB, shape (n_frames, 150) Typically from Moore2016SpecificLoudness output

Type:

torch.Tensor

Output Shape

------------

ltl

Long-term loudness in sone, shape (n_frames,)

Type:

torch.Tensor

stl_spec

Short-term specific loudness in sone/ERB, shape (n_frames, 150)

Type:

torch.Tensor (optional, if return_intermediate=True)

stl

Short-term loudness in sone, shape (n_frames,)

Type:

torch.Tensor (optional, if return_intermediate=True)

Notes

• The two-stage AGC provides hierarchical temporal integration: * Stage 1 (Short-term): Fast per-channel smoothing models early auditory adaptation * Stage 3 (Long-term): Slow scalar smoothing models perceptual loudness integration

• Stage 2 (frequency integration) converts 150-channel representation to scalar loudness

• All processing is causal (frame-by-frame) for real-time compatibility

• Frame rate is typically ~62.5 Hz (fs/512 for hop_length=512 at 32kHz sampling)

• Alpha coefficient time scales: * Short-term: Attack ~22 frames (0.35s), Release ~30 frames (0.48s) * Long-term: Attack ~100 frames (1.6s), Release ~752 frames (12.0s)

• The long-term AGC prevents abrupt loudness changes and models perceptual persistence

• Division by 4 in Stage 2 is ANSI S3.4-2007 calibration constant

See also

Moore2016AGC

Single-stage AGC module (used internally)

LoudnessIntegration

Glasberg2002 temporal integration (different approach)

Moore2016BinauralLoudness

Binaural loudness computation before temporal integration

Moore2016SpecificLoudness

Specific loudness computation before this stage

Examples

Basic usage with default Moore2016 parameters:

>>> import torch
>>> from torch_amt.common.loudness import Moore2016TemporalIntegration
>>>
>>> # Create module
>>> temporal = Moore2016TemporalIntegration()
>>> params = temporal.get_parameters()
>>> print(f"STL AGC: {params['stl']}")
STL AGC: {'attack_alpha': 0.045, 'release_alpha': 0.033}
>>> print(f"LTL AGC: {params['ltl']}")
LTL AGC: {'attack_alpha': 0.01, 'release_alpha': 0.00133}

Apply two-stage temporal integration:

>>> # Simulate instantaneous specific loudness (from Moore2016SpecificLoudness)
>>> n_frames = 200
>>> inst_spec_loud = torch.randn(n_frames, 150).abs() * 8  # 200 frames, 150 channels
>>>
>>> # Get long-term loudness (default output)
>>> ltl = temporal(inst_spec_loud)
>>> print(f"LTL shape: {ltl.shape}, range: [{ltl.min():.2f}, {ltl.max():.2f}] sone")
LTL shape: torch.Size([200]), range: [0.00, 95.34] sone

Return intermediate outputs:

>>> ltl, stl_spec, stl = temporal(inst_spec_loud, return_intermediate=True)
>>> print(f"ST specific loudness: {stl_spec.shape}")
ST specific loudness: torch.Size([200, 150])
>>> print(f"ST loudness: {stl.shape}, range: [{stl.min():.2f}, {stl.max():.2f}] sone")
ST loudness: torch.Size([200]), range: [0.00, 145.67] sone
>>> print(f"LT loudness: {ltl.shape}, range: [{ltl.min():.2f}, {ltl.max():.2f}] sone")
LT loudness: torch.Size([200]), range: [0.00, 95.34] sone
>>> print(f"LTL smooths STL: STL_max={stl.max():.2f}, LTL_max={ltl.max():.2f}")
LTL smooths STL: STL_max=145.67, LTL_max=95.34

Multi-scale temporal smoothing (impulse response):

>>> # Create impulse: sudden onset then decay
>>> impulse = torch.zeros(300, 150)
>>> impulse[50:70, :] = 20.0  # 20-frame pulse
>>>
>>> ltl_imp, stl_spec_imp, stl_imp = temporal(impulse, return_intermediate=True)
>>>
>>> # Short-term rises and falls quickly
>>> print(f"STL at frames 50-55: {stl_imp[50:56].tolist()}")
STL at frames 50-55: [0.0, 135.0, 450.0, 650.0, 720.0, 750.0]
>>> print(f"STL at frames 70-75: {stl_imp[70:76].tolist()}")
STL at frames 70-75: [750.0, 725.3, 701.5, 678.9, 657.2, 636.3]
>>>
>>> # Long-term rises slowly and persists longer
>>> print(f"LTL at frames 50-55: {ltl_imp[50:56]:.2f} (slower rise)")
LTL at frames 50-55: [0.0, 1.35, 5.85, 12.35, 20.12, 28.95] (slower rise)
>>> print(f"LTL at frames 70-80: {ltl_imp[70:81]:.2f} (much slower decay)")
LTL at frames 70-80: [65.43, 65.31, 65.20, 65.08, 64.97, ...] (much slower decay)

Integration into Moore2016 pipeline:

>>> # Complete Moore2016 pipeline:
>>> # 1. Gammatone filterbank → cochleagram
>>> # 2. Moore2016SpecificLoudness(cochleagram_L/R) → N_inst_L, N_inst_R
>>> # 3. Moore2016BinauralLoudness(N_inst_L, N_inst_R) → L_binaural
>>> # 4. Moore2016TemporalIntegration(L_binaural) → L_LT (long-term loudness)

Learnable parameters for model adaptation:

>>> temporal_learn = Moore2016TemporalIntegration(learnable=True)
>>> # Can train all 4 alpha coefficients with backpropagation
>>> optimizer = torch.optim.Adam(temporal_learn.parameters(), lr=1e-4)

Custom time constants (e.g., faster long-term):

>>> temporal_custom = Moore2016TemporalIntegration(
...     stl_attack=0.045, stl_release=0.033,  # Standard short-term
...     ltl_attack=0.05, ltl_release=0.01      # Faster long-term
... )
>>> params_custom = temporal_custom.get_parameters()
>>> print(f"Custom LTL: {params_custom['ltl']}")
Custom LTL: {'attack_alpha': 0.05, 'release_alpha': 0.01}

References

[1]

Moore, B. C. J., Glasberg, B. R., & Schlittenlacher, J. (2016). A model of binaural loudness perception based on the outputs of an auditory periphery model. Acta Acustica united with Acustica, 102(5), 824-837.

[2]

ANSI S3.4-2007. Procedure for the Computation of Loudness of Steady Sounds. American National Standards Institute.

__init__(stl_attack=0.045, stl_release=0.033, ltl_attack=0.01, ltl_release=0.00133, learnable=False, dtype=torch.float32)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• stl_attack (float)

• stl_release (float)

• ltl_attack (float)

• ltl_release (float)

• learnable (bool)

• dtype (dtype)

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:

None

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:

Iterator[tuple[str, Tensor]]

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:

Iterator[tuple[str, Parameter]]

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:

None

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:

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

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

3. 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:

None

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:

None

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:

None

share_memory()

See torch.Tensor.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

forward(inst_spec_loud, return_intermediate=False)

Apply two-stage temporal integration.

Parameters:

• inst_spec_loud (Tensor) – Instantaneous specific loudness. Shape: (n_frames, 150)

• return_intermediate (bool) – If True, return intermediate STL values. Default: False

Return type:

Tensor

Returns:

• ltl (torch.Tensor) – Long-term loudness in sone. Shape: (n_frames,)

• stl_spec (torch.Tensor (optional)) – Short-term specific loudness. Shape: (n_frames, 150)

• stl (torch.Tensor (optional)) – Short-term loudness. Shape: (n_frames,)

get_parameters()

Get all temporal integration parameters.

Returns:

Dictionary with STL and LTL parameters

Return type:

dict

extra_repr()

Extra representation string for module printing.

Returns:

String containing key module parameters.

Return type:

str