espnet2.enh.layers.dprnn.SingleRNN

source

class espnet2.enh.layers.dprnn.SingleRNN(rnn_type, input_size, hidden_size, dropout=0, bidirectional=False)

Bases: Module

Container module for a single RNN layer.

This class implements a single RNN layer that can be configured to use either RNN, LSTM, or GRU types. It includes a linear projection layer and dropout functionality.

rnn_type

The type of RNN (‘RNN’, ‘LSTM’, or ‘GRU’).

• Type: str

input_size

The dimension of the input features.

• Type: int

hidden_size

The dimension of the hidden state.

• Type: int

num_direction

The number of directions for the RNN (1 for unidirectional, 2 for bidirectional).

• Type: int

rnn

The RNN layer instance.

• Type: nn.Module

dropout

The dropout layer.

• Type: nn.Dropout

proj

The linear projection layer.

• Type: nn.Linear

• Parameters:

  • rnn_type (str) – Type of RNN to use. Must be ‘RNN’, ‘LSTM’, or ‘GRU’.
  • input_size (int) – Dimension of the input feature. The input should have shape (batch, seq_len, input_size).
  • hidden_size (int) – Dimension of the hidden state.
  • dropout (float , optional) – Dropout ratio. Default is 0.
  • bidirectional (bool , optional) – Whether the RNN layers are bidirectional. Default is False.

• Raises:AssertionError – If rnn_type is not one of ‘RNN’, ‘LSTM’, or ‘GRU’.

####### Examples

>>> rnn_layer = SingleRNN('LSTM', input_size=128, hidden_size=64)
>>> input_tensor = torch.randn(32, 10, 128)  # (batch, seq_len, input_size)
>>> output, state = rnn_layer(input_tensor)
>>> output.shape
torch.Size([32, 10, 128])  # output shape is (batch, seq_len, input_size)

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

forward(input, state=None)

Perform a forward pass through the SingleRNN layer.

This method takes the input tensor and an optional hidden state, processes the input through the RNN layer, applies dropout, and projects the output back to the input feature dimension.

• Parameters:
  • input (torch.Tensor) – Input tensor of shape (batch, seq_len, dim).
  • state (torch.Tensor , optional) – The initial hidden state of the RNN. If not provided, it defaults to None.
• Returns: A tuple containing: : - output (torch.Tensor): The output tensor after processing, with shape (batch, seq_len, input_size).
  • state (torch.Tensor): The hidden state after processing, with shape (num_layers * num_directions, batch, hidden_size).
• Return type: tuple

####### Examples

>>> rnn = SingleRNN(rnn_type='LSTM', input_size=10, hidden_size=20)
>>> input_tensor = torch.randn(5, 15, 10)  # (batch_size, seq_len, input_size)
>>> output, hidden_state = rnn(input_tensor)

NOTE

The input tensor should have shape (batch, seq_len, input_size). The hidden state should match the dimensions of the RNN layer.