espnet2.asr.transducer.rnnt_multi_blank.utils.rnnt_helper.get_workspace_size

source

espnet2.asr.transducer.rnnt_multi_blank.utils.rnnt_helper.get_workspace_size(maxT: int, maxU: int, minibatch: int, gpu: bool) → Tuple[int | None, RNNTStatus]

Calculate the required workspace size for the RNNT (Recurrent Neural Network Transducer) model based on input parameters.

This function computes the amount of memory needed for the RNNT model to perform forward and backward passes during training or inference, considering both CPU and GPU execution environments. It takes into account the maximum sequence lengths, the number of tokens, and the size of the minibatch.

espnet2.asr.transducer.rnnt_multi_blank.utils.rnnt_helper.maxT

The maximum length of the input sequences.

• Type: int

espnet2.asr.transducer.rnnt_multi_blank.utils.rnnt_helper.maxU

The maximum number of tokens (including blanks).

• Type: int

espnet2.asr.transducer.rnnt_multi_blank.utils.rnnt_helper.minibatch

The number of sequences processed in parallel.

• Type: int

espnet2.asr.transducer.rnnt_multi_blank.utils.rnnt_helper.gpu

A flag indicating whether to calculate for GPU or CPU.

• Type: bool

• Parameters:

  • maxT (int) – Maximum time steps in the input sequences.
  • maxU (int) – Maximum number of unique labels (including blanks).
  • minibatch (int) – Number of sequences processed in a single pass.
  • gpu (bool) – Flag indicating if the calculations should be done for GPU.

• Returns: A tuple containing the computed workspace size in bytes (or None if invalid) and the status of the operation.

• Return type: Tuple[Optional[int], global_constants.RNNTStatus]

• Raises:

  • ValueError – If any of the input parameters (minibatch, maxT, maxU) are
  • less than or equal to zero. –

Examples

>>> size, status = get_workspace_size(100, 50, 32, True)
>>> print(size)  # Expected output: Size in bytes required for workspace
>>> print(status)  # Expected output: RNNT_STATUS_SUCCESS

NOTE

The calculated workspace size is essential for memory management in deep learning applications to avoid runtime errors related to memory allocation.