[Wav2Vec2] Improve SpecAugment function by converting numpy based function to pytorch based function
# 🚀 Feature request

As can be seen here: https://github.com/huggingface/transformers/blob/11655fafdd42eb56ad94e09ecd84d4dc2d1041ae/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L47,

the function `_compute_mask_indices` (responsible for spec augment) of Wav2Vec2 is written in numpy which means that the function is not GPU compatible. 
The function could simply be rewritten in PyTorch, which should make training on GPU faster.

This "Good First Issue" is about converting `_compute_mask_indices` to PyTorch while keeping the same functionality. 

## Your contribution

I'm happy to guide the contributor through the PR!

@patrickvonplaten  You mean it function definition will become def _compute_mask_indices(
    shape: Tuple[int, int],
    mask_prob: float,
    mask_length: int,
    attention_mask: Optional[torch.Tensor] = None,
    min_masks: int = 0,
) -> torch.tensor:

?? Of course internal working also needs to be changed according. Just trying to learn what need to be done<|||||>essentially all `np....` operations should be replaced by `torch....` operations :-) <|||||>```
def _compute_mask_indices(
    shape: Tuple[int, int],
    mask_prob: float,
    mask_length: int,
    attention_mask: Optional[torch.Tensor] = None,
    min_masks: int = 0,
) -> torch.tensor:

bsz, all_sz = shape
#mask = np.full((bsz, all_sz), False)
mask = torch.Tensor(bsz, all_sz).fill_(False)

all_num_mask = int(
        # add a random number for probabilistic rounding
        mask_prob * all_sz / float(mask_length)
        + torch.rand()
    )
all_num_mask = max(min_masks, all_num_mask)

    mask_idcs = []
    padding_mask = attention_mask.ne(1) if attention_mask is not None else None

    for i in range(bsz):
            if padding_mask is not None:
                sz = all_sz - padding_mask[i].long().sum().item()
                num_mask = int(
                    # add a random number for probabilistic rounding
                    mask_prob * sz / float(mask_length)
                    + torch.rand()
                )
                num_mask = max(min_masks, num_mask)
            else:
                sz = all_sz
                num_mask = all_num_mask

            lengths = torch.Tensor(num_mask).fill_(mask_length)

            if sum(lengths) == 0:
                lengths[0] = min(mask_length, sz - 1)

            min_len = min(lengths)
            if sz - min_len <= num_mask:
                min_len = sz - num_mask - 1

            #mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)
            mask_idc = torch.randperm(sz - min_len)[:num_mask]
            #mask_idc = np.asarray([mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(lengths[j])])
            mask_idc = torch.from_numpy(np.asarray([mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(lengths[j])]))
            #mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))
            mask_idcs.append(torch.unique(mask_idc[mask_idc < sz]))

        min_len = min([len(m) for m in mask_idcs])
        for i, mask_idc in enumerate(mask_idcs):
            if len(mask_idc) > min_len:
                #mask_idc = np.random.choice(mask_idc, min_len, replace=False)
                mask_idc = torch.randperm(mask_idc)[:min_len]
            mask[i, mask_idc] = True

        return mask

```<|||||>Can you guide me if I am doing anything wrong above ? <|||||>This looks nice already - do you want to open a PR for this? 

It would be ideal to replace the `for ....` loops with pure PyTorch vector operations as well<|||||>yes sure. Let me try. I will send PR tomorrow. Bit late to work on it for now.<|||||>@patrickvonplaten  Can you please help pass all checks. I have ran "make style", "make fixup" and did respective changes. But make quality is failing on master itself. Do you  have any suggestion for code improvement ? <|||||>@patrickvonplaten 
Hello. I am Master 1student in japan.
I want to fine tuning with local small data.
Should I fix this your code???
https://github.com/huggingface/transformers/blob/master/examples/research_projects/wav2vec2/run_asr.py

Sorry for low level question.

p.s.　Is it okay to ask such a question here?<|||||>> @patrickvonplaten
> Hello. I am Master 1student in japan.
> I want to fine tuning with local small data.
> Should I fix this your code???
> https://github.com/huggingface/transformers/blob/master/examples/research_projects/wav2vec2/run_asr.py
> 
> Sorry for low level question.
> 
> p.s.　Is it okay to ask such a question here?

I'll open-source an explicit notebook on how to fine-tune Wav2Vec2 in a ~1,2 weeks under the hugging face blog. If you haven't seen it by then, please ping me here again.<|||||>Hi @patrickvonplaten. I'd like to contribute to this issue. Is it still open?<|||||>Hey @amalad,

Yes the PR is still open since the other PR mentioned here seems to be stuck - feel free to open a new one :-)<|||||>Since Pytorch has no equivalent function to `np.random.choice`, there're only workarounds. Some [discussions](https://github.com/pytorch/pytorch/issues/16897) about this issue.

Anyway here's my take.

```python3
import torch
from typing import Optional, Tuple
import random

def _compute_specaugment_mask_indices(
    shape: Tuple[int, int],
    mask_prob: float,
    mask_length: int,
    attention_mask: Optional[torch.Tensor] = None,
    min_masks: int = 0,
) -> torch.Tensor:
    """
    Computes random mask spans for a given shape

    Args:
        shape: the the shape for which to compute masks.
            should be of size 2 where first element is batch size and 2nd is timesteps
        attention_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
        mask_length: size of the mask
        min_masks: minimum number of masked spans

    Adapted from `fairseq's data_utils.py
    <https://github.com/pytorch/fairseq/blob/e0788f7007a8473a76db573985031f3c94201e79/fairseq/data/data_utils.py#L376>`__.
    """
    bsz, all_sz = shape
    mask = torch.full((bsz, all_sz), False)

    all_num_mask = int(
        # add a random number for probabilistic rounding
        mask_prob * all_sz / float(mask_length)
        + random.random()
    )

    all_num_mask = max(min_masks, all_num_mask)
    if all_num_mask == 0:
        return mask

    mask_idcs = []
    padding_mask = attention_mask.ne(1) if attention_mask is not None else None

    for i in range(bsz):
        if padding_mask is not None:
            sz = all_sz - padding_mask[i].long().sum().item()
            num_mask = int(
                # add a random number for probabilistic rounding
                mask_prob * sz / float(mask_length)
                + random.random()
            )
            num_mask = max(min_masks, num_mask)
        else:
            sz = all_sz
            num_mask = all_num_mask

        lengths = torch.full([num_mask], mask_length)

        if sum(lengths) == 0:
            lengths[0] = min(mask_length, sz - 1)

        min_len = int(min(lengths)) if not lengths.nelement() == 0 else 0
        if sz - min_len <= num_mask:
            min_len = sz - num_mask - 1

        #mask_idc = torch.randint(sz - min_len, [num_mask]) # TODO: should sample w/o replacement
        mask_idc = random.sample(range(sz - min_len), num_mask)
        mask_idc = torch.Tensor([mask_idc[j] + offset for j in range(num_mask) for offset in range(lengths[j])])
        mask_idcs.append(torch.unique(mask_idc[mask_idc < sz]))

    min_len = min([len(m) for m in mask_idcs])
    for i, mask_idc in enumerate(mask_idcs):
        if len(mask_idc) > min_len:
            mask_idc = mask_idc.gather(dim=0, index=torch.multinomial(mask_idc, min_len, replacement=False))
        mask[i, mask_idc.long()] = True

    return mask
```<|||||>Hey @chutaklee, 

This looks nice! Could you maybe open a PR for it and measure the speed improvement when training on a GPU? :-)<|||||>Hi, @patrickvonplaten can I still work on this issue? or Is @chutaklee working on it? Actually, I have been working with fairseq's wav2vec and would like to give this issue a go.<|||||>@01-vyom Hi Vyom, I'm stuck at vectorizing the mask generation. So feel free to try it.<|||||>Feel free to give it a go @01-vyom  :-)<|||||>@patrickvonplaten Made a PR.