PyTorch中存储神经网络前向传播状态的内存泄漏问题及疑问

我在Google Colaboratory上使用PyTorch实现一个模型，需存储多次前向传播状态用于后续损失计算与反向传播。最初将生成的张量存入Python列表，却发现该方法会引发内存泄漏——整个计算图被复制到列表中，导致GPU显存被大量占用，仅处理少量张量后就无法继续运行。我找到使用torch.utils.checkpoint的解决方案，但对此并不满意，想询问：PyTorch为何不提供存储张量引用列表的方式，而非要强制复制操作？

参考代码片段（注释append方法后显存占用大幅降低，添加break用于调试） #

""" Run 01 epoch of training
"""
def fit_epoch(self):
    epoch_training_loss = []

    for batch_idx, (normal_bag, abnormal_bag) in enumerate(self._DATALOADER):
        self._OPTIM.zero_grad()
        self._MODEL.train()

        normal_scores = []
        for idx, segment in enumerate(normal_bag):
            segment = segment.type(torch.FloatTensor).to(device)
            score = self._MODEL.forward(segment)
            # normal_scores.append(score)

            del segment, score

            # Lookout memory leaks

        abnormal_scores = []
        for idx, segment in enumerate(abnormal_bag):
            segment = segment.type(torch.FloatTensor).to(device)
            score = self._MODEL.forward(segment)
            # abnormal_scores.append(score)

            del segment, score

        break

        # Compute loss
        loss = self._LOSS_FN(normal_scores, abnormal_scores)

        # Back-propagation
        loss.backward()
        self._OPTIM.step()

        epoch_training_loss.append(loss.item())

    self._save_model()
    return epoch_training_loss

GPU显存占用情况 #

内容的提问来源于stack exchange，提问作者Khoa LT