深度学习框架中存储中间值用于反向传播的优化实践

1次阅读

没有评论

共计 2183 个字符，预计需要花费 6 分钟才能阅读完成。

在深度学习模型的训练过程中，反向传播算法需要依赖前向传播计算得到的中间值来计算梯度。这些中间值通常需要保存在内存中，以便在反向传播时使用。随着模型规模的增大和训练数据的增多，这些中间值占用的内存空间迅速膨胀，成为训练过程中的主要瓶颈之一。

大模型训练中，中间值可能占用数十 GB 甚至上百 GB 的内存
内存不足会导致训练过程中断或效率大幅下降
传统方法简单存储所有中间值，造成大量内存浪费

动态计算是指在反向传播时重新计算所需的中间值，而不是存储它们。这种策略可以显著减少内存使用，但会增加计算开销。

优点：内存占用极低
缺点：计算时间可能增加 50%-200%
适用场景：计算量小但内存需求大的操作

选择性存储是指只保存那些计算成本高或重计算困难的中间值，其余的在反向传播时重新计算。

优点：内存和计算的较好平衡
缺点：需要仔细分析每个操作的计算 / 内存特性
适用场景：大多数中等规模的模型

梯度检查点技术将计算图分成若干段，只保存段边界处的激活值，段内的激活值在反向传播时重新计算。

优点：可以将内存使用降低到 O(√n)
缺点：实现复杂，需要修改计算图
适用场景：超大模型训练

我们实现了一个基于 PyTorch 的选择性存储方案，其核心思想是：

分析模型中各层的计算特性和内存需求
为每个操作定义存储策略
实现自定义的 Function 类控制存储行为

关键数据结构：

class StoragePolicy:
    def __init__(self):
        self.ops_policy = {
            'conv2d': 'store',
            'relu': 'recompute',
            'batchnorm': 'store',
            # 其他操作策略...
        }

import torch
import torch.nn as nn
import torch.nn.functional as F

class SelectiveReLU(torch.autograd.Function):
    """
    选择性存储的 ReLU 实现
    在前向传播时不保存输入，反向传播时重新计算
    """
    @staticmethod
    def forward(ctx, input):
        # 前向传播不保存输入
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        # 需要时重新计算输入
        # 实际应用中这里需要重构前向计算逻辑
        # 为简化示例，我们假设能获取到原始输入
        input = ... # 从上层获取或重新计算
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

# 在模型中使用
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # 普通 ReLU
        x = self.pool(SelectiveReLU.apply(self.conv2(x)))  # 选择性 ReLU
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

我们在 CIFAR-10 数据集上测试了不同策略的效果：