动态batch和静态batch的原理和代码详解

原理：

在wenet中支持两种batch的方式，

第一种是常规的batch方案，但该方案当语音长短差异过大时，显存利用率低，同时带来显存oom的风险；

第二种是动态batch，不指定batch_size的大小，只限制了batch中的最大总帧数，这样就能够充分利用显存，同时不会有内存oom的风险。

论文：

https://arxiv.org/pdf/2102.01547.pdf

代码位置：

wenet/processor.py at main · wenet-e2e/wenet · GitHub

代码实现

dynamic_batch:

参数：

data: Iterable[{key,feat,label}]

max_frames_in_batch: 整个batch中能包含的总帧数不大于该值

返回：

Iterable[{key,feat,label}]

代码思路：

  步骤一：遍历data，获取每个样本sample

  步骤二：获取该样本的帧数，并更新最大帧数，然后获取padding后的总帧数

for sample in data:
    assert 'feat' in sample
    assert isinstance(sample['feat'], torch.Tensor)
    new_sample_frames = sample['feat'].size(0)
    longest_frames = max(longest_frames, new_sample_frames)
    frames_after_padding = longest_frames * (len(buf) + 1)

  步骤三：若大于batch中的最大帧，则将buf添加到迭代器中，否则一个buf没满，等待下次

if frames_after_padding > max_frames_in_batch:
    yield buf
    buf = [sample]
    longest_frames = new_sample_frames
else:
    buf.append(sample)

  步骤四：遍历结束后将剩余的buf也添加到迭代器中

static_batch:

参数：

  data：Iterable[{key,feat,label}]

  batch_size: batch size

返回：

Iterable[list{key,feat,label}]

代码思路：

  步骤一：遍历data获取每个样本sample

  步骤二：将sample取出来，当buf添加了batch_size数目样本之后，将buf传给迭代器，然后清空buf

buf = []
for sample in data:
    buf.append(sample)
    if len(buf) >= batch_size:
        yield buf
        buf = []

  步骤三：遍历结束后将剩余的buf也添加到迭代器中