PyTorch Gradient Accumulation 小数据集的最后一步

Question

我正在一个相对较小的数据集上训练 BERT 模型，并且不能丢失任何标记样本，因为它们必须全部用于训练。由于 GPU 内存的限制，我使用梯度累积来训练更大的批次（例如 32）。根据PyTorch documentation，梯度累加的实现方式如下：

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            loss = loss / iters_to_accumulate

        # Accumulates scaled gradients.
        scaler.scale(loss).backward()

        if (i + 1) % iters_to_accumulate == 0:
            # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)

            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

但是，如果您正在使用例如110 个训练样本，batch size 8 和累积 step 4（即有效批大小 32），这种方法只会训练前 96 个样本（即 32 x 3），即浪费 14 个样本。为了避免这种情况，我想将代码修改如下（注意更改为最终的 if 语句）：

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            loss = loss / iters_to_accumulate

        # Accumulates scaled gradients.
        scaler.scale(loss).backward()

        if (i + 1) % iters_to_accumulate == 0 or (i + 1) == len(data):
            # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)

            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

这是正确的，真的那么简单，还是会有任何副作用？这对我来说似乎很简单，但我以前从未见过它。任何帮助表示赞赏！

Answer 1

正如Lucas Ramos 已经提到的，当使用DataLoader 时，基础数据集的大小不能被批次大小整除，默认行为是具有较小的最后一批：

drop_last（bool，可选）- 如果数据集大小不能被批次大小整除，则设置为 True 以删除最后一个不完整的批次。如果 False 并且数据集的大小不能被批大小整除，则最后一批将更小。 （默认：False）

您的计划基本上是结合drop_last=False 实施梯度累积 - 即最后一批比其他所有批次都小。
因此，原则上不同批次大小的训练没有任何问题。

但是，您需要在代码中修复一些问题：
损失是小批量的平均。因此，如果您以通常的方式处理小批量，则无需担心。但是，在累积梯度时，您可以通过将损失除以 iters_to_accumulate 来明确地做到这一点：

loss = loss / iters_to_accumulate

在最后一个 minibatch（较小的大小）中，您需要更改 iter_to_accumulate 的值以反映这个较小的 minibatch 大小！

我提出了这个修改后的代码，将训练循环分为两个：小批量的外循环，以及每个小批量累积梯度的内循环。请注意使用 iter over the DataLoader 如何帮助将训练循环分成两部分：

scaler = GradScaler()

for epoch in epochs: 
    bi = 0  # index batches
    # outer loop over minibatches
    data_iter = iter(data)
    while bi < len(data):
        # determine the range for this batch
        nbi = min(len(data), bi + iters_to_accumulate)
        # inner loop over the items of the mini batch - accumulating gradients
        for i in range(bi, nbi):
            input, target = data_iter.next()
            with autocast():
                output = model(input)
                loss = loss_fn(output, target)
                loss = loss / (nbi - bi)  # divide by the true batch size

            # Accumulates scaled gradients.
            scaler.scale(loss).backward()
        # done mini batch loop - gradients were accumulated, we can make an optimizatino step.
        
        # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
        bi = nbi 

Answer 2

我很确定我以前见过这样做。查看来自 Pytorch Lightning 的 this code（函数 _accumulated_batches_reached、_num_training_batches_reached 和 should_accumulate）。