Repository: pytorch/audio

Issue:
[Doc] Explain the effect of `length` in `Wav2Vec2Model`
### 🚀 The feature

More specific explanation in the docs.
in [here](https://github.com/pytorch/audio/blob/e885204ea0d343b8bddd7e829b9e7b29e7c5a9bb/torchaudio/models/wav2vec2/model.py#L83-L85), I need more in-detailed explanation about ``length``. Is it a sample rate?

### Motivation, pitch

It's confusing to understand what the variable/argument behaves.

### Alternatives

_No response_

### Additional context

_No response_


Discussion:
Hi @hihunjin 

When batching multiple audios with different duration, the resulting Tensor would have padding for shorter audios.
`length` parameter is to indicate what is the valid (unpadded) length of each batch sample.

Say I create a batch from 1 second audio and 0.8 second audio, both single channel and sampled at 16k Hz. The resulting batch Tensor will be in shape of `[2, 16000]`.
The second audio in the batch,  actually has `12800 ( == 16000 * 0.8)` valid samples and the rest 3200 samples are just padding. In this case the input `length`  Tensor should look like `torch.Tensor([16000, 12800])`.

By providing the length Tensor, `Wav2Vec2Model` will compute the appropriate mask when it go through transformer layer, so that artifacts from the padding portion will not affect the computation.

`length` parameter also helps to get the same sort of information for the output values. Since the `Wav2Vec2Model` changes the number of frames, it is not intuitively clear what are the valid output lengths just by looking at the shape of the output Tensor. When `length` is provided, it will compute the valid output lengths and return it.


Length computation in convolution layer

https://github.com/pytorch/audio/blob/e885204ea0d343b8bddd7e829b9e7b29e7c5a9bb/torchaudio/models/wav2vec2/components.py#L62-L65

Mask computation in Transformer layer

https://github.com/pytorch/audio/blob/e885204ea0d343b8bddd7e829b9e7b29e7c5a9bb/torchaudio/models/wav2vec2/components.py#L442-L449
Thanks a lot. I appreciate it.
Glad to help.
I will keep this open until we update the doc to include the above information. Thanks for your feedback.

Resolving patch:
diff --git a/torchaudio/models/tacotron2.py b/torchaudio/models/tacotron2.py
--- a/torchaudio/models/tacotron2.py
+++ b/torchaudio/models/tacotron2.py
@@ -1080,7 +1080,7 @@ def infer(self, tokens: Tensor, lengths: Optional[Tensor] = None) -> Tuple[Tenso
                 If ``None``, it is assumed that the all the tokens are valid. Default: ``None``
 
         Returns:
-            Tensor, Tensor, and Tensor:
+            (Tensor, Tensor, Tensor):
                 Tensor
                     The predicted mel spectrogram with shape `(n_batch, n_mels, max of mel_specgram_lengths)`.
                 Tensor
diff --git a/torchaudio/models/wav2vec2/model.py b/torchaudio/models/wav2vec2/model.py
--- a/torchaudio/models/wav2vec2/model.py
+++ b/torchaudio/models/wav2vec2/model.py
@@ -50,8 +50,14 @@ def extract_features(
         Args:
             waveforms (Tensor): Audio tensor of shape `(batch, frames)`.
             lengths (Tensor or None, optional):
-                Indicates the valid length of each audio sample in the batch.
+                Indicates the valid length of each audio in the batch.
                 Shape: `(batch, )`.
+                When the ``waveforms`` contains audios with different durations,
+                by providing ``lengths`` argument, the model will compute
+                the corresponding valid output lengths and apply proper mask in
+                transformer attention layer.
+                If ``None``, it is assumed that the entire audio waveform
+                length is valid.
             num_layers (int or None, optional):
                 If given, limit the number of intermediate layers to go through.
                 Providing `1` will stop the computation after going through one
@@ -59,13 +65,14 @@ def extract_features(
                 intermediate layers are returned.
 
         Returns:
-            List of Tensors and an optional Tensor:
+            (List[Tensor], Optional[Tensor]):
             List of Tensors
                 Features from requested layers.
-                Each Tensor is of shape: `(batch, frames, feature dimention)`
+                Each Tensor is of shape: `(batch, time frame, feature dimension)`
             Tensor or None
                 If ``lengths`` argument was provided, a Tensor of shape `(batch, )`
-                is retuned. It indicates the valid length of each feature in the batch.
+                is returned.
+                It indicates the valid length in time axis of each feature Tensor.
         """
         x, lengths = self.feature_extractor(waveforms, lengths)
         x = self.encoder.extract_features(x, lengths, num_layers)
@@ -81,17 +88,24 @@ def forward(
         Args:
             waveforms (Tensor): Audio tensor of shape `(batch, frames)`.
             lengths (Tensor or None, optional):
-                Indicates the valid length of each audio sample in the batch.
+                Indicates the valid length of each audio in the batch.
                 Shape: `(batch, )`.
+                When the ``waveforms`` contains audios with different duration,
+                by providing ``lengths`` argument, the model will compute
+                the corresponding valid output lengths and apply proper mask in
+                transformer attention layer.
+                If ``None``, it is assumed that all the audio in ``waveforms``
+                have valid length. Default: ``None``.
 
         Returns:
-            Tensor and an optional Tensor:
+            (Tensor, Optional[Tensor]):
             Tensor
                 The sequences of probability distribution (in logit) over labels.
                 Shape: `(batch, frames, num labels)`.
             Tensor or None
                 If ``lengths`` argument was provided, a Tensor of shape `(batch, )`
-                is retuned. It indicates the valid length of each feature in the batch.
+                is retuned.
+                It indicates the valid length in time axis of the output Tensor.
         """
         x, lengths = self.feature_extractor(waveforms, lengths)
         x = self.encoder(x, lengths)
diff --git a/torchaudio/models/wavernn.py b/torchaudio/models/wavernn.py
--- a/torchaudio/models/wavernn.py
+++ b/torchaudio/models/wavernn.py
@@ -341,16 +341,23 @@ def infer(self, specgram: Tensor, lengths: Optional[Tensor] = None) -> Tuple[Ten
             specgram (Tensor):
                 Batch of spectrograms. Shape: `(n_batch, n_freq, n_time)`.
             lengths (Tensor or None, optional):
-                Indicates the valid length in of each spectrogram in time axis.
-                Shape: `(n_batch, )`.
+                Indicates the valid length of each audio in the batch.
+                Shape: `(batch, )`.
+                When the ``specgram`` contains spectrograms with different duration,
+                by providing ``lengths`` argument, the model will compute
+                the corresponding valid output lengths.
+                If ``None``, it is assumed that all the audio in ``waveforms``
+                have valid length. Default: ``None``.
 
         Returns:
-            Tensor and optional Tensor:
+            (Tensor, Optional[Tensor]):
             Tensor
                 The inferred waveform of size `(n_batch, 1, n_time)`.
                 1 stands for a single channel.
             Tensor or None
-                The valid lengths of each waveform in the batch. Size `(n_batch, )`.
+                If ``lengths`` argument was provided, a Tensor of shape `(batch, )`
+                is retuned.
+                It indicates the valid length in time axis of the output Tensor.
         """
 
         device = specgram.device


Test patch: