RNN 中的梯度累积

在运行大型 RNN 网络时，我遇到了一些内存问题（GPU），但我想保持我的批量大小合理，所以我想尝试梯度累积。在一次性预测输出的网络中，这似乎是不言而喻的，但在 RNN 中，您为每个输入步骤执行多次前向传递。因此，我担心我的实施无法按预期进行。我从用户 albanD 的优秀示例开始here https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/20，但我认为在使用 RNN 时应该修改它们。我认为这是因为你积累了更多的梯度，因为你对每个序列进行了多次前向。

我当前的实现如下所示，同时允许在 PyTorch 1.6 中使用 AMP，这似乎很重要 - 一切都需要在正确的位置调用。请注意，这只是一个抽象版本，看起来可能有很多代码，但主要是注释。

def train(epochs):
    """Main training loop. Loops for `epoch` number of epochs. Calls `process`."""
    for epoch in range(1, epochs + 1):
        train_loss = process("train")
        valid_loss = process("valid")
        # ... check whether we improved over earlier epochs
        if lr_scheduler:
            lr_scheduler.step(valid_loss)
        
def process(do):
    """Do a single epoch run through the dataloader of the training or validation set. 
       Also takes care of optimizing the model after every `gradient_accumulation_steps` steps.
       Calls `step` for each batch where it gets the loss from."""
    if do == "train":
        model.train()
        torch.set_grad_enabled(True)
    else:
        model.eval()
        torch.set_grad_enabled(False)
    
    loss = 0.
    for batch_idx, batch in enumerate(dataloaders[do]):
        step_loss, avg_step_loss = step(batch)
        loss += avg_step_loss

        if do == "train":
            if amp:
                scaler.scale(step_loss).backward()

                if (batch_idx + 1) % gradient_accumulation_steps == 0:
                    # Unscales the gradients of optimizer's assigned params in-place
                    scaler.unscale_(optimizer)
                    # clip in-place
                    clip_grad_norm_(model.parameters(), 2.0)
                    scaler.step(optimizer)
                    scaler.update()
                    model.zero_grad()
            else:
                step_loss.backward()
                if (batch_idx + 1) % gradient_accumulation_steps == 0:
                    clip_grad_norm_(model.parameters(), 2.0)
                    optimizer.step()
                    model.zero_grad()
        
        # return average loss
        return loss / len(dataloaders[do])

    def step():
        """Processes one step (one batch) by forwarding multiple times to get a final prediction for a given sequence."""
        # do stuff... init hidden state and first input etc.
        loss = torch.tensor([0.]).to(device)
        
        for i in range(target_len):
            with torch.cuda.amp.autocast(enabled=amp):
                # overwrite previous decoder_hidden
                output, decoder_hidden = model(decoder_input, decoder_hidden)

                # compute loss between predicted classes (bs x classes) and correct classes for _this word_
                item_loss = criterion(output, target_tensor[i])

                # We calculate the gradients for the average step so that when
                # we do take an optimizer.step, it takes into account the mean step_loss
                # across batches. So basically (A+B+C)/3 = A/3 + B/3 + C/3
                loss += (item_loss / gradient_accumulation_steps)

            topv, topi = output.topk(1)
            decoder_input = topi.detach()
        
        return loss, loss.item() / target_len

上面的内容似乎没有像我希望的那样工作，即它仍然很快就会遇到内存不足的问题。我认为原因是step已经积累了这么多信息，但我不确定。

为了简单起见，我只关心amp启用梯度积累，没有放大器，想法是一样的。你提出的步骤在下面运行amp所以让我们坚持这一点。

step

In 关于 amp 的 PyTorch 文档 https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation你有一个梯度累积的例子。你应该在里面做step。每次你跑步的时候loss.backward()梯度在张量叶子内部累积，可以通过以下方式优化optimizer。因此，你的step应该看起来像这样（见评论）：

def step():
    """Processes one step (one batch) by forwarding multiple times to get a final prediction for a given sequence."""
    # You should not accumulate loss on `GPU`, RAM and CPU is better for that
    # Use GPU only for calculations, not for gathering metrics etc.
    loss = 0

    for i in range(target_len):
        with torch.cuda.amp.autocast(enabled=amp):
            # where decoder_input is from?
            # I assume there is one in real code
            output, decoder_hidden = model(decoder_input, decoder_hidden)
            # Here you divide by accumulation steps
            item_loss = criterion(output, target_tensor[i]) / (
                gradient_accumulation_steps * target_len
            )


        scaler.scale(item_loss).backward()
        loss += item_loss.detach().item()

        # Not sure what was topv for here
        _, topi = output.topk(1)
        decoder_input = topi.detach()

    # No need to return loss now as we did backward above
    return loss / target_len

As you detach decoder_input无论如何（所以它就像没有历史记录的全新隐藏输入，并且参数将基于此进行优化，不基于所有运行）没有必要backward进行中。另外，你可能不需要decoder_hidden，如果没有传递到网络，torch.tensor用零填充是隐式传递的。

我们还应该除以gradient_accumulation_steps * target_len因为这就是多少backward我们将在单个优化步骤之前运行。

由于你的一些变量定义不明确，我假设你只是对正在发生的事情制定了一个计划。

另外，如果你想保留历史记录，你不应该detach decoder_input，在这种情况下它看起来像这样：

def step():
    """Processes one step (one batch) by forwarding multiple times to get a final prediction for a given sequence."""
    loss = 0

    for i in range(target_len):
        with torch.cuda.amp.autocast(enabled=amp):
            output, decoder_hidden = model(decoder_input, decoder_hidden)
            item_loss = criterion(output, target_tensor[i]) / (
                gradient_accumulation_steps * target_len
            )

        _, topi = output.topk(1)
        decoder_input = topi

        loss += item_loss
    scaler.scale(loss).backward()
    return loss.detach().cpu() / target_len

这实际上会通过 RNN 多次，并且可能会引发 OOM，不知道你在这里追求什么。如果是这种情况，据我所知，您无能为力，因为 RNN 计算太长，无法适应 GPU。

process

仅提供了该代码的相关部分，因此它将是：

loss = 0.0
for batch_idx, batch in enumerate(dataloaders[do]):
    # Here everything is detached from graph so we're safe
    avg_step_loss = step(batch)
    loss += avg_step_loss

    if do == "train":
        if (batch_idx + 1) % gradient_accumulation_steps == 0:
            # You can use unscale as in the example in PyTorch's docs
            # just like you did
            scaler.unscale_(optimizer)
            # clip in-place
            clip_grad_norm_(model.parameters(), 2.0)
            scaler.step(optimizer)
            scaler.update()
            # IMO in this case optimizer.zero_grad is more readable
            # but it's a nitpicking
            optimizer.zero_grad()

# return average loss
return loss / len(dataloaders[do])

疑问式

[...] 在 RNN 中，您为每个输入步骤执行多次前向传递。 因此，我担心我的实施不会像 故意的。

不要紧。对于每一次前进，您通常应该向后执行一次（这里似乎就是这种情况，请参阅步骤以了解可能的选项）。之后我们（通常）不需要将损失连接到图表正如我们已经执行过的backpropagation，得到我们的梯度并准备优化参数。

这种损失需要有历史记录，因为它会回到流程循环 哪里会向后调用它。

无需打电话backward正在进行中。