Skip to main content

espnet2.asr.state_spaces.block.SequenceResidualBlock

About 5 min

espnet2.asr.state_spaces.block.SequenceResidualBlock

source

class espnet2.asr.state_spaces.block.SequenceResidualBlock(d_input, i_layer=None, prenorm=True, dropout=0.0, tie_dropout=False, transposed=False, layer=None, residual=None, norm=None, pool=None, drop_path=0.0)

Bases: SequenceModule

Residual block wrapper for a black box layer.

The SequenceResidualBlock class implements a generic transformation from (batch, length, d_input) to (batch, length, d_input). It provides configurable options for normalization, subsampling, and residual connections, making it versatile for various neural network architectures.

d_input

Input feature dimension.

  • Type: int

i_layer

Layer index, required for certain residuals.

  • Type: int, optional

prenorm

If True, applies normalization before the black box layer.

  • Type: bool

dropout

Dropout rate for the black box module.

  • Type: float

tie_dropout

If True, ties the dropout mask across the sequence.

  • Type: bool

transposed

If True, transposes inputs for the black box layer.

  • Type: bool

layer

Instance of the black box module.

  • Type: nn.Module

residual

Instance of the residual function.

  • Type: nn.Module, optional

norm

Instance of the normalization layer.

pool

Instance of the pooling layer.

  • Type: nn.Module, optional

drop_path

Instance for stochastic depth.

  • Type:StochasticDepth

  • Parameters:

    • d_input (int) – Input feature dimension.
    • i_layer (int , optional) – Layer index, for certain residuals.
    • prenorm (bool) – Apply normalization before the black box layer.
    • dropout (float) – Dropout rate for the black box module.
    • tie_dropout (bool) – Tie dropout mask across the sequence.
    • transposed (bool) – Transpose inputs for each layer.
    • layer (dict) – Config for the black box module.
    • residual (dict , optional) – Config for the residual function.
    • norm (dict , optional) – Config for the normalization layer.
    • pool (dict , optional) – Config for pooling layer per stage.
    • drop_path (float) – Drop ratio for stochastic depth.

  • Returns: Output tensor and updated state.

  • Return type: Tuple[Tensor, Any]

################# Examples

>>> block = SequenceResidualBlock(d_input=128, layer={'type': 'Conv1d'})
>>> output, state = block(input_tensor, state)

########## NOTE Ensure that the input tensor shape matches (batch, length, d_input).

  • Raises:ValueError – If the input tensor dimensions do not match expected shape.

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

property d_output

Residual block wrapper for black box layer.

This class implements a residual block that wraps around a black box layer, allowing for a flexible and configurable transformation of input sequences. The block supports various configurations for normalization, pooling, and residual connections, enabling advanced architectures in sequence modeling.

d_input

Input feature dimension.

  • Type: int

i_layer

Layer index for certain residuals like Decay.

  • Type: int, optional

prenorm

Indicates if normalization should be applied before the black box layer.

  • Type: bool

dropout

Dropout rate for the black box module.

  • Type: float

tie_dropout

If True, ties dropout mask across sequence.

  • Type: bool

transposed

If True, transposes inputs to (batch, dim, length).

  • Type: bool

layer

Configured black box module.

  • Type: nn.Module

residual

Configured residual function.

  • Type: nn.Module, optional

norm

Normalization layer.

pool

Pooling layer per stage.

  • Type: Pooling, optional

drop_path

Stochastic depth for dropout.

  • Type:StochasticDepth, optional

  • Parameters:

    • d_input (int) – Input feature dimension.
    • i_layer (int , optional) – Layer index for specific residual configurations.
    • prenorm (bool) – If True, applies normalization before the black box layer.
    • dropout (float) – Dropout probability for the black box module.
    • tie_dropout (bool) – If True, ties the dropout mask across sequences.
    • transposed (bool) – If True, transposes inputs for the layer.
    • layer (dict , optional) – Configuration for the black box module.
    • residual (dict , optional) – Configuration for the residual function.
    • norm (dict , optional) – Configuration for normalization.
    • pool (dict , optional) – Configuration for pooling layer.
    • drop_path (float) – Drop ratio for stochastic depth.

  • Returns: A tuple containing the output tensor and the updated state.

  • Return type: tuple

################# Examples

>>> block = SequenceResidualBlock(d_input=128, dropout=0.1)
>>> input_tensor = torch.randn(32, 10, 128)  # (batch, length, features)
>>> output, state = block(input_tensor)

########## NOTE

  • The block can be configured to use various normalization and residual options.
  • Ensure that the dimensions of the input tensor match d_input.
  • Raises:ValueError – If the provided configurations are incompatible.

property d_state

Residual state dimension property for the SequenceResidualBlock class.

This property retrieves the state dimension from the underlying black box layer within the SequenceResidualBlock. The state dimension is typically utilized for maintaining the hidden states across time steps in sequential models.

  • Returns: The dimension of the state from the black box layer.
  • Return type: int

################# Examples

>>> block = SequenceResidualBlock(d_input=128, layer={'type': 'some_layer'})
>>> state_dim = block.d_state
>>> print(state_dim)
128  # Assuming the layer's output state dimension is 128

########## NOTE This property assumes that the underlying layer has a defined d_state attribute. If the layer does not define d_state, an AttributeError will be raised.

default_state(*args, **kwargs)

Return the default state for the black box layer.

This method serves as a wrapper to retrieve the default state from the underlying black box layer. It can be useful for initializing the state before processing input data through the forward pass.

  • Parameters:
    • *args – Variable length argument list to be passed to the underlying layer’s default_state method.
    • **kwargs – Arbitrary keyword arguments to be passed to the underlying layer’s default_state method.
  • Returns: The default state of the black box layer.

################# Examples

>>> block = SequenceResidualBlock(d_input=128)
>>> state = block.default_state()
>>> print(state)  # Prints the initialized state from the layer

########## NOTE Ensure that the black box layer is properly initialized before calling this method, as it relies on the layer’s configuration.

forward(x, state=None, **kwargs)

Performs a forward pass through the sequence residual block.

This method applies the defined transformations to the input tensor x. It processes the input through normalization (if configured), a black box layer, optional residual connections, and pooling (if configured).

  • Parameters:
    • x (torch.Tensor) – Input tensor of shape (batch, length, d_input).
    • state (optional) – State information for the black box layer.
    • **kwargs – Additional keyword arguments passed to the black box layer.
  • Returns: A tuple containing: : - y (torch.Tensor): Output tensor after transformations. Shape is (batch, length, d_output).
    • state: Updated state information from the black box layer.
  • Return type: tuple

################# Examples

>>> block = SequenceResidualBlock(d_input=128)
>>> input_tensor = torch.randn(32, 10, 128)  # (batch, length, d_input)
>>> output, new_state = block.forward(input_tensor)

########## NOTE The method applies pre-norm and post-norm based on the configuration. If prenorm is set to True, normalization is applied before the black box layer; otherwise, it is applied afterward.

  • Raises:ValueError – If the input tensor x does not have the expected shape.

property state_to_tensor

Converts the internal state of the layer to a tensor format.

This property provides a tensor representation of the current state of the layer, which can be useful for various operations, such as logging, visualization, or further processing. The specific format of the tensor is determined by the underlying black box layer’s implementation.

  • Returns: A tensor representation of the layer’s state.
  • Return type: torch.Tensor

################# Examples

Assuming block is an instance of SequenceResidualBlock

state_tensor = block.state_to_tensor print(state_tensor.shape) # Output will depend on the layer’s state

step(x, state, **kwargs)

Performs a single step of the residual block transformation.

This method applies the transformation for a single input sample in the context of the SequenceResidualBlock. It includes the application of normalization, a black box layer, residual connections, and pooling.

  • Parameters:
    • x (torch.Tensor) – Input tensor of shape (batch, length, d_input).
    • state (Any) – State information from previous layers, which may be required for the black box layer.
    • **kwargs – Additional keyword arguments to be passed to the black box layer.
  • Returns: A tuple containing: : - y (torch.Tensor): The transformed output tensor of shape (batch, length, d_output).
    • state (Any): Updated state information after processing the input.
  • Return type: tuple

################# Examples

>>> block = SequenceResidualBlock(d_input=128, layer=my_layer_config)
>>> x = torch.randn(32, 10, 128)  # (batch_size, seq_length, d_input)
>>> initial_state = block.default_state()
>>> output, updated_state = block.step(x, initial_state)

########## NOTE Ensure that the dimensions of the input tensor match the expected input feature dimension d_input defined during initialization.

  • Raises:
    • ValueError – If the input tensor x does not have the expected
    • shape or if the state is incompatible with the layer's requirements. –
Copyright © 2024 ESPnet Community. All rights reserved.