RNN 中的梯度累积

Question

我在运行大型 RNN 网络时遇到了一些内存问题 (GPU)，但我想保持我的批量大小合理，所以我想尝试梯度累积。在一个一次性预测输出的网络中，这似乎是不言而喻的，但在 RNN 中，您为每个输入步骤执行多个前向传递。因此，我担心我的实现无法按预期工作。我从用户 albanD 的优秀示例 here 开始，但我认为在使用 RNN 时应该对其进行修改。我认为这是因为你积累了更多的梯度，因为你在每个序列中进行了多次转发。

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

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已经积累了这么多信息，但我不确定。

Answer 1

为简单起见，我只会处理amp启用梯度累积，没有amp的想法是一样的。你提出的步骤在amp 下运行，所以让我们坚持下去。

step

在PyTorch documentation about amp 你有一个梯度累积的例子。你应该在step 里面做。每次运行loss.backward() 时，梯度都会在张量叶中累积，可以通过optimizer 进行优化。因此，您的 step 应该如下所示（参见 cmets）：

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

正如你 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，不确定你在这里追求什么。如果是这种情况，那么您就无法进行 AFAIK 了，因为 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。