Python中PyTorch模型如何显存优化_使用梯度检查点减少显存占用

梯度检查点是通过只保存部分中间激活值、反向时重算前向来节省显存的技术能降低40%~60%显存但增加15%~30%训练时间要求模块前向可重入且无副作用。梯度检查点是什么为什么能省显存梯度检查点torch.utils.checkpoint.checkpoint不是“不存梯度”而是“不存中间激活值”。反向传播时需要前向计算的中间结果来算梯度常规训练会把所有层的输出全存着显存占用和网络深度线性增长而检查点只存部分层的输入反向时临时重跑对应前向——用时间换空间。典型节省比例ResNet-50 训练 batch size 可从 16 提到 32ViT-L 类模型显存常降 40%~60%。但注意它只对**前向可重入、无副作用**的模块有效比如不能包裹含 nn.Dropout训练态随机行为不可复现或修改全局状态的自定义层。怎么加 checkpoint最简安全写法别直接套整个模型先从单个 nn.Sequential 或自定义 forward 块开始。PyTorch 官方推荐方式是用 checkpoint.checkpoint 包裹函数调用而不是用装饰器后者容易隐式捕获非 tensor 参数。必须确保被包裹函数只接收 Tensor 参数且不依赖闭包变量如 self.training若模块含 training 切换逻辑如 Dropout改用 checkpoint.checkpoint_sequential 或手动拆分 torch.no_grad() 重算示例对 Transformer 层列表做检查点from torch.utils.checkpoint import checkpointdef custom_forward(x, layer): return layer(x)# 替换原循环x layer(x)x checkpoint(custom_forward, x, layer)常见报错和绕过方法RuntimeError: Trying to backward through the graph a second time说明 checkpoint 内部用了被复用的 Tensor比如共享 embedding或者你在检查点外又对同一张量调了 backward()。根本原因是计算图被意外保留。 VWO 一个A/B测试工具