speechbrain.lobes.models.transformer.TransformerASR module
Transformer for ASR in the SpeechBrain style.
Authors * Jianyuan Zhong 2020
Summary
Classes:
This is a wrapper of any ASR transformer encoder. |
|
This is an implementation of transformer model for ASR. |
Reference
- class speechbrain.lobes.models.transformer.TransformerASR.TransformerASR(tgt_vocab, input_size, d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, d_ffn=2048, dropout=0.1, activation=<class 'torch.nn.modules.activation.ReLU'>, positional_encoding='fixed_abs_sine', normalize_before=False, kernel_size: int | None = 31, bias: bool | None = True, encoder_module: str | None = 'transformer', conformer_activation: ~torch.nn.modules.module.Module | None = <class 'speechbrain.nnet.activations.Swish'>, branchformer_activation: ~torch.nn.modules.module.Module | None = <class 'torch.nn.modules.activation.GELU'>, attention_type: str | None = 'regularMHA', max_length: int | None = 2500, causal: bool | None = True, csgu_linear_units: int | None = 3072, gate_activation: ~torch.nn.modules.module.Module | None = <class 'torch.nn.modules.linear.Identity'>, use_linear_after_conv: bool | None = False)[source]
Bases:
TransformerInterfaceThis is an implementation of transformer model for ASR.
The architecture is based on the paper “Attention Is All You Need”: https://arxiv.org/pdf/1706.03762.pdf
- Parameters:
tgt_vocab (int) – Size of vocabulary.
input_size (int) – Input feature size.
d_model (int, optional) – Embedding dimension size. (default=512).
nhead (int, optional) – The number of heads in the multi-head attention models (default=8).
num_encoder_layers (int, optional) – The number of sub-encoder-layers in the encoder (default=6).
num_decoder_layers (int, optional) – The number of sub-decoder-layers in the decoder (default=6).
dim_ffn (int, optional) – The dimension of the feedforward network model (default=2048).
dropout (int, optional) – The dropout value (default=0.1).
activation (torch.nn.Module, optional) – The activation function of FFN layers. Recommended: relu or gelu (default=relu).
positional_encoding (str, optional) – Type of positional encoding used. e.g. ‘fixed_abs_sine’ for fixed absolute positional encodings.
normalize_before (bool, optional) – Whether normalization should be applied before or after MHA or FFN in Transformer layers. Defaults to True as this was shown to lead to better performance and training stability.
kernel_size (int, optional) – Kernel size in convolutional layers when Conformer is used.
bias (bool, optional) – Whether to use bias in Conformer convolutional layers.
encoder_module (str, optional) – Choose between Conformer and Transformer for the encoder. The decoder is fixed to be a Transformer.
conformer_activation (torch.nn.Module, optional) – Activation module used after Conformer convolutional layers. E.g. Swish, ReLU etc. it has to be a torch Module.
branchformer_activation (torch.nn.Module, optional) – Activation module used within the Branchformer Encoder. E.g. Swish, ReLU etc. it has to be a torch Module.
attention_type (str, optional) – Type of attention layer used in all Transformer or Conformer layers. e.g. regularMHA or RelPosMHA.
max_length (int, optional) – Max length for the target and source sequence in input. Used for positional encodings.
causal (bool, optional) – Whether the encoder should be causal or not (the decoder is always causal). If causal the Conformer convolutional layer is causal.
csgu_linear_units (int, optional) – Number of neurons in the hidden linear units of the CSGU Module. -> Branchformer
gate_activation (torch.nn.Module, optional) – Activation function used at the gate of the CSGU module. -> Branchformer
use_linear_after_conv (bool, optional) – If True, will apply a linear transformation of size input_size//2. -> Branchformer
Example
>>> src = torch.rand([8, 120, 512]) >>> tgt = torch.randint(0, 720, [8, 120]) >>> net = TransformerASR( ... 720, 512, 512, 8, 1, 1, 1024, activation=torch.nn.GELU ... ) >>> enc_out, dec_out = net.forward(src, tgt) >>> enc_out.shape torch.Size([8, 120, 512]) >>> dec_out.shape torch.Size([8, 120, 512])
- forward(src, tgt, wav_len=None, pad_idx=0)[source]
- Parameters:
src (torch.Tensor) – The sequence to the encoder.
tgt (torch.Tensor) – The sequence to the decoder.
wav_len (torch.Tensor, optional) – Torch Tensor of shape (batch, ) containing the relative length to padded length for each example.
pad_idx (int, optional) – The index for <pad> token (default=0).
- make_masks(src, tgt=None, wav_len=None, pad_idx=0)[source]
This method generates the masks for training the transformer model.
- Parameters:
src (tensor) – The sequence to the encoder (required).
tgt (tensor) – The sequence to the decoder.
pad_idx (int) – The index for <pad> token (default=0).
- decode(tgt, encoder_out, enc_len=None)[source]
This method implements a decoding step for the transformer model.
- Parameters:
tgt (torch.Tensor) – The sequence to the decoder.
encoder_out (torch.Tensor) – Hidden output of the encoder.
enc_len (torch.LongTensor) – The actual length of encoder states.
- encode(src, wav_len=None, pad_idx=0)[source]
Encoder forward pass
- Parameters:
src (torch.Tensor) – The sequence to the encoder.
wav_len (torch.Tensor, optional) – Torch Tensor of shape (batch, ) containing the relative length to padded length for each example.
- class speechbrain.lobes.models.transformer.TransformerASR.EncoderWrapper(transformer, *args, **kwargs)[source]
Bases:
ModuleThis is a wrapper of any ASR transformer encoder. By default, the TransformerASR .forward() function encodes and decodes. With this wrapper the .forward() function becomes .encode() only.
Important: The TransformerASR class must contain a .encode() function.
- Parameters:
transformer (sb.lobes.models.TransformerInterface) – A Transformer instance that contains a .encode() function.
Example
>>> src = torch.rand([8, 120, 512]) >>> tgt = torch.randint(0, 720, [8, 120]) >>> net = TransformerASR( ... 720, 512, 512, 8, 1, 1, 1024, activation=torch.nn.GELU ... ) >>> encoder = EncoderWrapper(net) >>> enc_out = encoder(src) >>> enc_out.shape torch.Size([8, 120, 512])