Skip to content

Incorrect transformer mask size #2344

Description

@egaznep

Describe the bug

PyTorch multi_head_attention code enforces that data size and mask size matches.

https://github.com/pytorch/pytorch/blob/df4e3d9d08f3d5d5439c3626be4bf29659488cdf/torch/nn/functional.py#L5442-L5444

However, speechbrain code generates masks according to the longest wav_len, which can be shorter than the batch size with padding, resulting in the exception given in the codeblock above.

src_key_padding_mask = ~length_to_mask(abs_len).bool()

Solution:

The function length_to_mask accepts an optional argument max_len, this could be used. Should I open a PR?

def length_to_mask(length, max_len=None, dtype=None, device=None):
"""Creates a binary mask for each sequence.
Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3
Arguments
---------
length : torch.LongTensor
Containing the length of each sequence in the batch. Must be 1D.
max_len : int
Max length for the mask, also the size of the second dimension.
dtype : torch.dtype, default: None
The dtype of the generated mask.
device: torch.device, default: None
The device to put the mask variable.

Expected behaviour

The mask should have been generated according to what PyTorch anticipates. Instead, an exception is thrown.

To Reproduce

No response

Environment Details

speechbrain==0.5.16  (checking the develop branch this still seem to be an issue)
pytorch==2.1.1

Relevant Log Output

No response

Additional Context

No response

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions