speechbrain.lobes.models.wav2vec module

Components necessary to build a wav2vec 2.0 architecture following the original paper: https://arxiv.org/abs/2006.11477.

Authors * Rudolf A Braun 2022 * Guillermo Cambara 2022 * Titouan Parcollet 2022

Summary

Classes:

EncoderWrapper	A wrapper that adds positional information, masks the input and then runs the latent encoder.
W2VLatentExtractor	Convolution based feature extractor from raw audio.
W2VTargetQuantiser	Wraps nnet.quantiser.GumbelVectorQuantizer, see for documentation on arguments.

Functions:

compute_mask	This creates the boolean mask for a target shape which respects the sample lengths and will half roughly mask_prob entries set to True.
sample_negatives	Samples negatives from target tensor y.
w2v_mask_collate_fn	This creates a batch from a list of samples and also creates the boolean mask that will be used to mask the inputs of the latent encoder.

Reference

class speechbrain.lobes.models.wav2vec.W2VLatentExtractor(out_channels=[512, 512, 512, 512, 512, 512, 512], kernel_sizes=[11, 3, 3, 3, 3, 3, 3], strides=[5, 2, 2, 2, 2, 2, 2], dropout=0.0, conv_init='kaiming')

Bases: Module

Convolution based feature extractor from raw audio. Channel numbers increasing is based on https://arxiv.org/abs/2109.06870

Parameters:

• out_channels (list of ints) – Out channels of convolutional layers.

• kernel_sizes (list of ints) – Kernels of convolutional layers.

• strides (list of ints) – Strides of convolutional layers.

• dropout (float) – Dropout of CNN.

Example

>>> extractor = W2VLatentExtractor()
>>> inputs = torch.rand(10, 5000)
>>> outputs = extractor(inputs)
>>> outputs.shape
torch.Size([10, 14, 512])

forward(x, normalize_signal=True)

Calculates latents from audio input.

get_output_lengths(input_lengths: LongTensor)

Calculates output lengths for given input lengths.

training: bool

class speechbrain.lobes.models.wav2vec.W2VTargetQuantiser(in_dim=512, out_dim=256, quantiser=<class 'speechbrain.nnet.quantisers.GumbelVectorQuantizer'>, num_vars=320, temperature_decay=(2.0, 0.25, 0.999995))

Bases: Module

Wraps nnet.quantiser.GumbelVectorQuantizer, see for documentation on arguments.

Example

>>> quantiser = W2VTargetQuantiser()
>>> inputs = torch.rand(10, 12, 512)
>>> output, meta = quantiser(inputs)
>>> output.shape
torch.Size([10, 12, 256])

forward(x)

Returns quantised targets plus meta information.

training: bool

class speechbrain.lobes.models.wav2vec.EncoderWrapper(in_dim, embedding_dim, latent_encoder, positional_encoding=<class 'speechbrain.lobes.models.transformer.Transformer.PositionalEncoding'>, dropout_encoder_input=0.05)

Bases: Module

A wrapper that adds positional information, masks the input and then runs the latent encoder. :param in_dim: Last dimension of input tensor. :type in_dim: int :param embedding_dim: Dimension to project input to and that the latent encoder will use. :type embedding_dim: int :param latent_encoder: Initialized latent encoder object. :type latent_encoder: torch.nn.module :param positional_encoding: Uninitialized nn.module for adding positional information, will use embedding_dim. :type positional_encoding: torch.nn.module :param dropout_encoder_input: Dropout on encoder input. :type dropout_encoder_input: float

Example

>>> from speechbrain.lobes.models.transformer.Transformer import TransformerEncoder
>>> encoder = TransformerEncoder(d_model=768, num_layers=4, nhead=4, d_ffn=1024)
>>> wrapper = EncoderWrapper(1024, 768, encoder)
>>> inputs = torch.rand(10, 12, 1024)
>>> outputs = wrapper(inputs)
>>> outputs["embeddings"].shape
torch.Size([10, 12, 768])

forward(latents, wav_lens=None, padding_mask=None, mask=None)

Parameters:

• latents (torch.Tensor, shape (B, T, C)) – Batch of latent representations (AKA frames) output from latent extractor.

• wav_lens (torch.Tensor, shape (B,)) – The actual (unpadded) relative lengths for each sample of the batch (0<wav_lens<1).

• padding_mask (Torch.Tensor, shape (B, T,)) – Can be provided instead of wav_lens.

• mask (torch.Tensor, shape (B, T)) – Boolean mask which decides which latent frames will be masked.

training: bool

speechbrain.lobes.models.wav2vec.compute_mask(shape, sample_lens, mask_prob, mask_length)

This creates the boolean mask for a target shape which respects the sample lengths and will half roughly mask_prob entries set to True.

Parameters:

• shape (list of ints, like (N, M)) – Shape of boolean mask to return.

• sample_lens (list of ints) – Absolute lengths of per sample lengths.

• mask_prob (float) – Percentage to mask.

• mask_length (int) – Length of contiguous subsequence to mask.

Returns:

mask – Boolean mask with shape of input argument shape.

Return type:

numpy.ndarray

speechbrain.lobes.models.wav2vec.sample_negatives(y, num_neg)

Samples negatives from target tensor y. :param y: Tensor of shape (B, T, C) :type y: torch.Tensor :param num_neg: Number of negatives to sample. :type num_neg: int

Returns:

negs – Negatives in shape (N, B, T, C)

Return type:

torch.Tensor

speechbrain.lobes.models.wav2vec.w2v_mask_collate_fn(samples_lst, get_out_len_fn, mask_prob, mask_length)

This creates a batch from a list of samples and also creates the boolean mask that will be used to mask the inputs of the latent encoder. To create the mask we need to know the output shape after the latent extractor, therefore the argument get_out_len_fn. One could also create masks per sample (when loading the audio file) and then collate them but at that time one doesn’t know the length of the shortest sample in the batch (which determines the number of masked frames) so it’s better this way.

Parameters:

• samples_lst (list) – List of samples returned by the audio_pipeline.

• get_out_len_fn (function) – Function that calculates length of sample after it passes through feature extractor.

• mask_prob (float) – Approximate percentage of frames to mask.

• mask_length (int) – Number of contiguous frames that will be masked.

Returns:

• wavs_padded (torch.Tensor, shape (B, T)) – Audio arrays with right-sided padding.

• wav_lens (torch.Tensor, shape (B,)) – For each sample the percentage of the array that is not padding.

• mask (torch.Tensor, shape (B, T)) – Boolean mask to mask frames.