Complete Auditory Models

torch_amt provides 6 complete end-to-end auditory models ready for research and applications.

Model Overview

Model	Year	Key Features	Primary Applications
Dau1997	1997	Adaptation loops, modulation	AM detection, temporal processing
Glasberg2002	2002	Specific loudness, integration	Loudness perception, hearing aids
Moore2016	2016	Binaural processing, spatial	Binaural loudness, spatial hearing
King2019	2019	Broken-stick compression, FM/AM	FM masking, modulation interactions
Osses2021	2021	Extended temporal integration	Speech perception, temporal resolution
Paulick2024	2024	Physiological IHC, CASP	Physiological modeling, cochlear implants

Dau1997

class torch_amt.Dau1997(fs, flow=80.0, fhigh=8000.0, learnable=False, return_stages=False, dtype=torch.float32, filterbank_kwargs=None, ihc_kwargs=None, adaptation_kwargs=None, modulation_kwargs=None)

Bases: Module

Dau et al. (1997) auditory model for amplitude modulation processing.

Implements the complete auditory processing pipeline from Dau, Kollmeier, and Kohlrausch (1997) for modeling detection and masking with amplitude-modulated signals. The model simulates peripheral auditory processing from cochlear filtering through neural adaptation to modulation detection.

This implementation is based on the MATLAB Auditory Modeling Toolbox (AMT) dau1997 function and provides a differentiable, GPU-accelerated version suitable for neural network training and optimization.

Algorithm Overview

The model implements a 4-stage auditory processing pipeline:

Stage 1: Gammatone Filterbank

Decomposes the input signal into frequency channels using 4th-order gammatone filters with 1-ERB spacing:

\[x_f(t) = \text{Gammatone}_{f_c}(x(t))\]

where \(f_c\) are the center frequencies from flow to fhigh.

Stage 2: Inner Hair Cell (IHC) Envelope Extraction

Extracts the envelope via half-wave rectification followed by 1000 Hz lowpass filtering (Dau 1996 method):

\[e_f(t) = \text{Lowpass}_{1000\,\text{Hz}}(\max(0, x_f(t)))\]

Stage 3: Adaptation Loops

Models neural adaptation using 5 cascaded feedback loops with time constants \(\tau = [5, 50, 129, 253, 500]\) ms and overshoot limit \(L=10\):

\[a_f(t) = \text{AdaptLoop}(e_f(t), \tau, L)\]

Stage 4: Modulation Filterbank

Analyzes temporal modulations using bandpass filters with \(Q=2\) up to 150 Hz modulation frequency:

\[m_{f,j}(t) = \text{ModFilter}_{f_{mod,j}}(a_f(t))\]

where \(j\) indexes modulation channels (frequency-dependent).

param fs:

Sampling rate in Hz. Typical values: 16000, 32000, 44100. Higher sampling rates provide better temporal resolution but increase computational cost.

type fs:

float

param flow:

Lower frequency bound for gammatone filterbank in Hz. Default: 80 Hz. Lower values extend low-frequency coverage but increase filter count. Typical range: [50, 200] Hz.

type flow:

float

param fhigh:

Upper frequency bound for gammatone filterbank in Hz. Default: 8000 Hz. Higher values extend high-frequency coverage but increase filter count. Must be less than Nyquist frequency (fs/2). Typical range: [4000, 16000] Hz.

type fhigh:

float

param learnable:

If True, all model stages (filterbank, IHC, adaptation, modulation) become trainable with gradient-based optimization. Default: False (fixed parameters). When True, enables end-to-end model training for task-specific optimization.

type learnable:

bool

param return_stages:

If True, returns intermediate processing stages along with final output. Default: False (only final modulation representation). Useful for visualization, analysis, and multi-stage training.

type return_stages:

bool

param dtype:

Data type for computations and parameters. Default: torch.float32. Use torch.float64 for higher precision if needed (increases memory/computation).

type dtype:

dtype

param **filterbank_kwargs:

Additional keyword arguments passed to GammatoneFilterbank. Common options:

• n (int): Filter order. Default: 4.

• erb_scale (float): ERB scale factor. Default: 1.0.

• Other parameters accepted by GammatoneFilterbank.

type **filterbank_kwargs:

Dict[str, Any]

param **ihc_kwargs:

Additional keyword arguments passed to IHCEnvelope. Common options:

• method (str): IHC method. Default: ‘dau1996’.

• Other parameters accepted by IHCEnvelope.

type **ihc_kwargs:

Dict[str, Any]

param **adaptation_kwargs:

Additional keyword arguments passed to AdaptLoop. Common options:

• tau (list or torch.Tensor): Time constants in seconds. Default: None (uses [0.005, 0.050, 0.129, 0.253, 0.500]).

• limit (float): Overshoot limit factor. Default: 10.0.

• minspl (float): Minimum SPL in dB. Default: 0.0.

• Other parameters accepted by AdaptLoop.

type **adaptation_kwargs:

Dict[str, Any]

param **modulation_kwargs:

Additional keyword arguments passed to ModulationFilterbank. Common options:

• Q (float): Quality factor for modulation filters. Default: 2.0.

• max_mfc (float): Maximum modulation frequency in Hz. Default: 150.0.

• filter_type (str): Filter type (‘efilt’ or ‘butterworth’). Default: ‘efilt’.

• Other parameters accepted by ModulationFilterbank.

type **modulation_kwargs:

Dict[str, Any]

fs

Sampling rate in Hz.

Type:

float

flow

Lower frequency bound in Hz.

Type:

float

fhigh

Upper frequency bound in Hz.

Type:

float

learnable

Whether model parameters are trainable.

Type:

bool

return_stages

Whether to return intermediate processing stages.

Type:

bool

dtype

Data type for computations.

Type:

torch.dtype

filterbank

Stage 1: Gammatone auditory filterbank module.

Type:

GammatoneFilterbank

ihc

Stage 2: Inner hair cell envelope extraction module.

Type:

IHCEnvelope

adaptation

Stage 3: Neural adaptation loops module.

Type:

AdaptLoop

modulation

Stage 4: Temporal modulation filterbank module.

Type:

ModulationFilterbank

fc

Center frequencies of auditory filterbank, shape (num_channels,) in Hz. Computed with 1-ERB spacing between flow and fhigh.

Type:

torch.Tensor

num_channels

Number of auditory frequency channels in the filterbank. Typically 20-40 channels for default frequency range (80-8000 Hz).

Type:

int

Input Shape

-----------

x

Audio signal with shape:

• \((B, T)\) - Batch of audio samples

• \((T,)\) - Single audio sample (mono)

where:

• \(B\) = batch size

• \(T\) = time samples

Type:

torch.Tensor

Output Shape

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

When ``return_stages=False`` (default)

List[torch.Tensor]

List of length num_channels (one per frequency channel). Each element has shape \((B, M_i, T)\) where:

• \(M_i\) = number of modulation channels for frequency channel \(i\)

• Varies per channel: low frequencies have fewer modulation channels

When ``return_stages=True``

Tuple[List[torch.Tensor], Dict[str, torch.Tensor]]

• First element: internal representation (as above)

• Second element: dict with keys [‘filterbank’, ‘ihc’, ‘adaptation’] containing intermediate outputs, each shape \((B, F, T)\) where \(F\) = num_channels

Examples

Basic usage:

>>> import torch
>>> from torch_amt.models import Dau1997
>>>
>>> # Create model
>>> model = Dau1997(fs=44100, flow=80, fhigh=8000)
>>> print(f"Frequency channels: {model.num_channels}")
Frequency channels: 31
>>>
>>> # Process audio (1 second, 44.1 kHz)
>>> audio = torch.randn(2, 44100)  # 2 batches
>>> output = model(audio)
>>> print(f"Output: {len(output)} frequency channels")
Output: 31 frequency channels
>>> print(f"First channel shape: {output[0].shape}")
First channel shape: torch.Size([2, 13, 44100])
>>> print(f"Last channel shape: {output[-1].shape}")
Last channel shape: torch.Size([2, 8, 44100])

With intermediate stages for visualization:

>>> model_debug = Dau1997(fs=44100, return_stages=True)
>>> output, stages = model_debug(audio)
>>>
>>> print(f"Available stages: {list(stages.keys())}")
Available stages: ['filterbank', 'ihc', 'adaptation']
>>> print(f"Filterbank output: {stages['filterbank'].shape}")
Filterbank output: torch.Size([2, 31, 44100])
>>> print(f"IHC output: {stages['ihc'].shape}")
IHC output: torch.Size([2, 31, 44100])
>>> print(f"Adaptation output: {stages['adaptation'].shape}")
Adaptation output: torch.Size([2, 31, 44100])

Single channel input (mono):

>>> audio_mono = torch.randn(44100)  # No batch dimension
>>> output_mono = model(audio_mono)
>>> print(f"Output shape (mono): {output_mono[0].shape}")
Output shape (mono): torch.Size([13, 44100])

Learnable model for optimization:

>>> model_learnable = Dau1997(fs=44100, learnable=True)
>>> n_params = sum(p.numel() for p in model_learnable.parameters())
>>> print(f"Trainable parameters: {n_params}")
Trainable parameters: 15234
>>>
>>> # Example training loop
>>> optimizer = torch.optim.Adam(model_learnable.parameters(), lr=1e-3)
>>> # ... training code ...

Different frequency ranges:

>>> # Extended low frequency
>>> model_low = Dau1997(fs=44100, flow=50, fhigh=8000)
>>> print(f"Channels with flow=50Hz: {model_low.num_channels}")
Channels with flow=50Hz: 35
>>>
>>> # Extended high frequency
>>> model_high = Dau1997(fs=44100, flow=80, fhigh=16000)
>>> print(f"Channels with fhigh=16kHz: {model_high.num_channels}")
Channels with fhigh=16kHz: 39

Custom submodule parameters via kwargs:

>>> # Custom filterbank: 6th-order filters instead of default 4th-order
>>> model_custom_fb = Dau1997(
...     fs=44100,
...     filterbank_kwargs={'n': 6}
... )
>>>
>>> # Custom adaptation: different overshoot limit and time constants
>>> model_custom_adapt = Dau1997(
...     fs=44100,
...     adaptation_kwargs={
...         'limit': 5.0,  # Lower overshoot limit
...         'tau': [0.01, 0.1, 0.2, 0.4, 0.8]  # Custom time constants
...     }
... )
>>>
>>> # Custom modulation: higher Q and different max frequency
>>> model_custom_mod = Dau1997(
...     fs=44100,
...     modulation_kwargs={
...         'Q': 1.0,  # Lower Q (broader filters)
...         'max_mfc': 200.0,  # Extend to 200 Hz
...         'filter_type': 'butterworth'  # Use Butterworth instead of efilt
...     }
... )
>>>
>>> # Combine multiple custom parameters
>>> model_fully_custom = Dau1997(
...     fs=44100,
...     flow=50,
...     fhigh=12000,
...     filterbank_kwargs={'n': 5, 'erb_scale': 1.2},
...     ihc_kwargs={'method': 'breebaart2001'},
...     adaptation_kwargs={'limit': 8.0, 'minspl': -10.0},
...     modulation_kwargs={'Q': 1.5, 'max_mfc': 180.0}
... )
>>> print(f"Fully custom model channels: {model_fully_custom.num_channels}")
Fully custom model channels: 43

Notes

Model Configuration:

The Dau1997 model uses specific configurations for each processing stage:

• Gammatone filterbank: 4th-order filters with 1-ERB spacing

• IHC envelope: Dau1996 method (half-wave rectification + 1 kHz lowpass)

• Adaptation: 5 loops with \(\\tau = [5, 50, 129, 253, 500]\) ms, overshoot limit \(L=10\), \(\\text{minspl}=0\) dB

• Modulation filterbank: \(Q=2\), max modulation frequency 150 Hz

Customizing Submodule Parameters:

All submodules can be customized through dedicated kwargs dictionaries:

• Use filterbank_kwargs to pass parameters to GammatoneFilterbank

• Use ihc_kwargs to pass parameters to IHCEnvelope

• Use adaptation_kwargs to pass parameters to AdaptLoop

• Use modulation_kwargs to pass parameters to ModulationFilterbank

The learnable and dtype parameters are always centralized and applied to all submodules automatically. Custom parameters override defaults while maintaining the Dau1997 model structure.

Example: To use 6th-order gammatone filters with a different adaptation limit:

model = Dau1997(
    fs=44100,
    filterbank_kwargs={'n': 6},
    adaptation_kwargs={'limit': 5.0}
)

Computational Complexity:

Processing time scales approximately as:

\[\begin{split}T_{compute} \\propto B \\cdot F \\cdot T \\cdot (1 + M)\end{split}\]

where \(F\) = num_channels (20-40), \(M\) = avg modulation channels (~10). For 1 second at 44.1 kHz: ~0.1-0.5 seconds on CPU, ~0.01-0.05 seconds on GPU.

Memory Requirements:

Peak memory scales with output representation size:

\[\begin{split}Memory \\approx B \\cdot F \\cdot M \\cdot T \\cdot 4\\,\\text{bytes}\end{split}\]

For batch=8, 1 second @ 44.1 kHz: ~50-100 MB.

Applications:

The model internal representation can be used for:

• Amplitude modulation detection and discrimination

• Monaural masking predictions

• Temporal modulation transfer function (TMTF) modeling

• Feature extraction for machine learning tasks

• Psychoacoustic model evaluation and fitting

See also

GammatoneFilterbank

Stage 1 - Cochlear filtering

IHCEnvelope

Stage 2 - Inner hair cell transduction

AdaptLoop

Stage 3 - Neural adaptation

ModulationFilterbank

Stage 4 - Modulation detection

References

[1]

T. Dau, B. Kollmeier, and A. Kohlrausch, “Modeling auditory processing of amplitude modulation. I. Detection and masking with narrow-band carriers,” J. Acoust. Soc. Am., vol. 102, no. 5, pp. 2892-2905, Nov. 1997.

[2]

T. Dau, B. Kollmeier, and A. Kohlrausch, “Modeling auditory processing of amplitude modulation. II. Spectral and temporal integration,” J. Acoust. Soc. Am., vol. 102, no. 5, pp. 2906-2919, Nov. 1997.

[3]

T. Dau, D. Püschel, and A. Kohlrausch, “A quantitative model of the ‘effective’ signal processing in the auditory system. I. Model structure,” J. Acoust. Soc. Am., vol. 99, no. 6, pp. 3615-3622, June 1996.

[4]

R. D. Patterson, I. Nimmo-Smith, J. Holdsworth, and P. Rice, “An efficient auditory filterbank based on the gammatone function,” APU Report 2341, MRC Applied Psychology Unit, Cambridge, UK, 1988.

[5]

P. Majdak, C. Hollomey, and R. Baumgartner, “AMT 1.x: A toolbox for reproducible research in auditory modeling,” Acta Acust., vol. 6, p. 19, 2022.

__init__(fs, flow=80.0, fhigh=8000.0, learnable=False, return_stages=False, dtype=torch.float32, filterbank_kwargs=None, ihc_kwargs=None, adaptation_kwargs=None, modulation_kwargs=None)

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

Parameters:

• fs (float)

• flow (float)

• fhigh (float)

• learnable (bool)

• return_stages (bool)

• dtype (dtype)

• filterbank_kwargs (Dict[str, Any])

• ihc_kwargs (Dict[str, Any])

• adaptation_kwargs (Dict[str, Any])

• modulation_kwargs (Dict[str, Any])

forward(x)

Process audio through the Dau1997 model.

Parameters:

x (Tensor) – Input audio signal. Shape: (B, T), (C, T), or (T,).

Returns:

If return_stages=False:

List of tensors (one per frequency channel), each shape (B, M, T).

If return_stages=True:

Tuple of (output, stages) where stages is a dict with intermediate outputs.

Return type:

List[Tensor] | tuple[List[Tensor], Dict[str, Any]]

extra_repr()

Extra representation for printing.

Returns:

String representation of module parameters.

Return type:

str

distribute_gradients()

Distribute gradients to grouped filter coefficients in FastModulationFilterbank.

Call this method after loss.backward() to ensure all filter coefficients in the modulation filterbank receive gradient updates. This is necessary when using FastModulationFilterbank with learnable=True, as filters are grouped for efficiency and gradients need to be shared across group members.

Notes

This method should be called in the training loop:

>>> model = Dau1997(fs=44100, learnable=True)
>>> output = model(input_signal)
>>> loss = criterion(output, target)
>>> loss.backward()
>>> model.distribute_gradients()  # ← Important for FastModulationFilterbank!
>>> optimizer.step()

If the modulation filterbank doesn’t have a distribute_gradients method (e.g., using standard ModulationFilterbank), this is a no-op.

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

Glasberg2002

class torch_amt.Glasberg2002(fs=32000, learnable=False, return_stages=False, multi_fft_kwargs=None, erb_kwargs=None, excitation_kwargs=None, specific_loudness_kwargs=None, loudness_integration_kwargs=None)

Bases: Module

Glasberg & Moore (2002) model for time-varying loudness perception.

Implements the complete loudness computation pipeline from Glasberg & Moore (2002), providing perceptual loudness measures in sone from audio waveforms. The model accounts for frequency-dependent hearing sensitivity (ISO 226), masking effects via asymmetric excitation spreading, and temporal integration with attack/release dynamics.

This implementation is based on the MATLAB Auditory Modeling Toolbox (AMT) glasberg2002 function and provides a differentiable, GPU-accelerated version suitable for neural network training and loudness-based optimization.

Algorithm Overview

The model implements a 5-stage loudness processing pipeline:

Stage 1: Multi-Resolution FFT

Performs time-frequency analysis with multiple FFT sizes to balance temporal and frequency resolution across the audible spectrum:

\[\begin{split}X(t, f) = \\text{FFT}_{N(f)}(x(t))\end{split}\]

where \(N(f)\) is frequency-dependent FFT size (larger for low frequencies). Outputs power spectral density (PSD) in \(\\text{Pa}^2/\\text{Hz}\).

Stage 2: ERB Integration

Maps PSD to perceptual ERB frequency scale with 1/4 ERB resolution:

\[\begin{split}E_{\\text{ERB}}(t, f_{\\text{ERB}}) = \\int P(t, f) \\cdot W_{\\text{ERB}}(f, f_{\\text{ERB}}) df\end{split}\]

where \(W_{\\text{ERB}}\) is the ERB weighting function. Output in dB SPL.

Stage 3: Excitation Pattern

Models asymmetric frequency spreading with level-dependent slopes:

\[\begin{split}E_{\\text{spread}}(t, f) = \\sum_g E_{\\text{ERB}}(t, f+g) \\cdot S(g, E)\end{split}\]

where \(S(g, E)\) is the spreading function (steeper upward, shallower downward).

Stage 4: Specific Loudness

Applies 3-regime loudness transformation (Moore & Glasberg 1997):

\[\begin{split}N(t, f) = \\begin{cases} 0 & E < E_{\\text{thrq}} \\\\ C \\cdot (E - E_{\\text{thrq}}) & E_{\\text{thrq}} < E < E_0 \\\\ C \\cdot E_0^{1-\\alpha} (E - E_{\\text{thrq}})^{\\alpha} & E > E_0 \\end{cases}\end{split}\]

with \(C=0.047\), \(\\alpha=0.2\), \(E_0=10\) dB above threshold.

Stage 5: Loudness Integration

Spatial integration (sum across ERB channels) followed by temporal integration with asymmetric attack/release filter:

\[\begin{split}\\text{STL}(t) = \\sum_f N(t, f), \\quad \\text{LTL}[n] = (1-\\alpha[n])\\text{STL}[n] + \\alpha[n]\\text{LTL}[n-1]\end{split}\]

where \(\\alpha = \\exp(-\\Delta t / \\tau)\) with \(\\tau_{\\text{attack}}=50\) ms, \(\\tau_{\\text{release}}=200\) ms.

param fs:

Sampling rate in Hz. Default: 32000 Hz. Higher sampling rates improve temporal resolution but increase computational cost. Typical values: 16000, 32000, 44100, 48000 Hz.

type fs:

int

param learnable:

If True, all model stages become trainable with gradient-based optimization. Default: False (fixed parameters). When True, enables end-to-end model training for task-specific optimization.

type learnable:

bool

param return_stages:

If True, returns intermediate processing stages along with final output. Default: False (only final long-term loudness). Useful for visualization, analysis, and multi-stage training.

type return_stages:

bool

param **multi_fft_kwargs:

Additional keyword arguments passed to MultiResolutionFFT. Common options:

• hop_length (int): Hop size for STFT in samples.

• n_ffts (list): FFT sizes for multi-resolution analysis.

• Other parameters accepted by MultiResolutionFFT.

type **multi_fft_kwargs:

Dict[str, Any]

param **erb_kwargs:

Additional keyword arguments passed to ERBIntegration. Common options:

• f_min (float): Minimum frequency in Hz. Default: 50.0.

• f_max (float): Maximum frequency in Hz. Default: 15000.0.

• erb_step (float): ERB frequency step. Default: 0.25.

• bandwidth_scale (float): Bandwidth scaling factor. Default: 1.0.

• Other parameters accepted by ERBIntegration.

type **erb_kwargs:

Dict[str, Any]

param **excitation_kwargs:

Additional keyword arguments passed to ExcitationPattern. Common options:

• upper_slope_base (float): Base upper spreading slope. Default: 27.0 dB/ERB.

• lower_slope_base (float): Base lower spreading slope. Default: 27.0 dB/ERB.

• upper_slope_per_db (float): Upper slope level dependency. Default: 0.0.

• lower_slope_per_db (float): Lower slope level dependency. Default: -0.4 dB/ERB per dB.

• Other parameters accepted by ExcitationPattern.

type **excitation_kwargs:

Dict[str, Any]

param **specific_loudness_kwargs:

Additional keyword arguments passed to SpecificLoudness. Common options:

• f_min (float): Minimum ERB frequency. Default: 50.0 Hz.

• f_max (float): Maximum ERB frequency. Default: 15000.0 Hz.

• erb_step (float): ERB step. Default: 0.25.

• Other parameters accepted by SpecificLoudness.

type **specific_loudness_kwargs:

Dict[str, Any]

param **loudness_integration_kwargs:

Additional keyword arguments passed to LoudnessIntegration. Common options:

• tau_attack (float): Attack time constant in seconds. Default: 0.05 (50 ms).

• tau_release (float): Release time constant in seconds. Default: 0.20 (200 ms).

• Other parameters accepted by LoudnessIntegration.

type **loudness_integration_kwargs:

Dict[str, Any]

fs

Sampling rate in Hz.

Type:

int

learnable

Whether model parameters are trainable.

Type:

bool

return_stages

Whether to return intermediate processing stages.

Type:

bool

multi_fft

Stage 1: Multi-resolution time-frequency analysis module.

Type:

MultiResolutionFFT

erb_integration

Stage 2: ERB frequency scale integration module.

Type:

ERBIntegration

excitation_pattern

Stage 3: Excitation pattern spreading module.

Type:

ExcitationPattern

specific_loudness

Stage 4: Specific loudness transformation module.

Type:

SpecificLoudness

loudness_integration

Stage 5: Spatial and temporal loudness integration module.

Type:

LoudnessIntegration

Input Shape

-----------

audio

Audio signal with shape:

• \((B, T)\) - Batch of audio samples

• \((T,)\) - Single audio sample (mono)

where:

• \(B\) = batch size

• \(T\) = time samples

Type:

torch.Tensor

Output Shape

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

When ``return_stages=False`` (default)

torch.Tensor

Long-term loudness in sone, shape \((B, F)\) where:

• \(F\) = number of time frames (depends on hop_length)

When ``return_stages=True``

Tuple[torch.Tensor, Dict[str, torch.Tensor]]

• First element: long-term loudness (as above)

• Second element: dict with keys:

  • 'stl': Short-term loudness, shape \((B, F)\) in sone

  • 'specific_loudness': Specific loudness, shape \((B, F, N_{\\text{ERB}})\) in sone/ERB

  • 'excitation': Excitation pattern, shape \((B, F, N_{\\text{ERB}})\) in dB SPL

  • 'erb_excitation': ERB-integrated excitation, shape \((B, F, N_{\\text{ERB}})\) in dB SPL

  • 'psd': Power spectral density, shape \((B, F, N_{\\text{freq}})\)

  • 'freqs': Frequency vector for PSD, shape \((N_{\\text{freq}},)\) in Hz

Examples

Basic usage:

>>> import torch
>>> from torch_amt.models import Glasberg2002
>>>
>>> # Create model
>>> model = Glasberg2002(fs=32000)
>>> n_erb = model.erb_integration.n_erb_bands
>>> print(f"ERB channels: {n_erb}")
ERB channels: 150
>>>
>>> # Process 1 second of audio
>>> audio = torch.randn(2, 32000)  # 2 batches
>>> ltl = model(audio)
>>> print(f"LTL shape: {ltl.shape}, range: [{ltl.min():.2f}, {ltl.max():.2f}] sone")
LTL shape: torch.Size([2, 62]), range: [0.23, 45.67] sone

With intermediate stages:

>>> model_debug = Glasberg2002(fs=32000, return_stages=True)
>>> ltl, stages = model_debug(audio)
>>>
>>> print(f"Available stages: {list(stages.keys())}")
Available stages: ['stl', 'specific_loudness', 'excitation', 'erb_excitation', 'psd', 'freqs']
>>> print(f"STL shape: {stages['stl'].shape}")
STL shape: torch.Size([2, 62])
>>> print(f"Specific loudness shape: {stages['specific_loudness'].shape}")
Specific loudness shape: torch.Size([2, 62, 150])
>>> print(f"Excitation shape: {stages['excitation'].shape}")
Excitation shape: torch.Size([2, 62, 150])

Single channel input:

>>> audio_mono = torch.randn(32000)  # No batch dimension
>>> ltl_mono = model(audio_mono)
>>> print(f"Output shape (mono): {ltl_mono.shape}")
Output shape (mono): torch.Size([62])

Learnable model for optimization:

>>> model_learnable = Glasberg2002(fs=32000, learnable=True)
>>> n_params = sum(p.numel() for p in model_learnable.parameters())
>>> print(f"Trainable parameters: {n_params}")
Trainable parameters: 8743
>>>
>>> # Example training loop
>>> optimizer = torch.optim.Adam(model_learnable.parameters(), lr=1e-3)
>>> # ... training code ...

Custom submodule parameters:

>>> # Custom ERB frequency range
>>> model_custom_erb = Glasberg2002(
...     fs=44100,
...     erb_kwargs={'f_min': 80.0, 'f_max': 12000.0, 'erb_step': 0.5}
... )
>>> print(f"ERB channels: {model_custom_erb.erb_integration.n_erb_bands}")
ERB channels: 75
>>>
>>> # Custom excitation spreading
>>> model_custom_exc = Glasberg2002(
...     fs=32000,
...     excitation_kwargs={
...         'upper_slope_base': 30.0,  # Steeper upper slope
...         'lower_slope_per_db': -0.5  # More level-dependent lower slope
...     }
... )
>>>
>>> # Custom temporal integration
>>> model_custom_temp = Glasberg2002(
...     fs=32000,
...     loudness_integration_kwargs={
...         'tau_attack': 0.03,  # Faster attack (30 ms)
...         'tau_release': 0.30  # Slower release (300 ms)
...     }
... )

Different sampling rates:

>>> model_44k = Glasberg2002(fs=44100)
>>> audio_44k = torch.randn(2, 44100)  # 1 second @ 44.1 kHz
>>> ltl_44k = model_44k(audio_44k)
>>> print(f"Output frames @ 44.1kHz: {ltl_44k.shape[1]}")
Output frames @ 44.1kHz: 86

Reset temporal state for new signal:

>>> # Process first signal
>>> signal1 = torch.randn(1, 32000)
>>> ltl1 = model(signal1)
>>>
>>> # Reset before processing unrelated second signal
>>> model.reset_state()
>>> signal2 = torch.randn(1, 32000)
>>> ltl2 = model(signal2)

Convert to loudness level (phon):

>>> ltl_sone = model(audio)
>>> ltl_phon = model.compute_loudness_level(ltl_sone)
>>> print(f"Loudness: {ltl_sone.mean():.2f} sone = {ltl_phon.mean():.2f} phon")
Loudness: 12.34 sone = 54.32 phon

Notes

Model Configuration:

The Glasberg2002 model uses specific configurations for each processing stage:

• Multi-resolution FFT: Multiple FFT sizes (frequency-dependent) for balanced resolution

• ERB integration: 1/4 ERB steps from 50 Hz to 15 kHz (150 channels)

• Excitation pattern: Asymmetric spreading (27 dB/ERB base, -0.4 dB/ERB per dB lower slope)

• Specific loudness: 3-regime transformation (\(C=0.047\), \(\\alpha=0.2\), \(E_0=10\) dB)

• Loudness integration: Attack 50 ms, release 200 ms

Customizing Submodule Parameters:

All submodules can be customized through dedicated kwargs dictionaries:

• Use multi_fft_kwargs to pass parameters to MultiResolutionFFT

• Use erb_kwargs to pass parameters to ERBIntegration

• Use excitation_kwargs to pass parameters to ExcitationPattern

• Use specific_loudness_kwargs to pass parameters to SpecificLoudness

• Use loudness_integration_kwargs to pass parameters to LoudnessIntegration

The learnable and dtype parameters are always centralized and applied to all submodules automatically. Custom parameters override defaults while maintaining the Glasberg2002 model structure.

Computational Complexity:

Processing time scales approximately as:

\[\begin{split}T_{\\text{compute}} \\propto B \\cdot F \\cdot N_{\\text{ERB}} \\cdot \\log N_{\\text{FFT}}\end{split}\]

where \(F\) = number of time frames (~60 per second), \(N_{\\text{ERB}}=150\). For 1 second at 32 kHz: ~0.05-0.2 seconds on CPU, ~0.005-0.02 seconds on GPU.

Memory Requirements:

Peak memory scales with intermediate representations:

\[\begin{split}Memory \\approx B \\cdot F \\cdot N_{\\text{ERB}} \\cdot 4\\,\\text{bytes}\end{split}\]

For batch=8, 1 second @ 32 kHz: ~20-40 MB.

Differences from MATLAB AMT:

• This implementation uses PyTorch tensors for GPU acceleration

• Supports batch processing natively

• All stages are differentiable for gradient-based optimization

• Output frames depend on hop_length (not fixed downsampling)

Loudness Units:

• Sone: Perceptual loudness unit. 1 sone = loudness of 1 kHz tone at 40 dB SPL

• Phon: Loudness level unit. Equal to dB SPL at 1 kHz

• Conversion: \(L_{\\text{phon}} = 40 + 10\\log_2(L_{\\text{sone}})\)

Applications:

The model output can be used for:

• Perceptual loudness measurement and normalization

• Audio quality assessment (loudness-based metrics)

• Dynamic range compression/expansion

• Hearing aid fitting and evaluation

• Psychoacoustic model validation

• Feature extraction for machine learning

See also

MultiResolutionFFT

Stage 1 - Time-frequency analysis

ERBIntegration

Stage 2 - ERB frequency scale

ExcitationPattern

Stage 3 - Excitation spreading

SpecificLoudness

Stage 4 - Loudness transformation

LoudnessIntegration

Stage 5 - Spatial and temporal integration

References

[1]

B. R. Glasberg and B. C. J. Moore, “A Model of Loudness Applicable to Time-Varying Sounds,” J. Audio Eng. Soc., vol. 50, no. 5, pp. 331-342, May 2002.

[2]

B. C. J. Moore and B. R. Glasberg, “A Model for the Prediction of Thresholds, Loudness, and Partial Loudness,” J. Audio Eng. Soc., vol. 45, no. 4, pp. 224-240, Apr. 1997.

[3]

B. R. Glasberg and B. C. J. Moore, “Derivation of auditory filter shapes from notched-noise data,” Hear. Res., vol. 47, no. 1-2, pp. 103-138, Aug. 1990.

[4]

B. C. J. Moore and B. R. Glasberg, “Formulae describing frequency selectivity as a function of frequency and level, and their use in calculating excitation patterns,” Hear. Res., vol. 28, no. 2-3, pp. 209-225, 1987.

[5]

ISO 226:2003, “Acoustics - Normal equal-loudness-level contours,” International Organization for Standardization, 2003.

[6]

P. Majdak, C. Hollomey, and R. Baumgartner, “AMT 1.x: A toolbox for reproducible research in auditory modeling,” Acta Acust., vol. 6, p. 19, 2022.

__init__(fs=32000, learnable=False, return_stages=False, multi_fft_kwargs=None, erb_kwargs=None, excitation_kwargs=None, specific_loudness_kwargs=None, loudness_integration_kwargs=None)

Initialize Glasberg & Moore (2002) loudness model.

Parameters:

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

• learnable (bool) – If True, all model parameters become trainable. Default: False.

• return_stages (bool) – If True, return intermediate processing stages. Default: False.

• **kwargs (dict) – Additional keyword arguments for submodules (see class docstring).

• multi_fft_kwargs (Dict[str, Any])

• erb_kwargs (Dict[str, Any])

• excitation_kwargs (Dict[str, Any])

• specific_loudness_kwargs (Dict[str, Any])

• loudness_integration_kwargs (Dict[str, Any])

forward(audio)

Process audio through the Glasberg2002 loudness model.

Parameters:

audio (Tensor) – Input audio signal. Shape: (B, T) or (T,).

Returns:

If return_stages=False:

Long-term loudness in sone, shape (B, F) or (F,).

If return_stages=True:

Tuple of (ltl, stages) where stages is a dict with: - ‘stl’: Short-term loudness (B, F) in sone - ‘specific_loudness’: Specific loudness (B, F, N_ERB) in sone/ERB - ‘excitation’: Excitation pattern (B, F, N_ERB) in dB SPL - ‘erb_excitation’: ERB-integrated excitation (B, F, N_ERB) in dB SPL - ‘psd’: Power spectral density (B, F, N_freq) - ‘freqs’: Frequency vector (N_freq,) in Hz

Return type:

Tensor | Tuple[Tensor, Dict[str, Any]]

reset_state()

Reset temporal integration state for processing discontinuous signals.

Clears the internal state of the temporal integration filter (LoudnessIntegration). Call this method when processing multiple unrelated audio signals sequentially to prevent temporal blending between signals.

Examples

>>> model = Glasberg2002(fs=32000)
>>> signal1 = torch.randn(1, 32000)
>>> ltl1 = model(signal1)
>>>
>>> # Reset before processing new signal
>>> model.reset_state()
>>> signal2 = torch.randn(1, 32000)
>>> ltl2 = model(signal2)  # No temporal carryover from signal1

get_erb_frequencies()

Get ERB channel center frequencies.

Returns the center frequencies (in Hz) of the ERB-spaced frequency channels used throughout the model pipeline. Useful for frequency-domain visualization and analysis.

Returns:

Center frequencies of ERB channels, shape (n_erb_bands,) in Hz. Typically 150 channels from 50 Hz to 15 kHz with 1/4 ERB spacing.

Return type:

Tensor

Examples

>>> model = Glasberg2002(fs=32000)
>>> fc = model.get_erb_frequencies()
>>> print(f"ERB frequencies: {fc.shape}, range: [{fc.min():.1f}, {fc.max():.1f}] Hz")
ERB frequencies: torch.Size([150]), range: [50.0, 14999.2] Hz

get_learnable_parameters()

Get all learnable parameters organized by model component.

Returns a nested dictionary containing the current values of all trainable parameters in each pipeline stage. Only returns parameters when model is initialized with learnable=True.

Returns:

Dictionary with component names as keys and parameter dicts as values:

• 'multi_fft': MultiResolutionFFT parameters (hop_length, n_ffts)

• 'erb_integration': ERBIntegration parameters (fc_erb, bandwidth_scale)

• 'excitation_pattern': ExcitationPattern parameters (slopes)

• 'specific_loudness': SpecificLoudness parameters (C, alpha, E0, thresholds)

• 'loudness_integration': LoudnessIntegration parameters (tau_attack, tau_release)

Empty dict if learnable=False.

Return type:

Dict[str, Any]

Examples

>>> model = Glasberg2002(fs=32000, learnable=True)
>>> params = model.get_learnable_parameters()
>>> print(f"Components: {list(params.keys())}")
Components: ['multi_fft', 'erb_integration', 'excitation_pattern', 'specific_loudness', 'loudness_integration']
>>>
>>> # Access specific parameters
>>> print(f"Attack time: {params['loudness_integration']['tau_attack']:.3f} s")
Attack time: 0.050 s
>>> print(f"Bandwidth scale: {params['erb_integration']['bandwidth_scale']:.2f}")
Bandwidth scale: 1.00

compute_loudness_level(ltl)

Convert loudness from sone to loudness level in phon.

Applies Stevens’ power law to convert perceptual loudness (sone) to loudness level (phon), which is equivalent to dB SPL at 1 kHz.

Parameters:

ltl (Tensor) – Loudness in sone, shape (B, F) or (F,). Values should be non-negative.

Returns:

Loudness level in phon, same shape as input. Formula: \(L_{\text{phon}} = 40 + 10\log_2(L_{\text{sone}})\)

Return type:

Tensor

Notes

Loudness Units:

• Sone: Perceptual loudness. 1 sone = loudness of 1 kHz tone at 40 dB SPL. Doubling sone value = doubling perceived loudness.

• Phon: Loudness level. Equal to dB SPL of equally loud 1 kHz tone. 40 phon = 40 dB SPL @ 1 kHz = 1 sone.

Conversion Examples:

• 1 sone = 40 phon (reference)

• 2 sone = 50 phon (10 dB louder)

• 4 sone = 60 phon (20 dB louder)

• 0.5 sone = 30 phon (10 dB quieter)

Examples

>>> model = Glasberg2002(fs=32000)
>>> audio = torch.randn(2, 32000)
>>> ltl_sone = model(audio)
>>> ltl_phon = model.compute_loudness_level(ltl_sone)
>>>
>>> print(f"Loudness: {ltl_sone.mean():.2f} sone")
Loudness: 12.34 sone
>>> print(f"Loudness level: {ltl_phon.mean():.2f} phon")
Loudness level: 54.32 phon

References

[1]

S. S. Stevens, “Perceived level of noise by Mark VII and decibels (E),” J. Acoust. Soc. Am., vol. 51, no. 2B, pp. 575-601, 1972.

extra_repr()

Extra representation for printing.

Returns:

String representation of 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

Moore2016

class torch_amt.Moore2016(fs=32000, learnable=False, return_stages=False, dtype=torch.float32, outer_middle_ear_kwargs=None, spectrum_kwargs=None, excitation_kwargs=None, specific_loudness_kwargs=None, temporal_integration_kwargs=None, binaural_loudness_kwargs=None, ltl_agc_kwargs=None)

Bases: Module

Moore et al. (2016) model for binaural loudness perception.

Implements the complete binaural loudness processing pipeline from Moore, Glasberg, and Schlittenlacher (2016), providing perceptual loudness measures in sone from stereo audio waveforms. The model extends monaural loudness computation with binaural inhibition effects, accounting for both diotic summation and dichotic interactions between ears.

This implementation is a 1:1 port of the MATLAB Auditory Modeling Toolbox (AMT) moore2016 function and provides a differentiable, GPU-accelerated version suitable for neural network training and binaural loudness optimization.

Algorithm Overview

The model implements an 8-stage binaural loudness processing pipeline:

Stage 1: Outer/Middle Ear Filtering (per ear)

Simulates transmission through outer and middle ear using ANSI S3.4-2007 transfer function:

\[\begin{split}x_L'(t) = h_{\\text{ear}}(t) * x_L(t), \\quad x_R'(t) = h_{\\text{ear}}(t) * x_R(t)\end{split}\]

where \(h_{\\text{ear}}\) is the tfOuterMiddle2007 filter (1025 coefficients).

Stage 2: Multi-Resolution Spectral Analysis (per ear)

Uses 6 time windows (2.048ms to 65.536ms) with frequency-dependent selection:

\[\begin{split}S(t, f) = \\text{FFT}_{N(f)}(w(t) \\cdot x'(t))\end{split}\]

where \(N(f)\) varies from 64 to 2048 samples based on frequency. Output: sparse spectrum with ERB-spaced frequencies.

Stage 3: Excitation Pattern (per ear)

Models basilar membrane response with rounded-exponential (roex) filters on 150 ERB-spaced channels (80 Hz - 15 kHz):

\[\begin{split}E(t, f_{ERB}) = \\sum_k S(t, f_k) \\cdot W_{roex}(f_k, f_{ERB})\end{split}\]

Output in dB SPL.

Stage 4: Specific Loudness (per ear)

Applies instantaneous nonlinear compression via 88-value lookup table:

\[\begin{split}N_{inst}(t, f) = \\text{LUT}(E(t, f))\end{split}\]

Output: instantaneous specific loudness in sone/ERB.

Stage 5: Short-Term Temporal Integration (per ear)

AGC smoothing with asymmetric attack/release dynamics:

\[\begin{split}N_{STL}(t, f) = \\alpha N_{inst}(t, f) + (1-\\alpha) N_{STL}(t-1, f)\end{split}\]

where \(\\alpha = 0.045\) (attack), \(\\alpha = 0.033\) (release).

Stage 6: Binaural Inhibition (frame-by-frame)

Spatial smoothing followed by cross-ear inhibition:

\[\begin{split}N_{smooth}(f) = \\sum_g N_{STL}(f+g) \\cdot G(g, \\sigma=0.08)\end{split}\]

\[\begin{split}N_{inhib,L} = N_{smooth,L} \\cdot \\text{sech}(p \\cdot N_{smooth,R}), \\quad p=1.5978\end{split}\]

Integration yields scalar loudness per ear: \(L_L = \\sum_f N_{inhib,L}(f) / 4\).

Stage 7: Long-Term Temporal Integration (per ear)

AGC on scalar loudness with slower dynamics:

\[\begin{split}L_{LTL}(t) = \\alpha L_{STL}(t) + (1-\\alpha) L_{LTL}(t-1)\end{split}\]

where \(\\alpha = 0.01\) (attack), \(\\alpha = 0.00133\) (release).

Stage 8: Binaural Output

\[\begin{split}L_{STL} = L_{STL,L} + L_{STL,R}, \\quad L_{LTL} = L_{LTL,L} + L_{LTL,R}\end{split}\]

Output: short-term loudness (STL), long-term loudness (LTL), max loudness.

param fs:

Sampling rate in Hz. Must be 32000 Hz. Default: 32000. This requirement is enforced by the outer/middle ear filter design (ANSI S3.4-2007). If you have audio at different sampling rates, resample to 32 kHz before processing.

type fs:

int

param learnable:

If True, all model stages become trainable with gradient-based optimization. Default: False (fixed parameters). When True, enables end-to-end model training for task-specific optimization.

type learnable:

bool

param return_stages:

If True, returns intermediate processing stages along with final output. Default: False (only final STL, LTL, mLoud). Useful for visualization, analysis, and multi-stage training.

type return_stages:

bool

param dtype:

Data type for computations and parameters. Default: torch.float32. Use torch.float64 for higher precision if needed (increases memory/computation).

type dtype:

dtype

param **outer_middle_ear_kwargs:

Additional keyword arguments passed to OuterMiddleEarFilter. Common options:

• compensation_type (str): Filter type. Default: ‘tfOuterMiddle2007’.

• field_type (str): Sound field type. Default: ‘free’.

• Other parameters accepted by OuterMiddleEarFilter.

type **outer_middle_ear_kwargs:

Dict[str, Any]

param **spectrum_kwargs:

Additional keyword arguments passed to Moore2016Spectrum. Common options:

• hop_length (int): Hop size in samples. Default: 512.

• n_windows (int): Number of window sizes. Default: 6.

• Other parameters accepted by Moore2016Spectrum.

type **spectrum_kwargs:

Dict[str, Any]

param **excitation_kwargs:

Additional keyword arguments passed to Moore2016ExcitationPattern. Common options:

• erb_lower (float): Lower ERB limit. Default: 1.75.

• erb_upper (float): Upper ERB limit. Default: 39.0.

• erb_step (float): ERB step size. Default: 0.25.

• Other parameters accepted by Moore2016ExcitationPattern.

type **excitation_kwargs:

Dict[str, Any]

param **specific_loudness_kwargs:

Additional keyword arguments passed to Moore2016SpecificLoudness. Parameters for the 88-value loudness lookup table.

type **specific_loudness_kwargs:

Dict[str, Any]

param **temporal_integration_kwargs:

Additional keyword arguments passed to Moore2016TemporalIntegration. Common options:

• attack_alpha (float): STL attack coefficient. Default: 0.045.

• release_alpha (float): STL release coefficient. Default: 0.033.

• Other parameters accepted by Moore2016TemporalIntegration.

type **temporal_integration_kwargs:

Dict[str, Any]

param **binaural_loudness_kwargs:

Additional keyword arguments passed to Moore2016BinauralLoudness. Common options:

• sigma (float): Spatial smoothing width. Default: 0.08.

• p (float): Inhibition strength parameter. Default: 1.5978.

• Other parameters accepted by Moore2016BinauralLoudness.

type **binaural_loudness_kwargs:

Dict[str, Any]

param **ltl_agc_kwargs:

Additional keyword arguments passed to Moore2016AGC (for LTL). Common options:

• attack_alpha (float): LTL attack coefficient. Default: 0.01.

• release_alpha (float): LTL release coefficient. Default: 0.00133.

• Applied to both left and right LTL AGCs.

• Applied to both left and right LTL AGCs.

type **ltl_agc_kwargs:

Dict[str, Any]

fs

Sampling rate in Hz (fixed at 32000).

Type:

int

learnable

Whether model parameters are trainable.

Type:

bool

return_stages

Whether to return intermediate processing stages.

Type:

bool

dtype

Data type for computations.

Type:

torch.dtype

outer_middle_ear

Stage 1: Outer and middle ear transfer function filter.

Type:

OuterMiddleEarFilter

spectrum

Stage 2: Multi-resolution spectral analysis module.

Type:

Moore2016Spectrum

excitation

Stage 3: Auditory excitation pattern computation module.

Type:

Moore2016ExcitationPattern

specific_loudness

Stage 4: Instantaneous specific loudness module.

Type:

Moore2016SpecificLoudness

temporal_integration

Stage 5: Short-term temporal AGC module.

Type:

Moore2016TemporalIntegration

binaural_loudness

Stage 6: Binaural inhibition and spatial smoothing module.

Type:

Moore2016BinauralLoudness

ltl_agc_left

Stage 7: Long-term AGC for left ear.

Type:

Moore2016AGC

ltl_agc_right

Stage 7: Long-term AGC for right ear.

Type:

Moore2016AGC

Input Shape

-----------

audio

Stereo audio signal with shape:

• \((B, 2, T)\) - Batch of stereo audio samples

• \((2, T)\) - Single stereo sample

where:

• \(B\) = batch size

• \(2\) = channels (left, right)

• \(T\) = time samples at 32 kHz

Important: Audio must be in Pascal (Pa) units. For conversion: \(\\text{Pa} = 2 \\times 10^{-5} \\times 10^{\\text{dB SPL}/20}\).

Type:

torch.Tensor

Output Shape

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

When ``return_stages=False`` (default)

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

Three tensors:

• sLoud: Short-term binaural loudness, shape \((B, F)\) in sone

• lLoud: Long-term binaural loudness, shape \((B, F)\) in sone

• mLoud: Maximum long-term loudness, shape \((B,)\) in sone

where \(F\) = number of time frames (depends on hop_length).

When ``return_stages=True``

Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], Dict[str, torch.Tensor]]

• First element: (sLoud, lLoud, mLoud) as above

• Second element: dict with keys:

  • 'filtered_left': After outer/middle ear, shape \((B, T')\)

  • 'filtered_right': shape \((B, T')\)

  • 'freqs_left': Sparse spectrum frequencies, shape \((B, F, N_{comp})\)

  • 'levels_left': Sparse spectrum levels in dB, shape \((B, F, N_{comp})\)

  • 'freqs_right', 'levels_right': same for right ear

  • 'excitation_left': Excitation pattern, shape \((B, F, 150)\) in dB SPL

  • 'excitation_right': shape \((B, F, 150)\)

  • 'inst_spec_loud_left': Instantaneous specific loudness, shape \((B, F, 150)\) sone/ERB

  • 'inst_spec_loud_right': shape \((B, F, 150)\)

  • 'stl_spec_loud_left': Short-term specific loudness, shape \((B, F, 150)\)

  • 'stl_spec_loud_right': shape \((B, F, 150)\)

  • 'stl_loud_left': Short-term scalar loudness left, shape \((B, F)\) sone

  • 'stl_loud_right': shape \((B, F)\)

  • 'ltl_loud_left': Long-term scalar loudness left, shape \((B, F)\)

  • 'ltl_loud_right': shape \((B, F)\)

Examples

Basic usage:

>>> import torch
>>> import numpy as np
>>> from torch_amt.models import Moore2016
>>>
>>> # Create model
>>> model = Moore2016(fs=32000)
>>>
>>> # Generate stereo tone at 1 kHz, 60 dB SPL
>>> t = np.linspace(0, 1, 32000)
>>> tone = np.sin(2 * np.pi * 1000 * t)
>>> tone_pa = tone * 0.02  # 60 dB SPL
>>> audio_stereo = torch.from_numpy(np.stack([tone_pa, tone_pa])).float().unsqueeze(0)
>>>
>>> # Process
>>> sLoud, lLoud, mLoud = model(audio_stereo)
>>> print(f"STL: {sLoud.mean():.2f} sone")
STL: 6.17 sone

With stages:

>>> model_debug = Moore2016(fs=32000, return_stages=True)
>>> (sLoud, lLoud, mLoud), stages = model_debug(audio_stereo)
>>> print(f"Excitation: {stages['excitation_left'].shape}")
Excitation: torch.Size([1, 62, 150])

Notes

Model Configuration:

• Outer/Middle Ear: ANSI S3.4-2007 tfOuterMiddle2007 transfer function

• Spectrum: 6 windows (64-2048 samples), hop=512

• Excitation: 150 ERB channels (1.75-39.0 ERB, 80 Hz - 15 kHz)

• STL AGC: attack=0.045, release=0.033

• Binaural: Gaussian σ=0.08, inhibition p=1.5978

• LTL AGC: attack=0.01, release=0.00133

Customizing Submodules:

All submodules accept kwargs through dedicated dictionaries. The learnable and dtype parameters are centralized. See Parameters section for details.

Sampling Rate:

Must be 32000 Hz due to outer/middle ear filter design. Resample if needed:

import torchaudio
audio = torchaudio.functional.resample(audio, fs_orig, 32000)

Binaural Effects:

1. Summation: Diotic stimuli ~2x louder than monaural

2. Inhibition: Dichotic stimuli show cross-ear suppression

Computational Complexity:

Frame-by-frame binaural processing requires nested loops. For 1s stereo @ 32kHz: ~0.5-2s CPU, ~0.05-0.2s GPU.

See also

OuterMiddleEarFilter

Stage 1 - Outer/middle ear

Moore2016Spectrum

Stage 2 - Spectral analysis

Moore2016ExcitationPattern

Stage 3 - Excitation

Moore2016SpecificLoudness

Stage 4 - Specific loudness

Moore2016TemporalIntegration

Stage 5 - STL AGC

Moore2016BinauralLoudness

Stage 6 - Binaural inhibition

Moore2016AGC

Stage 7 - LTL AGC

References

[1]

B. C. J. Moore, B. R. Glasberg, and J. Schlittenlacher, “A model of binaural loudness perception based on the banded loudness model,” Acta Acust. united with Acust., vol. 102, no. 5, pp. 824-837, Sep. 2016.

[2]

B. C. J. Moore and B. R. Glasberg, “Modeling binaural loudness,” J. Acoust. Soc. Am., vol. 121, no. 3, pp. 1604-1612, Mar. 2007.

[3]

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

[4]

B. R. Glasberg and B. C. J. Moore, “Derivation of auditory filter shapes from notched-noise data,” Hear. Res., vol. 47, no. 1-2, pp. 103-138, Aug. 1990.

[5]

ISO 226:2003, “Acoustics - Normal equal-loudness-level contours,” International Organization for Standardization, 2003.

[6]

P. Majdak, C. Hollomey, and R. Baumgartner, “AMT 1.x: A toolbox for reproducible research in auditory modeling,” Acta Acust., vol. 6, p. 19, 2022.

__init__(fs=32000, learnable=False, return_stages=False, dtype=torch.float32, outer_middle_ear_kwargs=None, spectrum_kwargs=None, excitation_kwargs=None, specific_loudness_kwargs=None, temporal_integration_kwargs=None, binaural_loudness_kwargs=None, ltl_agc_kwargs=None)

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

Parameters:

• fs (int)

• learnable (bool)

• return_stages (bool)

• dtype (dtype)

• outer_middle_ear_kwargs (Dict[str, Any])

• spectrum_kwargs (Dict[str, Any])

• excitation_kwargs (Dict[str, Any])

• specific_loudness_kwargs (Dict[str, Any])

• temporal_integration_kwargs (Dict[str, Any])

• binaural_loudness_kwargs (Dict[str, Any])

• ltl_agc_kwargs (Dict[str, Any])

forward(audio)

Process stereo audio through complete Moore2016 binaural loudness model.

Parameters:

audio (Tensor) – Stereo audio input. Shape: (batch, 2, n_samples) or (2, n_samples). Audio must be in Pascal (Pa) units and sampled at 32 kHz.

Returns:

If return_stages=False:

(sLoud, lLoud, mLoud) where: - sLoud: Short-term binaural loudness, shape (batch, n_frames) in sone - lLoud: Long-term binaural loudness, shape (batch, n_frames) in sone - mLoud: Maximum long-term loudness, shape (batch,) in sone

If return_stages=True:

((sLoud, lLoud, mLoud), stages) where stages is a dict with: - ‘filtered_left’, ‘filtered_right’: After outer/middle ear - ‘freqs_left’, ‘levels_left’, ‘freqs_right’, ‘levels_right’: Sparse spectrum - ‘excitation_left’, ‘excitation_right’: Excitation patterns - ‘inst_spec_loud_left’, ‘inst_spec_loud_right’: Instantaneous specific loudness - ‘stl_spec_loud_left’, ‘stl_spec_loud_right’: Short-term specific loudness - ‘stl_loud_left’, ‘stl_loud_right’: Short-term scalar loudness - ‘ltl_loud_left’, ‘ltl_loud_right’: Long-term scalar loudness

Return type:

Tuple[Tensor, Tensor, Tensor] | Tuple[Tuple[Tensor, Tensor, Tensor], Dict[str, Any]]

get_parameters()

Get all model parameters.

Returns:

Dictionary with model parameters: - ‘fs’: Sampling rate (32000 Hz) - ‘learnable’: Whether parameters are trainable - ‘return_stages’: Whether intermediate stages are returned

Return type:

Dict[str, Any]

extra_repr()

Extra representation for printing.

Returns:

String representation of 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

King2019

class torch_amt.King2019(fs, flow=80.0, fhigh=8000.0, basef=None, compression_type='brokenstick', compression_n=0.3, compression_knee_db=30.0, dboffset=100.0, adt_hp_fc=3.0, adt_hp_order=1, mflow=2.0, mfhigh=150.0, modbank_nmod=None, modbank_qfactor=1.0, lp_150hz=False, subfs=None, learnable=False, return_stages=False, dtype=torch.float32, filterbank_kwargs=None, compression_kwargs=None, ihc_kwargs=None, adaptation_kwargs=None, modulation_kwargs=None)

Bases: Module

King et al. (2019) auditory model with nonlinear compression.

Implements a computational model of the auditory periphery designed for studying masking of frequency modulation (FM) by amplitude modulation (AM). The model features explicit nonlinear compression (broken-stick or power-law) and adaptation stages, with a modulation filterbank for extracting temporal modulation content.

This implementation follows the MATLAB Auditory Modeling Toolbox (AMT) king2019 configuration and provides a differentiable, GPU-accelerated version suitable for neural network training and optimization.

Algorithm Overview

The model implements a 6-stage auditory processing pipeline:

Stage 1: Gammatone Filterbank

Decomposes signal into \(N\) frequency channels with 1-ERB spacing:

\[\begin{split}y_i(t) = g_i(t) * x(t), \\quad i=1,\\ldots,N\end{split}\]

where \(g_i(t) = t^{3} e^{-2\\pi b_i t} \\cos(2\\pi f_i t)\) is a 4th-order gammatone impulse response.

Stage 2: Nonlinear Compression

Two compression types are available:

Broken-stick compression (default):

\[\begin{split}c_i(t) = \\begin{cases} |y_i(t)|, & \\text{if } L_i(t) < L_{\\text{knee}} \\\\ 10^{(L_{\\text{knee}}/20)} \\cdot \\left(\\frac{|y_i(t)|}{10^{(L_{\\text{knee}}/20)}}\\right)^n, & \\text{if } L_i(t) \\geq L_{\\text{knee}} \\end{cases}\end{split}\]

where \(L_i(t) = 20\\log_{10}(|y_i(t)|/p_{\\text{ref}}) + \\text{dboffset}\), \(L_{\\text{knee}}\) is the knee point (default 30 dB), and \(n\) is the compression exponent (default 0.3).

Power-law compression:

\[c_i(t) = |y_i(t)|^n\]

Stage 3: Inner Hair Cell (IHC) Envelope

Extracts envelope via half-wave rectification and lowpass filtering:

\[\begin{split}e_i(t) = h_{\\text{lp}}(t) * \\max(0, c_i(t))\end{split}\]

King2019 uses 1000 Hz cutoff (1st-order Butterworth lowpass).

Stage 4: Adaptation

High-pass filtering removes DC and slow drifts:

\[\begin{split}a_i(t) = h_{\\text{hp}}(t) * e_i(t)\end{split}\]

Default: 3 Hz cutoff, 1st-order Butterworth highpass.

Stage 5: Optional 150 Hz Lowpass

Pre-filtering before modulation analysis (if lp_150hz=True):

\[a'_i(t) = h_{150}(t) * a_i(t)\]

Stage 6: Modulation Filterbank

Extracts modulation frequencies using bandpass filters:

\[\begin{split}m_{i,k}(t) = h_{\\text{mod},k}(t) * a'_i(t)\end{split}\]

Modulation center frequencies \(f_{\\text{mod},k}\) are logarithmically spaced from \(f_{\\text{low}}\) (default 2 Hz) to \(f_{\\text{high}}\) (default 150 Hz) based on Q-factor. Each filter is a 2nd-order Butterworth bandpass with bandwidth \(\\text{BW} = f_{\\text{mod},k} / Q\).

Output: Tensor of shape \((B, T, F, M)\) where \(M\) is the number of modulation channels.

param fs:

Sampling rate in Hz. Must match the audio sampling rate. Common values: 44100, 48000, 32000 Hz.

type fs:

float

param flow:

Lower frequency bound for gammatone filterbank in Hz. Default: 80 Hz. Determines the lowest auditory channel center frequency.

type flow:

float

param fhigh:

Upper frequency bound for gammatone filterbank in Hz. Default: 8000 Hz. Determines the highest auditory channel center frequency.

type fhigh:

float

param basef:

Base frequency in Hz for centered analysis. If provided, flow and fhigh are automatically computed as basef ± 2 ERB, creating a narrow frequency range centered on basef. Default: None (use flow/fhigh).

type basef:

float | None

param compression_type:

Type of nonlinear compression. Default: ‘brokenstick’.

• 'brokenstick': Two-segment compression with linear below knee, power-law above knee (physiologically realistic)

• 'power': Simple power-law compression \(y = x^n\)

type compression_type:

str

param compression_n:

Compression exponent for power-law segment. Default: 0.3. Typical physiological values: 0.2-0.4.

type compression_n:

float

param compression_knee_db:

Knee point in dB SPL for broken-stick compression. Default: 30 dB. Signals below this level are uncompressed; above are compressed.

type compression_knee_db:

float

param dboffset:

dB full scale convention (calibration). Default: 100 dB SPL.

• 100 dB: MATLAB AMT convention (0 dBFS = 100 dB SPL)

• 94 dB: Alternative convention for specific AMT signals

Must match the signal’s reference level for correct compression.

type dboffset:

float

param adt_hp_fc:

Adaptation highpass cutoff frequency in Hz. Default: 3.0 Hz. Removes DC offset and very slow drifts.

type adt_hp_fc:

float

param adt_hp_order:

Adaptation highpass filter order. Default: 1 (1st-order Butterworth). Higher orders provide steeper roll-off but may introduce artifacts.

type adt_hp_order:

int

param mflow:

Minimum modulation frequency for modulation filterbank in Hz. Default: 2.0 Hz. Lowest temporal modulation captured.

type mflow:

float

param mfhigh:

Maximum modulation frequency for modulation filterbank in Hz. Default: 150.0 Hz. Highest temporal modulation captured.

type mfhigh:

float

param modbank_nmod:

Number of modulation filters. If None, automatically determined by Q-factor spacing (typically 5-10 filters). Default: None (automatic).

type modbank_nmod:

int | None

param modbank_qfactor:

Q-factor for modulation filters, controlling bandwidth: \(\\text{BW} = f_{\\text{mod}} / Q\). Default: 1.0. Higher Q → narrower filters, better frequency resolution.

type modbank_qfactor:

float

param lp_150hz:

Apply 150 Hz lowpass filter before modulation filterbank. Default: False. Useful for limiting analysis to lower modulation frequencies.

type lp_150hz:

bool

param subfs:

Target sampling rate in Hz for downsampling output. If None, no downsampling is applied. Default: None (keep original fs). Reduces computational cost for downstream processing.

type subfs:

float | None

param learnable:

If True, all model stages become trainable with gradient-based optimization. Default: False (fixed parameters). Enables end-to-end model training for task-specific optimization.

type learnable:

bool

param return_stages:

If True, returns intermediate processing stages along with final output. Default: False (only final modulation output). Useful for visualization, analysis, and multi-stage training.

type return_stages:

bool

param dtype:

Data type for computations and parameters. Default: torch.float32. Use torch.float64 for higher precision if needed.

type dtype:

dtype

param **filterbank_kwargs:

Additional keyword arguments passed to GammatoneFilterbank. Common options:

• n (int): Filter order. Default: 4.

• betamul (float): Beta multiplier. Default: 1.0186.

• Other parameters accepted by GammatoneFilterbank.

type **filterbank_kwargs:

Dict[str, Any] | None

param **compression_kwargs:

Additional keyword arguments passed to BrokenStickCompression or PowerCompression depending on compression_type.

For BrokenStickCompression:

• No additional parameters typically needed beyond those in main signature.

For PowerCompression:

• No additional parameters typically needed beyond those in main signature.

type **compression_kwargs:

Dict[str, Any] | None

param **ihc_kwargs:

Additional keyword arguments passed to IHCEnvelope. Common options:

• method (str): IHC method. Fixed to ‘king2019’ for this model.

• Other parameters accepted by IHCEnvelope.

type **ihc_kwargs:

Dict[str, Any] | None

param **adaptation_kwargs:

Additional keyword arguments passed to ButterworthFilter (adaptation highpass filter).

• No additional parameters typically needed beyond fc and order.

type **adaptation_kwargs:

Dict[str, Any] | None

param **modulation_kwargs:

Additional keyword arguments passed to King2019ModulationFilterbank.

• No additional parameters typically needed beyond those in main signature.

type **modulation_kwargs:

Dict[str, Any] | None

fs

Sampling rate in Hz.

Type:

float

flow

Lower frequency bound for filterbank.

Type:

float

fhigh

Upper frequency bound for filterbank.

Type:

float

dboffset

dB full scale calibration.

Type:

float

compression_type

Type of compression (‘brokenstick’ or ‘power’).

Type:

str

learnable

Whether model parameters are trainable.

Type:

bool

return_stages

Whether to return intermediate processing stages.

Type:

bool

dtype

Data type for computations.

Type:

torch.dtype

filterbank

Stage 1: Gammatone auditory filterbank.

Type:

GammatoneFilterbank

compression

Stage 2: Nonlinear compression module.

Type:

BrokenStickCompression or PowerCompression

ihc

Stage 3: Inner hair cell envelope extraction.

Type:

IHCEnvelope

adaptation_filter

Stage 4: Adaptation highpass filter.

Type:

ButterworthFilter

lp150_filter

Stage 5: Optional 150 Hz lowpass filter.

Type:

ButterworthFilter or None

modulation_filterbank

Stage 6: Modulation filterbank module.

Type:

King2019ModulationFilterbank

fc

Center frequencies of auditory channels, shape (num_channels,) in Hz.

Type:

torch.Tensor

mfc

Center frequencies of modulation channels, shape (num_mod_filters,) in Hz.

Type:

torch.Tensor

num_channels

Number of auditory frequency channels (depends on flow, fhigh, ERB spacing).

Type:

int

Input Shape

-----------

x

Audio signal with shape:

• \((B, T)\) - Batch of signals

• \((C, T)\) - Multi-channel input

• \((T,)\) - Single signal

where:

• \(B\) = batch size

• \(C\) = channels

• \(T\) = time samples

Type:

torch.Tensor

Output Shape

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

When ``return_stages=False`` (default)

torch.Tensor

Output tensor with shape \((B, T', F, M)\):

• \(B\) = batch size

• \(T'\) = time samples (possibly downsampled if subfs specified)

• \(F\) = num_channels (number of auditory frequency channels)

• \(M\) = number of modulation channels (depends on mflow, mfhigh, qfactor)

When ``return_stages=True``

Tuple[torch.Tensor, Dict[str, torch.Tensor]]

• First element: modulation output (shape as above)

• Second element: dict with keys:

  • 'gtone_response': After gammatone, shape \((B, F, T)\)

  • 'compressed_response': After compression, shape \((B, F, T)\)

  • 'ihc': After IHC envelope, shape \((B, F, T)\)

  • 'adapted_response': After adaptation, shape \((B, F, T)\)

Examples

Basic usage:

>>> import torch
>>> from torch_amt.models import King2019
>>>
>>> # Create model
>>> model = King2019(fs=48000, basef=5000)
>>>
>>> # Generate 0.5 second tone
>>> audio = torch.randn(2, 24000) * 0.01
>>>
>>> # Process
>>> output = model(audio)
>>> print(f"Input: {audio.shape}, Output: {output.shape}")
Input: torch.Size([2, 24000]), Output: torch.Size([2, 24000, 5, 5])
>>> print(f"Frequency channels: {model.num_channels}")
Frequency channels: 5
>>> print(f"Modulation channels: {len(model.mfc)}")
Modulation channels: 5

With intermediate stages:

>>> model_debug = King2019(fs=48000, basef=5000, return_stages=True)
>>> output, stages = model_debug(audio)
>>>
>>> print(f"Available stages: {list(stages.keys())}")
Available stages: ['gtone_response', 'compressed_response', 'ihc', 'adapted_response']
>>> print(f"After gammatone: {stages['gtone_response'].shape}")
After gammatone: torch.Size([2, 5, 24000])
>>> print(f"After compression: {stages['compressed_response'].shape}")
After compression: torch.Size([2, 5, 24000])

Batch processing:

>>> # Process multiple signals
>>> batch_audio = torch.randn(8, 48000) * 0.01
>>> output_batch = model(batch_audio)
>>> print(f"Batch output: {output_batch.shape}")
Batch output: torch.Size([8, 48000, 5, 5])

Using basef for frequency-specific analysis:

>>> # Analyze around 1 kHz (±2 ERB)
>>> model_1k = King2019(fs=48000, basef=1000)
>>> print(f"Channels: {model_1k.num_channels}")
Channels: 5
>>> print(f"Center frequencies: {model_1k.fc}")
Center frequencies: tensor([ 794.3, 1000.0, 1259.2, ...])

Power-law compression:

>>> # Simple power-law instead of broken-stick
>>> model_power = King2019(fs=48000, compression_type='power', compression_n=0.4)
>>> output_power = model_power(audio)

Custom modulation filterbank:

>>> # Wider modulation range with higher Q
>>> model_mod = King2019(
...     fs=48000,
...     mflow=1.0,
...     mfhigh=200.0,
...     modbank_qfactor=2.0,
...     lp_150hz=True
... )
>>> print(f"Modulation frequencies: {model_mod.mfc}")
Modulation frequencies: tensor([  1.00,   3.00,   9.00,  27.00, ...])

With downsampling:

>>> # Downsample output to 1000 Hz for efficiency
>>> model_ds = King2019(fs=48000, subfs=1000)
>>> audio_long = torch.randn(1, 480000) * 0.01  # 10 seconds
>>> output_ds = model_ds(audio_long)
>>> print(f"Downsampled output: {output_ds.shape}")
Downsampled output: torch.Size([1, 10000, 31, 10])  # T reduced from 480k to 10k

Learnable model for optimization:

>>> model_learnable = King2019(fs=48000, learnable=True)
>>> n_params = sum(p.numel() for p in model_learnable.parameters())
>>> print(f"Trainable parameters: {n_params}")
Trainable parameters: 2156
>>>
>>> # Example training loop
>>> optimizer = torch.optim.Adam(model_learnable.parameters(), lr=1e-3)
>>> # ... training code ...

Accessing center frequencies:

>>> model = King2019(fs=48000)
>>> print(f"Auditory fc (Hz): {model.fc[:5]}")
Auditory fc (Hz): tensor([  80.0,  100.8,  126.9,  159.9,  201.4])
>>> print(f"Modulation mfc (Hz): {model.mfc}")
Modulation mfc (Hz): tensor([  2.00,   5.24,  13.71,  35.89,  93.96])

Notes

Model Configuration:

• Filterbank: Gammatone 4th-order, 1-ERB spacing

• Compression: Broken-stick (knee=30dB, n=0.3) or power-law

• IHC: Half-wave rectification + 1000 Hz lowpass (king2019 method)

• Adaptation: 3 Hz highpass, 1st-order Butterworth

• Modulation: 2-150 Hz, Q=1.0, 2nd-order Butterworth bandpass

Compression Calibration:

The compression stage requires correct calibration via dboffset:

• 100 dB (default): MATLAB AMT convention (0 dBFS = 100 dB SPL)

• 94 dB: Alternative convention for specific AMT signals

The dboffset must match your signal’s reference level. Incorrect calibration leads to improper compression behavior and inaccurate modulation representations.

Example calibration:

# For signals calibrated to 100 dB SPL at 0 dBFS
model = King2019(fs=48000, dboffset=100.0)

# For AMT-specific signals at 94 dB SPL
model = King2019(fs=48000, dboffset=94.0)

basef Parameter:

When basef is specified, flow and fhigh are automatically computed to create a narrow frequency range:

\[\begin{split}\\text{ERB}_{\\text{base}} = \\text{fc2erb}(\\text{basef})\end{split}\]

\[\begin{split}f_{\\text{low}} = \\text{erb2fc}(\\text{ERB}_{\\text{base}} - 2)\end{split}\]

\[\begin{split}f_{\\text{high}} = \\text{erb2fc}(\\text{ERB}_{\\text{base}} + 2)\end{split}\]

This creates ~5 channels centered on basef, useful for frequency-specific analysis (e.g., studying FM masking at a particular carrier frequency).

Computational Complexity:

Processing time scales as:

\[\begin{split}T_{\\text{compute}} \\propto T \\cdot (N_{\\text{filt}} + N_{\\text{filt}} \\cdot N_{\\text{mod}})\end{split}\]

where \(T\) = signal length, \(N_{\\text{filt}}\) = auditory channels (~5-31), \(N_{\\text{mod}}\) = modulation channels (~5-10).

For 1 second @ 48 kHz: ~0.05-0.2 seconds on CPU, ~0.01-0.05 seconds on GPU.

Memory Requirements:

Peak memory with intermediate stages:

\[\begin{split}\\text{Memory} \\approx B \\cdot T \\cdot (F + F \\cdot M) \\cdot 4\\,\\text{bytes}\end{split}\]

For batch=4, 1 second @ 48 kHz, F=5, M=5: ~5-10 MB.

Applications:

The model is particularly suited for:

• Frequency modulation (FM) masking studies

• Amplitude modulation (AM) masking studies

• FM-AM interaction analysis

• Temporal modulation transfer functions (TMTF)

• Psychoacoustic feature extraction

• Auditory scene analysis

See also

GammatoneFilterbank

Stage 1 - Auditory filterbank

BrokenStickCompression

Stage 2 - Nonlinear compression

PowerCompression

Stage 2 - Alternative compression

IHCEnvelope

Stage 3 - Inner hair cell envelope

ButterworthFilter

Stages 4-5 - Filtering

King2019ModulationFilterbank

Stage 6 - Modulation analysis

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, Apr. 2019.

[2]

T. Dau, B. Kollmeier, and A. Kohlrausch, “Modeling auditory processing of amplitude modulation. I. Detection and masking with narrow-band carriers,” J. Acoust. Soc. Am., vol. 102, no. 5, pp. 2892-2905, Nov. 1997.

[3]

R. D. Patterson, I. Nimmo-Smith, J. Holdsworth, and P. Rice, “An efficient auditory filterbank based on the gammatone function,” in Proc. Meet. IOC Speech Group Auditory Modelling, 1988.

[4]

P. Majdak, C. Hollomey, and R. Baumgartner, “AMT 1.x: A toolbox for reproducible research in auditory modeling,” Acta Acust., vol. 6, p. 19, 2022.

__init__(fs, flow=80.0, fhigh=8000.0, basef=None, compression_type='brokenstick', compression_n=0.3, compression_knee_db=30.0, dboffset=100.0, adt_hp_fc=3.0, adt_hp_order=1, mflow=2.0, mfhigh=150.0, modbank_nmod=None, modbank_qfactor=1.0, lp_150hz=False, subfs=None, learnable=False, return_stages=False, dtype=torch.float32, filterbank_kwargs=None, compression_kwargs=None, ihc_kwargs=None, adaptation_kwargs=None, modulation_kwargs=None)

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

Parameters:

• fs (float)

• flow (float)

• fhigh (float)

• basef (float | None)

• compression_type (str)

• compression_n (float)

• compression_knee_db (float)

• dboffset (float)

• adt_hp_fc (float)

• adt_hp_order (int)

• mflow (float)

• mfhigh (float)

• modbank_nmod (int | None)

• modbank_qfactor (float)

• lp_150hz (bool)

• subfs (float | None)

• learnable (bool)

• return_stages (bool)

• dtype (dtype)

• filterbank_kwargs (Dict[str, Any] | None)

• compression_kwargs (Dict[str, Any] | None)

• ihc_kwargs (Dict[str, Any] | None)

• adaptation_kwargs (Dict[str, Any] | None)

• modulation_kwargs (Dict[str, Any] | None)

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(x)

Process audio through the KING2019 model.

Parameters:

x (Tensor) – Input audio signal. Shape: (B, T), (C, T), or (T,).

Returns:

If return_stages=False:

Output tensor of shape (B, T’, F, M) where:

T’ = time samples (possibly downsampled) F = number of frequency channels M = number of modulation channels

If return_stages=True:

Tuple of (output, stages) where stages is a dict with intermediate outputs.

Return type:

Tensor | tuple[Tensor, Dict[str, Any]]

extra_repr()

Extra representation for printing.

Returns:

String representation of key parameters.

Return type:

str

Osses2021

class torch_amt.Osses2021(fs, flow=80.0, fhigh=8000.0, phase_type='minimum', learnable=False, return_stages=False, dtype=torch.float32, headphone_kwargs=None, middleear_kwargs=None, filterbank_kwargs=None, ihc_kwargs=None, adaptation_kwargs=None, modulation_kwargs=None)

Bases: Module

Osses et al. (2021) auditory model with realistic peripheral filtering.

Implements a computational model of the auditory periphery designed for fluctuation strength prediction and other psychoacoustic tasks. The model extends the Dau1997 framework by incorporating headphone and middle ear transfer functions, providing more accurate peripheral modeling suitable for headphone-presented stimuli and perceptual predictions.

This implementation follows the MATLAB Auditory Modeling Toolbox (AMT) osses2021 configuration and provides a differentiable, GPU-accelerated version suitable for neural network training and optimization.

Algorithm Overview

The model implements a 6-stage auditory processing pipeline:

Stage 0a: Headphone Filter

Compensates for headphone and outer ear characteristics using Pralong & Carlile (1996):

\[x_{hp}(t) = h_{hp}(t) * x(t)\]

where \(h_{hp}\) models outer ear + headphone frequency response.

Stage 0b: Middle Ear Filter

Simulates middle ear transmission using Jepsen et al. (2008) model:

\[x_{me}(t) = h_{me}(t) * x_{hp}(t)\]

Bandpass characteristic approximating middle ear transfer function.

Stage 1: Gammatone Filterbank

Decomposes signal into \(N\) frequency channels with 1-ERB spacing:

\[\begin{split}y_i(t) = g_i(t) * x_{me}(t), \\quad i=1,\\ldots,N\end{split}\]

where \(g_i(t) = t^{n-1} e^{-2\\pi b t} \\cos(2\\pi f_c t + \\phi)\), typically \(n=4\) (4th-order).

Stage 2: Inner Hair Cell (IHC) Envelope

Extracts envelope via Breebaart et al. (2001) method:

\[e_i(t) = |y_i(t)|^2 * h_{lp}(t)\]

where \(h_{lp}\) is a lowpass filter extracting the envelope.

Stage 3: Adaptation

Five parallel adaptation loops with time constants \(\\tau_j\):

\[\begin{split}v_i(t) = \\sum_{j=1}^5 a_j v_{i,j}(t)\end{split}\]

where each \(v_{i,j}\) follows:

\[\begin{split}\\frac{dv_{i,j}}{dt} = \\frac{e_i(t) - v_{i,j}(t)}{\\tau_j}\end{split}\]

Osses2021 preset uses \(\\text{limit}=5.0\) (vs 10.0 in Dau1997).

Stage 4: Modulation Filterbank

Extracts modulation content using jepsen2008 preset:

\[m_{i,k}(t) = h_{mod,k}(t) * v_i(t)\]

Configuration: 150 Hz lowpass, attenuation factor \(1/\\sqrt{2}\).

Output: List of \(N\) tensors, tensor \(i\) has shape \((B, M_i, T)\) where \(M_i\) varies per frequency channel.

param fs:

Sampling rate in Hz. Must match the audio sampling rate. Common values: 44100, 48000, 32000 Hz.

type fs:

float

param flow:

Lower frequency bound for gammatone filterbank in Hz. Default: 80 Hz. Determines the lowest auditory channel center frequency.

type flow:

float

param fhigh:

Upper frequency bound for gammatone filterbank in Hz. Default: 8000 Hz. Determines the highest auditory channel center frequency.

type fhigh:

float

param phase_type:

Filter phase characteristic for peripheral filters. Default: ‘minimum’. Options:

• 'minimum': Minimum-phase FIR (causal, introduces group delay)

• 'zero': Zero-phase via filtfilt (non-causal, no phase distortion)

Use ‘minimum’ for real-time/causal processing, ‘zero’ for offline analysis.

type phase_type:

str

param learnable:

If True, all model stages become trainable with gradient-based optimization. Default: False (fixed parameters). When True, enables end-to-end model training for task-specific optimization.

type learnable:

bool

param return_stages:

If True, returns intermediate processing stages along with final output. Default: False (only final modulation output). Useful for visualization, analysis, and multi-stage training.

type return_stages:

bool

param dtype:

Data type for computations and parameters. Default: torch.float32. Use torch.float64 for higher precision if needed (increases memory/computation).

type dtype:

dtype

param **headphone_kwargs:

Additional keyword arguments passed to HeadphoneFilter. Common options:

• filter_type (str): Filter characteristic. Default: ‘pralong1996’.

• Other parameters accepted by HeadphoneFilter.

type **headphone_kwargs:

Dict[str, Any]

param **middleear_kwargs:

Additional keyword arguments passed to MiddleEarFilter. Common options:

• filter_type (str): Filter model. Default: ‘jepsen2008’.

• normalize (bool): Normalize filter response. Default: True.

• Other parameters accepted by MiddleEarFilter.

type **middleear_kwargs:

Dict[str, Any]

param **filterbank_kwargs:

Additional keyword arguments passed to GammatoneFilterbank. Common options:

• n (int): Filter order. Default: 4 (4th-order gammatone).

• bandwidth_factor (float): Bandwidth scaling. Default: 1.0 (1-ERB).

• Other parameters accepted by GammatoneFilterbank.

type **filterbank_kwargs:

Dict[str, Any]

param **ihc_kwargs:

Additional keyword arguments passed to IHCEnvelope. Common options:

• method (str): Extraction method. Default: ‘breebaart2001’.

• cutoff (float): Lowpass cutoff frequency in Hz.

• Other parameters accepted by IHCEnvelope.

type **ihc_kwargs:

Dict[str, Any]

param **adaptation_kwargs:

Additional keyword arguments passed to AdaptLoop. Common options:

• preset (str): Configuration preset. Default: ‘osses2021’.

• limit (float): Adaptation limit. Default: 5.0.

• num_loops (int): Number of adaptation loops. Default: 5.

• Other parameters accepted by AdaptLoop.

type **adaptation_kwargs:

Dict[str, Any]

param **modulation_kwargs:

Additional keyword arguments passed to ModulationFilterbank. Common options:

• preset (str): Configuration preset. Default: ‘jepsen2008’.

• lowpass_cutoff (float): Lowpass frequency in Hz. Default: 150.

• att_factor (float): Attenuation factor. Default: 1/√2.

• Other parameters accepted by ModulationFilterbank.

• Other parameters accepted by ModulationFilterbank.

type **modulation_kwargs:

Dict[str, Any]

fs

Sampling rate in Hz.

Type:

float

flow

Lower frequency bound for filterbank.

Type:

float

fhigh

Upper frequency bound for filterbank.

Type:

float

phase_type

Filter phase characteristic (‘minimum’ or ‘zero’).

Type:

str

learnable

Whether model parameters are trainable.

Type:

bool

return_stages

Whether to return intermediate processing stages.

Type:

bool

dtype

Data type for computations.

Type:

torch.dtype

headphone

Stage 0a: Headphone and outer ear filter.

Type:

HeadphoneFilter

middleear

Stage 0b: Middle ear transmission filter.

Type:

MiddleEarFilter

filterbank

Stage 1: Gammatone auditory filterbank.

Type:

GammatoneFilterbank

ihc

Stage 2: Inner hair cell envelope extraction.

Type:

IHCEnvelope

adaptation

Stage 3: Adaptation loops module.

Type:

AdaptLoop

modulation

Stage 4: Modulation filterbank.

Type:

ModulationFilterbank

fc

Center frequencies of auditory channels, shape (num_channels,) in Hz.

Type:

torch.Tensor

num_channels

Number of auditory frequency channels (depends on flow, fhigh, ERB spacing).

Type:

int

Input Shape

-----------

x

Audio signal with shape:

• \((B, T)\) - Batch of signals

• \((C, T)\) - Multi-channel input

• \((T,)\) - Single signal

where:

• \(B\) = batch size

• \(C\) = channels

• \(T\) = time samples

Type:

torch.Tensor

Output Shape

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

When ``return_stages=False`` (default)

List[torch.Tensor]

List of \(N\) tensors (one per frequency channel), where tensor \(i\) has shape \((B, M_i, T)\):

• \(N\) = num_channels (number of auditory frequency channels)

• \(M_i\) = number of modulation channels for frequency channel \(i\)

• \(T\) = time samples

Note: \(M_i\) varies across frequency channels due to jepsen2008 preset.

When ``return_stages=True``

Tuple[List[torch.Tensor], Dict[str, torch.Tensor]]

• First element: modulation output (List as above)

• Second element: dict with keys:

  • 'headphone': After headphone filter, shape \((B, T)\)

  • 'middleear': After middle ear filter, shape \((B, T)\)

  • 'filterbank': After gammatone, shape \((B, N, T)\)

  • 'ihc': After IHC envelope, shape \((B, N, T)\)

  • 'adaptation': After adaptation, shape \((B, N, T)\)

Examples

Basic usage:

>>> import torch
>>> from torch_amt.models import Osses2021
>>>
>>> # Create model
>>> model = Osses2021(fs=44100)
>>>
>>> # Generate 1 second tone
>>> audio = torch.randn(1, 44100) * 0.01
>>>
>>> # Process
>>> output = model(audio)
>>> print(f"Number of frequency channels: {len(output)}")
Number of frequency channels: 31
>>> print(f"First channel shape (B, M, T): {output[0].shape}")
First channel shape (B, M, T): torch.Size([1, 13, 44100])
>>> print(f"Last channel shape: {output[-1].shape}")
Last channel shape: torch.Size([1, 8, 44100])

With intermediate stages:

>>> model_debug = Osses2021(fs=44100, return_stages=True)
>>> output, stages = model_debug(audio)
>>>
>>> print(f"Available stages: {list(stages.keys())}")
Available stages: ['headphone', 'middleear', 'filterbank', 'ihc', 'adaptation']
>>> print(f"Filterbank output: {stages['filterbank'].shape}")
Filterbank output: torch.Size([1, 31, 44100])
>>> print(f"After adaptation: {stages['adaptation'].shape}")
After adaptation: torch.Size([1, 31, 44100])

Batch processing:

>>> # Process multiple signals
>>> batch_audio = torch.randn(4, 44100) * 0.01
>>> output_batch = model(batch_audio)
>>> print(f"Batch output: {output_batch[0].shape}")
Batch output: torch.Size([4, 13, 44100])

Zero-phase filters for offline analysis:

>>> model_zero = Osses2021(fs=44100, phase_type='zero')
>>> output_zero = model_zero(audio)
>>> # No phase distortion, suitable for offline processing

Custom frequency range:

>>> # Narrow frequency range
>>> model_narrow = Osses2021(fs=44100, flow=500, fhigh=4000)
>>> print(f"Channels: {model_narrow.num_channels}")
Channels: 15
>>> print(f"Center frequencies: {model_narrow.fc}")
Center frequencies: tensor([ 500.,  631., ..., 3175., 4000.])

Custom submodule parameters:

>>> # Custom adaptation limit
>>> model_custom = Osses2021(
...     fs=44100,
...     adaptation_kwargs={'limit': 10.0, 'num_loops': 7}
... )
>>>
>>> # Custom modulation filterbank
>>> model_mod = Osses2021(
...     fs=44100,
...     modulation_kwargs={'lowpass_cutoff': 200.0}
... )

Learnable model for optimization:

>>> model_learnable = Osses2021(fs=44100, learnable=True)
>>> n_params = sum(p.numel() for p in model_learnable.parameters())
>>> print(f"Trainable parameters: {n_params}")
Trainable parameters: 2648
>>>
>>> # Example training loop
>>> optimizer = torch.optim.Adam(model_learnable.parameters(), lr=1e-3)
>>> # ... training code ...

Accessing center frequencies:

>>> model = Osses2021(fs=44100)
>>> print(f"Center frequencies (Hz): {model.fc[:5]}")
Center frequencies (Hz): tensor([  80.00,  100.79,  126.96,  159.91,  201.42])
>>> print(f"Frequency range: {model.fc[0]:.1f} - {model.fc[-1]:.1f} Hz")
Frequency range: 80.0 - 8000.0 Hz

Notes

Model Configuration:

• Headphone: Pralong & Carlile (1996) outer ear + headphone compensation

• Middle Ear: Jepsen et al. (2008) bandpass characteristic

• Filterbank: Gammatone 4th-order, 1-ERB spacing

• IHC: Breebaart et al. (2001) envelope extraction method

• Adaptation: 5 loops, limit=5.0 (osses2021 preset)

• Modulation: jepsen2008 preset (150 Hz lowpass, att_factor=1/√2)

Customizing Submodule Parameters:

All submodules can be customized through dedicated kwargs dictionaries:

• Use headphone_kwargs to pass parameters to HeadphoneFilter

• Use middleear_kwargs to pass parameters to MiddleEarFilter

• Use filterbank_kwargs to pass parameters to GammatoneFilterbank

• Use ihc_kwargs to pass parameters to IHCEnvelope

• Use adaptation_kwargs to pass parameters to AdaptLoop

• Use modulation_kwargs to pass parameters to ModulationFilterbank

The learnable, dtype, and phase_type parameters are always centralized and applied to all submodules automatically. Custom parameters override defaults while maintaining the Osses2021 model structure.

Phase Type Selection:

The phase_type parameter controls peripheral filter characteristics:

• ‘minimum’ (default): Causal minimum-phase FIR filters

  • Introduces frequency-dependent group delay

  • Suitable for real-time or causal processing

  • Matches physiological phase response

• ‘zero’: Zero-phase filtering via filtfilt

  • No phase distortion (symmetric impulse response)

  • Non-causal (requires future samples)

  • Better for offline analysis and visualization

Choose based on application: real-time → ‘minimum’, offline → ‘zero’.

Output Format (List of Tensors):

Unlike other models that return a single tensor, Osses2021 returns a List[torch.Tensor] because each frequency channel has a different number of modulation channels:

output = model(audio)  # List of N tensors
for i, channel_output in enumerate(output):
    print(f"Channel {i}: {channel_output.shape}")
    # Channel 0: torch.Size([B, 13, T])
    # Channel 1: torch.Size([B, 12, T])
    # ...

This reflects the jepsen2008 modulation filterbank configuration where higher frequency channels have fewer modulation channels. To work with a single tensor, you can concatenate or pad as needed for your application.

Computational Complexity:

Processing time scales as:

\[\begin{split}T_{compute} \\propto T \\cdot (N_{filter} + N_{filt} \\cdot N_{mod})\end{split}\]

where \(T\) = signal length, \(N_{filter}\) = peripheral filter taps, \(N_{filt}\) = number of frequency channels (~31), \(N_{mod}\) = modulation channels per frequency (~8-13).

For 1 second @ 44.1 kHz: ~0.1-0.5 seconds on CPU, ~0.01-0.05 seconds on GPU.

Memory Requirements:

Peak memory with intermediate stages:

\[\begin{split}Memory \\approx B \\cdot T \\cdot (N + \\sum_i M_i) \\cdot 4\\,\\text{bytes}\end{split}\]

For batch=8, 1 second @ 44.1 kHz: ~40-80 MB.

Applications:

The model is particularly suited for:

• Fluctuation strength prediction (original application)

• Roughness and amplitude modulation detection

• Headphone-presented stimuli modeling

• Psychoacoustic feature extraction

• Perceptual quality assessment

• Machine learning auditory features

See also

HeadphoneFilter

Stage 0a - Headphone and outer ear filtering

MiddleEarFilter

Stage 0b - Middle ear transmission

GammatoneFilterbank

Stage 1 - Auditory filterbank

IHCEnvelope

Stage 2 - Inner hair cell envelope

AdaptLoop

Stage 3 - Adaptation loops

ModulationFilterbank

Stage 4 - Modulation analysis

References

[1]

A. Osses Vecchi, L. E. García, and A. Kohlrausch, “Modelling the sensation of fluctuation strength,” in Proc. Forum Acusticum 2020, Lyon, France, 2020, pp. 367-371.

[2]

D. Pralong and S. Carlile, “The role of individualized headphone calibration for the generation of high fidelity virtual auditory space,” J. Acoust. Soc. Am., vol. 100, no. 6, pp. 3785-3793, Dec. 1996.

[3]

M. L. Jepsen, S. D. Ewert, and T. Dau, “A computational model of human auditory signal processing and perception,” J. Acoust. Soc. Am., vol. 124, no. 1, pp. 422-438, Jul. 2008.

[4]

J. Breebaart, S. van de Par, and A. Kohlrausch, “Binaural processing model based on contralateral inhibition. I. Model structure,” J. Acoust. Soc. Am., vol. 110, no. 2, pp. 1074-1088, Aug. 2001.

[5]

T. Dau, D. Püschel, and A. Kohlrausch, “A quantitative model of the ‘effective’ signal processing in the auditory system. I. Model structure,” J. Acoust. Soc. Am., vol. 99, no. 6, pp. 3615-3622, Jun. 1996.

[6]

P. Majdak, C. Hollomey, and R. Baumgartner, “AMT 1.x: A toolbox for reproducible research in auditory modeling,” Acta Acust., vol. 6, p. 19, 2022.

__init__(fs, flow=80.0, fhigh=8000.0, phase_type='minimum', learnable=False, return_stages=False, dtype=torch.float32, headphone_kwargs=None, middleear_kwargs=None, filterbank_kwargs=None, ihc_kwargs=None, adaptation_kwargs=None, modulation_kwargs=None)

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

Parameters:

• fs (float)

• flow (float)

• fhigh (float)

• phase_type (str)

• learnable (bool)

• return_stages (bool)

• dtype (dtype)

• headphone_kwargs (Dict[str, Any])

• middleear_kwargs (Dict[str, Any])

• filterbank_kwargs (Dict[str, Any])

• ihc_kwargs (Dict[str, Any])

• adaptation_kwargs (Dict[str, Any])

• modulation_kwargs (Dict[str, Any])

forward(x)

Process audio through the Osses2021 model.

Parameters:

x (Tensor) – Input audio signal. Shape: (B, T), (C, T), or (T,).

Returns:

If return_stages=False:

List of tensors (one per frequency channel), each shape (B, M, T).

If return_stages=True:

Tuple of (output, stages) where stages is a dict with intermediate outputs.

Return type:

List[Tensor] | tuple[List[Tensor], Dict[str, Any]]

extra_repr()

Extra representation for printing.

Returns:

String representation of module parameters.

Return type:

str

distribute_gradients()

Distribute gradients to grouped filter coefficients in FastModulationFilterbank.

Call this method after loss.backward() to ensure all filter coefficients in the modulation filterbank receive gradient updates. This is necessary when using FastModulationFilterbank with learnable=True, as filters are grouped for efficiency and gradients need to be shared across group members.

Notes

This method should be called in the training loop:

>>> model = Osses2021(fs=44100, learnable=True)
>>> output = model(input_signal)
>>> loss = criterion(output, target)
>>> loss.backward()
>>> model.distribute_gradients()  # ← Important for FastModulationFilterbank!
>>> optimizer.step()

If the modulation filterbank doesn’t have a distribute_gradients method (e.g., using standard ModulationFilterbank), this is a no-op.

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

Paulick2024

class torch_amt.Paulick2024(fs=44100, flow=125.0, fhigh=8000.0, n_channels=50, use_outerear=True, learnable=False, return_stages=False, filter_type='efilt', dtype=torch.float32, outerear_kwargs=None, drnl_kwargs=None, ihc_kwargs=None, adaptation_kwargs=None, modulation_kwargs=None)

Bases: Module

Paulick et al. (2024) CASP model for auditory processing.

Implements the revised Computational Auditory Signal Processing and Perception (CASP) model, an advanced auditory periphery simulation with physiologically accurate nonlinear processing and integrated decision-making capabilities for psychophysical task modeling.

The model extends Jepsen et al. (2008) with improved IHC transduction, frequency-dependent adaptation, and comprehensive decision-making methods for detection, discrimination, and masking studies. It provides both the auditory internal representation and psychophysical decision mechanisms in a single unified framework.

This implementation follows the MATLAB Auditory Modeling Toolbox (AMT) implementation and provides a differentiable, GPU-accelerated version suitable for neural network training and optimization.

Algorithm Overview

The model implements a 6-stage auditory processing pipeline with optional decision-making post-processing:

Stage 0: Outer/Middle Ear Filter (Optional)

Applies free-field or diffuse-field head-related transfer function:

\[\begin{split}x_{\\text{ear}}(t) = h_{\\text{outer}}(t) * x(t)\end{split}\]

Transfer function from Lopezpoveda & Meddis (2001), compensating for headphone presentation to simulate natural listening conditions.

Stage 1: DRNL Filterbank

Dual Resonance NonLinear (DRNL) filterbank with 50 channels spanning 125-8000 Hz with ERB spacing. Each channel combines linear and nonlinear paths:

Linear path:

\[\begin{split}y_{\\text{lin},i}(t) = g_{\\text{lin}} \\cdot [\\text{GT}_{\\text{lin}}(t) * x_{\\text{ear}}(t)]\end{split}\]

Nonlinear path with compression:

\[\begin{split}y_{\\text{nl},i}(t) = g_{\\text{nl}} \\cdot [\\text{LP}(t) * (\\text{GT}_{\\text{nl}}(t) * x_{\\text{ear}}(t))^p]\end{split}\]

Combined output:

\[\begin{split}y_i(t) = y_{\\text{lin},i}(t) + y_{\\text{nl},i}(t)\end{split}\]

where GT = 4th-order gammatone, LP = lowpass filter, \(p \\approx 0.2\) (compression).

Stage 2: IHC Transduction (Paulick2024-specific)

3-stage physiological inner hair cell model with asymmetric compression:

Stage 2a - Half-wave rectification:

\[\begin{split}v_1(t) = \\max(0, y_i(t))\end{split}\]

Stage 2b - Asymmetric compression:

\[\begin{split}v_2(t) = \\text{sign}(v_1) \\cdot |v_1|^{0.23}\end{split}\]

Stage 2c - 1st-order lowpass (1500 Hz):

\[\begin{split}v_{\\text{IHC}}(t) = h_{\\text{LP},1500}(t) * v_2(t)\end{split}\]

Stage 3: Adaptation Loops (5 loops, frequency-dependent)

Multi-stage adaptation using parallel feedback loops:

\[\begin{split}v_{\\text{adapt}}(t) = \\sum_{k=1}^{5} a_{1,k}(f_c) \\cdot [v_{\\text{IHC}}(t) - b_{0,k}(f_c) \\cdot s_k(t)]\end{split}\]

where \(s_k(t)\) is the state of loop \(k\) with time constant \(\\tau_k\).

Time constants (Paulick2024 preset): - Loop 1: \(\\tau_1 = 0.005\) s (5 ms, fast) - Loop 2: \(\\tau_2 = 0.050\) s (50 ms, medium) - Loop 3: \(\\tau_3 = 0.129\) s (129 ms, slow) - Loop 4: \(\\tau_4 = 0.253\) s (253 ms, very slow) - Loop 5: \(\\tau_5 = 0.500\) s (500 ms, ultra slow)

Stage 4: Resampling

Downsample from 44100 Hz to 11025 Hz (÷4) using sinc interpolation:

\[\begin{split}v_{\\text{resamp}}[n] = v_{\\text{adapt}}(t) \\Big|_{t=n/f_{s,\\text{new}}}\end{split}\]

where \(f_{s,\\text{new}} = f_s / 4 = 11025\) Hz.

Stage 5: Modulation Filterbank

Extracts amplitude modulation content with 8 channels (Paulick2024 preset):

• Lowpass: 0 Hz (DC, cutoff 2.5 Hz)

• Bandpass: Geometric progression with ratio 5/3

  \[\begin{split}f_{\\text{mod},k} = f_{\\text{mod},1} \\cdot (5/3)^{k-1}, \\quad k=1,\\ldots,7\end{split}\]

Center frequencies: [5, 8.33, 13.89, 23.15, 38.58, 64.30, 107.17, 128.6] Hz

Number of modulation filters varies per auditory channel based on upper limit:

\[\begin{split}f_{\\text{mod,max}}(f_c) = 0.25 \\cdot f_c\end{split}\]

Stage 6: Decision-Making (Optional post-processing)

Psychophysical decision mechanisms:

1. ROI Selection: Extract time/frequency/modulation regions of interest

2. Template Correlation: Cross-correlate with internal template

3. Decision Variable: Compute metric (RMS, mean, max, L2)

4. Binary Decision: Threshold or ML-based classification

Output: List of \(N_{\\text{chan}}\) tensors, each with shape \((B, M_i, T_{\\text{resamp}})\) where \(M_i\) ∈ [6, 8] depends on \(f_c\).

param fs:

Sampling rate in Hz. Must match the audio sampling rate. Default: 44100 Hz. Common values: 44100, 48000 Hz.

Note: Resampling to fs/4 occurs internally after adaptation.

type fs:

float

param flow:

Lower frequency bound for DRNL filterbank in Hz. Default: 125 Hz. Determines the lowest auditory channel center frequency.

Typical range: 80-200 Hz for speech/music applications.

type flow:

float

param fhigh:

Upper frequency bound for DRNL filterbank in Hz. Default: 8000 Hz. Determines the highest auditory channel center frequency.

Typical range: 4000-12000 Hz. Higher values increase computational cost.

type fhigh:

float

param n_channels:

Number of auditory frequency channels (DRNL filterbank). Default: 50.

More channels → better frequency resolution but higher computational cost. Paulick2024 paper uses 50 channels for detailed spectral analysis.

type n_channels:

int

param use_outerear:

Whether to apply outer/middle ear filtering. Default: True.

• True: Applies free-field HRTF (simulates natural listening)

• False: Skip ear filtering (for already compensated signals)

Use True for headphone presentation, False for loudspeaker or pre-compensated stimuli.

type use_outerear:

bool

param learnable:

If True, all model stages become trainable with gradient-based optimization. Default: False (fixed physiological parameters).

Enables end-to-end model training for task-specific optimization or hearing loss parameter fitting.

type learnable:

bool

param return_stages:

If True, returns intermediate processing stages along with final output. Default: False (only final modulation representation).

Useful for visualization, analysis, debugging, and multi-stage training. Returns dict with keys: [‘outerear’, ‘drnl’, ‘ihc’, ‘adaptation’, ‘resampled’].

type return_stages:

bool

param filter_type:

Type of modulation filters. Default: ‘efilt’.

• 'efilt': MATLAB AMT compatible (complex frequency-shifted lowpass, resonant response, asymmetric frequency response). Original implementation.

• 'butterworth': Conceptually correct symmetric bandpass filters (2nd-order, better frequency resolution).

Use ‘efilt’ for exact MATLAB AMT compatibility, ‘butterworth’ for cleaner frequency responses.

type filter_type:

str

param dtype:

Data type for computations and parameters. Default: torch.float32. Use torch.float64 for higher numerical precision if needed (slower).

type dtype:

dtype

param **outerear_kwargs:

Additional keyword arguments passed to OuterMiddleEarFilter. Common options:

• compensation_type (str): Type of compensation. Default: ‘tfOuterMiddle1997’.

• field_type (str): Field type (‘free’ or ‘diffuse’). Default: ‘free’.

• Other parameters accepted by OuterMiddleEarFilter.

type **outerear_kwargs:

Dict[str, Any] | None

param **drnl_kwargs:

Additional keyword arguments passed to DRNLFilterbank. Common options:

• subject (str): Subject type (‘NH’ for normal hearing). Default: ‘NH’.

• model (str): Model version. Default: ‘paulick2024’.

• Other parameters accepted by DRNLFilterbank.

type **drnl_kwargs:

Dict[str, Any] | None

param **ihc_kwargs:

Additional keyword arguments passed to IHCPaulick2024.

• No additional parameters typically needed (Paulick2024-specific IHC).

type **ihc_kwargs:

Dict[str, Any] | None

param **adaptation_kwargs:

Additional keyword arguments passed to AdaptLoop. Common options:

• preset (str): Adaptation preset. Default: ‘paulick2024’.

• Other parameters accepted by AdaptLoop.

type **adaptation_kwargs:

Dict[str, Any] | None

param **modulation_kwargs:

Additional keyword arguments passed to ModulationFilterbank. Common options:

• preset (str): Modulation preset. Default: ‘paulick2024’.

• filter_type (str): Filter type. Inherited from main parameter.

• Other parameters accepted by ModulationFilterbank.

type **modulation_kwargs:

Dict[str, Any] | None

fs

Sampling rate in Hz.

Type:

float

flow

Lower frequency bound for DRNL filterbank.

Type:

float

fhigh

Upper frequency bound for DRNL filterbank.

Type:

float

n_channels

Number of auditory frequency channels.

Type:

int

use_outerear

Whether outer/middle ear filtering is applied.

Type:

bool

learnable

Whether model parameters are trainable.

Type:

bool

return_stages

Whether to return intermediate processing stages.

Type:

bool

filter_type

Type of modulation filters (‘efilt’ or ‘butterworth’).

Type:

str

dtype

Data type for computations.

Type:

torch.dtype

outer_middle_ear

Stage 0: Outer/middle ear filtering module (if use_outerear=True).

Type:

OuterMiddleEarFilter or None

drnl

Stage 1: Dual Resonance NonLinear filterbank.

Type:

DRNLFilterbank

ihc

Stage 2: Inner hair cell transduction module (Paulick2024-specific).

Type:

IHCPaulick2024

adaptation

Stage 3: Multi-stage adaptation loops module.

Type:

AdaptLoop

modulation

Stage 5: Modulation filterbank module (operates at resampled rate).

Type:

ModulationFilterbank

resampler

Stage 4: Resampling module (÷4 downsampling).

Type:

torchaudio.transforms.Resample or None

fc

Center frequencies of DRNL channels, shape (n_channels,) in Hz.

Type:

torch.Tensor

num_channels

Number of auditory frequency channels (same as n_channels).

Type:

int

resample_factor

Resampling factor (default: 4 for 44100 Hz → 11025 Hz).

Type:

int

fs_resampled

Sampling rate after resampling (fs / resample_factor = 11025 Hz).

Type:

float

Input Shape

-----------

x

Audio signal with shape:

• \((B, T)\) - Batch of signals

• \((T,)\) - Single signal

where:

• \(B\) = batch size

• \(T\) = time samples at original sampling rate fs

Type:

torch.Tensor

Output Shape

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

When ``return_stages=False`` (default)

List[torch.Tensor]

List of length \(N_{\\text{channels}}\) (default: 50), one tensor per auditory frequency channel.

Each tensor has shape:

• \((B, M_i, T_{\\text{resamp}})\) for batched input

• \((M_i, T_{\\text{resamp}})\) for single signal input

where:

• \(M_i\) ∈ [6, 8] = number of modulation filters for channel \(i\) (varies per channel due to dynamic upper frequency limit)

• \(T_{\\text{resamp}} = T / 4\) (for fs=44100 → 11025 Hz)

Example: For 1-second audio at 44100 Hz: - Input: (B, 44100) - Output: List of 50 tensors, each (B, ~7-8, 11025)

When ``return_stages=True``

Tuple[List[torch.Tensor], Dict[str, torch.Tensor]]

• First element: modulation representation (as above)

• Second element: dict with intermediate stages:

  • 'outerear': After ear filtering, shape \((B, T)\)

  • 'drnl': After DRNL filterbank, shape \((B, F, T)\)

  • 'ihc': After IHC transduction, shape \((B, F, T)\)

  • 'adaptation': After adaptation, shape \((B, F, T)\)

  • 'resampled': After resampling, shape \((B, F, T_{\\text{resamp}})\)

Examples

Basic usage:

>>> import torch
>>> from torch_amt.models import Paulick2024
>>>
>>> # Create model
>>> model = Paulick2024(fs=44100)
>>>
>>> # Generate 1 second audio
>>> audio = torch.randn(2, 44100) * 0.01
>>>
>>> # Process
>>> internal_repr = model(audio)
>>> print(f"Number of frequency channels: {len(internal_repr)}")
Number of frequency channels: 50
>>> print(f"First channel shape: {internal_repr[0].shape}")
First channel shape: torch.Size([2, 8, 11025])
>>> print(f"Modulation filters in channel 0: {internal_repr[0].shape[1]}")
Modulation filters in channel 0: 8

With intermediate stages:

>>> model_debug = Paulick2024(fs=44100, return_stages=True)
>>> internal_repr, stages = model_debug(audio)
>>>
>>> print(f"Available stages: {list(stages.keys())}")
Available stages: ['outerear', 'drnl', 'ihc', 'adaptation', 'resampled']
>>> print(f"After DRNL: {stages['drnl'].shape}")
After DRNL: torch.Size([2, 50, 44100])
>>> print(f"After resampling: {stages['resampled'].shape}")
After resampling: torch.Size([2, 50, 11025])

Detection task:

>>> # Simple detection with threshold
>>> signal = torch.randn(1, 44100) * 0.01
>>> decision = model.detection_task(signal, threshold=0.5)
>>> print(f"Detection: {'Signal detected' if decision.item() else 'No signal'}")
Detection: Signal detected
>>>
>>> # Detection with custom metric
>>> decision_max = model.detection_task(signal, threshold=0.3, metric='max')

Discrimination task:

>>> # Two-interval forced choice
>>> signal1 = torch.randn(1, 44100) * 0.01
>>> signal2 = torch.randn(1, 44100) * 0.015  # Slightly louder
>>> choice = model.discrimination_task(signal1, signal2, criterion='rms')
>>> print(f"Chose interval: {choice.item() + 1}")
Chose interval: 2

ROI selection:

>>> # Extract specific time window and frequency range
>>> full_repr = model(audio)
>>>
>>> # Focus on 200-500 ms, channels 10-20, slow modulations (0-3)
>>> roi = model.roi_selection(
...     full_repr,
...     time_window=(0.2, 0.5),
...     channel_range=(10, 20),
...     modulation_range=(0, 3)
... )
>>> print(f"ROI channels: {len(roi)}")
ROI channels: 10
>>> print(f"ROI shape: {roi[0].shape}")
ROI shape: torch.Size([2, 3, 3307])  # 3 mod filters, ~0.3s at 11025 Hz

Template correlation:

>>> # Compute cross-correlation with internal template
>>> test_signal = torch.randn(1, 44100) * 0.01
>>> template_signal = torch.randn(1, 44100) * 0.01
>>>
>>> test_repr = model(test_signal)
>>> template_repr = model(template_signal)
>>>
>>> correlation = model.template_correlation(test_repr, template_repr)
>>> print(f"Template correlation: {correlation.item():.4f}")
Template correlation: 0.8523

Custom decision variable:

>>> # Compute RMS energy with channel weighting
>>> channel_weights = torch.ones(50)
>>> channel_weights[20:30] *= 2.0  # Emphasize mid-frequency channels
>>>
>>> dv = model.compute_decision_variable(
...     full_repr,
...     metric='rms',
...     channel_weights=channel_weights
... )
>>> print(f"Decision variable: {dv}")
Decision variable: tensor([0.0234, 0.0198])

Batch processing:

>>> # Process multiple signals
>>> batch_audio = torch.randn(8, 44100) * 0.01
>>> batch_repr = model(batch_audio)
>>> print(f"Batch output - Channel 0: {batch_repr[0].shape}")
Batch output - Channel 0: torch.Size([8, 8, 11025])

Without outer ear filtering:

>>> # For pre-compensated or loudspeaker signals
>>> model_noear = Paulick2024(fs=44100, use_outerear=False)
>>> output_noear = model_noear(audio)

Custom frequency range:

>>> # Focus on specific frequency region
>>> model_lowfreq = Paulick2024(fs=44100, flow=80, fhigh=2000, n_channels=20)
>>> print(f"Frequency channels: {model_lowfreq.num_channels}")
Frequency channels: 20
>>> print(f"Center freq range: {model_lowfreq.fc[0]:.1f} - {model_lowfreq.fc[-1]:.1f} Hz")
Center freq range: 80.0 - 2000.0 Hz

Learnable model for optimization:

>>> model_learnable = Paulick2024(fs=44100, learnable=True)
>>> n_params = sum(p.numel() for p in model_learnable.parameters() if p.requires_grad)
>>> print(f"Trainable parameters: {n_params}")
Trainable parameters: 17223
>>>
>>> # Example training loop
>>> optimizer = torch.optim.Adam(model_learnable.parameters(), lr=1e-4)
>>> # ... training code ...

Butterworth modulation filters:

>>> # Use symmetric bandpass filters instead of efilt
>>> model_butter = Paulick2024(fs=44100, filter_type='butterworth')
>>> output_butter = model_butter(audio)

Accessing center frequencies:

>>> model = Paulick2024(fs=44100)
>>> print(f"Auditory fc (first 5): {model.fc[:5]}")
Auditory fc (first 5): tensor([ 125.0,  145.8,  170.1,  198.5,  231.6])
>>> print(f"Resampled rate: {model.fs_resampled} Hz")
Resampled rate: 11025.0 Hz

Notes

Model Configuration:

• Filterbank: DRNL 50 channels, 125-8000 Hz, dual-path nonlinear

• IHC: Paulick2024-specific 3-stage (rectify → compress → lowpass)

• Adaptation: 5 loops with frequency-dependent parameters

• Resampling: ÷4 (44100 → 11025 Hz) via sinc interpolation

• Modulation: 8 channels, 0-128.6 Hz, geometric spacing 5/3

Computational Complexity:

Processing time scales as:

\[\begin{split}T_{\\text{compute}} \\propto T \\cdot (N_{\\text{filt}} + N_{\\text{filt}} \\cdot N_{\\text{mod}})\end{split}\]

where \(T\) = signal length, \(N_{\\text{filt}}\) = 50 (auditory channels), \(N_{\\text{mod}}\) ≈ 6-8 (modulation channels).

For 1 second @ 44100 Hz: ~0.2-0.5 seconds on CPU, ~0.05-0.15 seconds on GPU.

Memory Requirements:

Peak memory with intermediate stages:

\[\begin{split}\\text{Memory} \\approx B \\cdot T \\cdot F \\cdot (1 + M \\cdot 0.25) \\cdot 8\\,\\text{bytes}\end{split}\]

For batch=4, 1 second @ 44100 Hz, F=50, M=8: ~50-70 MB (float64). With float32: ~25-35 MB.

Applications:

The model is particularly suited for:

• Psychophysical detection threshold prediction

• Discrimination task modeling (2AFC, 3AFC, etc.)

• Forward/backward masking studies

• Spectral and temporal masking

• Amplitude modulation detection/discrimination

• Hearing aid processing evaluation

• Hearing loss simulation and compensation

• Speech intelligibility prediction

Decision-Making Methods:

The model includes 6 decision-making methods for psychophysical modeling:

1. roi_selection: Extract time/frequency/modulation regions of interest

2. template_correlation: Cross-correlate with internal template

3. compute_decision_variable: Compute metric (RMS, mean, max, L2)

4. make_decision: Binary threshold or ML-based decision

5. detection_task: Complete detection task (signal vs. noise)

6. discrimination_task: Complete discrimination task (2AFC, etc.)

These methods operate on the internal representation and provide a complete framework for modeling psychophysical experiments.

Learnable Parameters (if learnable=True):

• OuterMiddleEarFilter: ~16000 parameters (FIR filter taps)

• DRNLFilterbank: ~400 parameters (nonlinearity coefficients, gains)

• IHCPaulick2024: 13 parameters (physiological constants)

• AdaptLoop: 10 parameters (time constants a1, b0 for 5 loops)

• ModulationFilterbank: ~800 parameters (filter coefficients)

Total: ~17223 trainable parameters

See also

DRNLFilterbank

Stage 1 - Dual Resonance NonLinear filterbank

OuterMiddleEarFilter

Stage 0 - Outer/middle ear filtering

IHCPaulick2024

Stage 2 - Inner hair cell transduction

AdaptLoop

Stage 3 - Multi-stage adaptation

ModulationFilterbank

Stage 5 - Modulation analysis

References

[1]

L. Paulick, H. Relaño-Iborra, and T. Dau, “The Computational Auditory Signal Processing and Perception Model (CASP): A Revised Version,” bioRxiv, 2024.

[2]

M. Jepsen, S. Ewert, and T. Dau, “A computational model of human auditory signal processing and perception,” J. Acoust. Soc. Am., vol. 124, no. 1, pp. 422-438, Jul. 2008.

[3]

T. Dau, B. Kollmeier, and A. Kohlrausch, “Modeling auditory processing of amplitude modulation. I. Detection and masking with narrow-band carriers,” J. Acoust. Soc. Am., vol. 102, no. 5, pp. 2892-2905, Nov. 1997.

[4]

P. Majdak, C. Hollomey, and R. Baumgartner, “AMT 1.x: A toolbox for reproducible research in auditory modeling,” Acta Acust., vol. 6, p. 19, 2022.

__init__(fs=44100, flow=125.0, fhigh=8000.0, n_channels=50, use_outerear=True, learnable=False, return_stages=False, filter_type='efilt', dtype=torch.float32, outerear_kwargs=None, drnl_kwargs=None, ihc_kwargs=None, adaptation_kwargs=None, modulation_kwargs=None)

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

Parameters:

• fs (float)

• flow (float)

• fhigh (float)

• n_channels (int)

• use_outerear (bool)

• learnable (bool)

• return_stages (bool)

• filter_type (str)

• dtype (dtype)

• outerear_kwargs (Dict[str, Any] | None)

• drnl_kwargs (Dict[str, Any] | None)

• ihc_kwargs (Dict[str, Any] | None)

• adaptation_kwargs (Dict[str, Any] | None)

• modulation_kwargs (Dict[str, Any] | None)

forward(x)

Process audio through the Paulick2024 model.

Parameters:

x (Tensor) – Input audio signal. Shape: (B, T) or (T,).

Returns:

If return_stages=False:

List of tensors (one per frequency channel), each shape (B, M_i, T_resample).

If return_stages=True:

Tuple of (internal_repr, stages) where stages is a dict with intermediate outputs.

Return type:

List[Tensor] | Tuple[List[Tensor], Dict[str, Tensor | List[Tensor]]]

roi_selection(internal_repr, time_window=None, channel_range=None, modulation_range=None)

Extract Region of Interest (ROI) from internal representation.

Useful for focusing decision metrics on specific: - Time windows (e.g., signal interval vs. silence) - Auditory channels (e.g., low-frequency vs. high-frequency) - Modulation channels (e.g., slow vs. fast modulations)

Return type:

List[Tensor]

Parameters:

internal_reprList[torch.Tensor]

Internal representation from forward pass (list of 50 tensors)

time_windowTuple[float, float], optional

Time window in seconds (start, end). If None, use full duration.

channel_rangeTuple[int, int], optional

Auditory channel range (start_idx, end_idx). If None, use all channels.

modulation_rangeTuple[int, int], optional

Modulation channel range (start_idx, end_idx). If None, use all modulation channels.

Returns:

: roi_repr : List[torch.Tensor]

Selected region, same format as internal_repr but with reduced dimensions

type internal_repr:

List[Tensor]

param internal_repr:

type time_window:

Tuple[float, float] | None

param time_window:

type channel_range:

Tuple[int, int] | None

param channel_range:

type modulation_range:

Tuple[int, int] | None

param modulation_range:

template_correlation(internal_repr, template, normalize=True)

Compute correlation between internal representation and a reference template.

Useful for template-matching decision strategies in psychophysical tasks (e.g., signal detection, discrimination).

Return type:

Tensor

Parameters:

internal_reprList[torch.Tensor]

Internal representation from test stimulus

templateList[torch.Tensor]

Reference template (same format as internal_repr)

normalizebool, optional

If True, compute normalized correlation (Pearson-like). Default: True

Returns:

: correlation : torch.Tensor

Correlation values with shape [B] (one value per batch item)

type internal_repr:

List[Tensor]

param internal_repr:

type template:

List[Tensor]

param template:

type normalize:

bool

param normalize:

compute_decision_variable(internal_repr, metric='rms', channel_weights=None, modulation_weights=None)

Compute decision variable from internal representation using specified metric.

Common metrics for psychophysical modeling: - ‘rms’: Root-mean-square energy - ‘mean’: Average activity - ‘max’: Maximum activity - ‘l2’: L2 norm

Return type:

Tensor

Parameters:

internal_reprList[torch.Tensor]

Internal representation from forward pass

metricstr, optional

Decision metric to compute. Default: ‘rms’

channel_weightstorch.Tensor, optional

Weights for auditory channels [50]. If None, uniform weighting.

modulation_weightstorch.Tensor, optional

Weights for modulation channels [max_Nmod]. If None, uniform weighting.

Returns:

: decision_var : torch.Tensor

Decision variable with shape [B]

type internal_repr:

List[Tensor]

param internal_repr:

type metric:

str

param metric:

type channel_weights:

Tensor | None

param channel_weights:

type modulation_weights:

Tensor | None

param modulation_weights:

make_decision(decision_var, threshold=0.0, criterion='greater')

Make binary decision based on decision variable and threshold.

Implements the final decision stage in signal detection theory.

Return type:

Tensor

Parameters:

decision_vartorch.Tensor

Decision variable from compute_decision_variable() [B]

thresholdfloat, optional

Decision threshold. Default: 0.0

criterionstr, optional

Decision criterion: ‘greater’ (DV > threshold) or ‘less’ (DV < threshold). Default: ‘greater’

Returns:

: decisions : torch.Tensor

Binary decisions with shape [B], dtype=torch.bool

type decision_var:

Tensor

param decision_var:

type threshold:

float

param threshold:

type criterion:

str

param criterion:

detection_task(signal_audio, noise_audio=None, metric='rms', threshold=None, return_decision_var=False)

Perform signal detection task (signal vs. noise).

High-level method combining forward pass, decision variable computation, and threshold-based decision.

Return type:

Tensor | Tuple[Tensor, Tensor]

Parameters:

signal_audiotorch.Tensor

Test audio (signal + noise or noise-only) [B, T]

noise_audiotorch.Tensor, optional

Reference noise-only audio for template subtraction [1 or B, T]

metricstr, optional

Decision metric. Default: ‘rms’

thresholdfloat, optional

Decision threshold. If None, return decision variable only.

return_decision_varbool, optional

If True, return (decisions, decision_var). Default: False

Returns:

: decisions : torch.Tensor (if threshold provided)

Binary decisions [B]

decision_vartorch.Tensor (if threshold is None or return_decision_var=True)

Decision variable [B]

type signal_audio:

Tensor

param signal_audio:

type noise_audio:

Tensor | None

param noise_audio:

type metric:

str

param metric:

type threshold:

float | None

param threshold:

type return_decision_var:

bool

param return_decision_var:

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

discrimination_task(stimulus1, stimulus2, metric='rms', decision_rule='greater', return_decision_var=False)

Perform discrimination task (which stimulus is louder/different?).

Return type:

Tensor | Tuple[Tensor, Tensor]

Parameters:

stimulus1torch.Tensor

First stimulus [B, T]

stimulus2torch.Tensor

Second stimulus [B, T]

metricstr, optional

Decision metric. Default: ‘rms’

decision_rulestr, optional

‘greater’: choose stimulus1 if DV1 > DV2 ‘less’: choose stimulus1 if DV1 < DV2 Default: ‘greater’

return_decision_varbool, optional

If True, return (decisions, delta_dv). Default: False

Returns:

: decisions : torch.Tensor

Binary decisions [B]: True if stimulus1 chosen, False if stimulus2 chosen

delta_dvtorch.Tensor (if return_decision_var=True)

Difference in decision variables [B]: DV1 - DV2

type stimulus1:

Tensor

param stimulus1:

type stimulus2:

Tensor

param stimulus2:

type metric:

str

param metric:

type decision_rule:

str

param decision_rule:

type return_decision_var:

bool

param return_decision_var:

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

extra_repr()

Extra representation for printing.

Returns:

String representation of module parameters.

Return type:

str