pytorch 中 利用自定义函数 get_mask_from_lengths(lengths, max_len)获取每个batch的mask

在pytorch中，经常会需要通过batch进行批量处理数据，由于每个batch中各个样本之间存在差异，经常会需要进行先padding后mask的操作。

尤其是在自然语言处理任务中，每个batch中的每个句子是不等长的。一般都是先通过填充0的方式将每个batch中每一句padding成和最长的句子等长的形式；再模型中或者计算loss的时候，再将padding成0的部分mask掉，从而避免padding带来的影响。

所以，每个batch在读取数据时，需要保存下该batch中每个句子的长度（以及最长的句子长度），也就是 lengths 和 max_len；通过保存的句子长度就可以得到每个batch对应的mask，然后就可以对数据输出进行mask操作。

那如何通过batch中每个句子的长度lengths 和最长的句子长度 max_len 获得mask呢？可以通过如下的自定义函数 get_mask_from_lengths(lengths, max_len) 获得。具体如下：

import torch

def get_mask_from_lengths(lengths, max_len=None):
    '''
    param:
        lengths --- [Batch_size]
    return:
        mask --- [Batch_size, max_len]
    '''
    batch_size = lengths.shape[0]  
    if max_len is None:
        max_len = torch.max(lengths).item()  

    ids = torch.arange(0, max_len).unsqueeze(0).expand(batch_size, -1)
    # ids = torch.arange(0, max_len).unsqueeze(0).expand(batch_size, -1).cuda()  ## 实际需要注意device
    mask = ids >= lengths.unsqueeze(1).expand(-1, max_len)    ## True 或 False

    return mask

## 实例
## batch_size = 4 , 每句话的长度分别是 [2, 4, 3, 2]
lengths = torch.tensor([2,4,3,2])
mask = get_mask_from_lengths(lengths)
print(mask)

tensor([[False, False,  True,  True],
        [False, False, False, False],
        [False, False, False,  True],
        [False, False,  True,  True]])

实际应用场景：

比如，语音合成时，每个batch中每个句子的音素长度，每个batch中每个句子的梅尔谱的长度等等。