幻方萤火 | 显存节省利器 CPUOffload

CPUOffload

 在深度学习模型的训练中，研究者与开发者们常常会碰到显存不足的问题 (OOM, out of memory) ，比如模型参数规模大，或者训练过程中产生的额外开销大，又或者是程序代码的问题，没有足够的显存资源，对我们科研与开发产生了诸多限制。

以往，我们通过代码优化、梯度累计、半精度等等一系列方法，以降低深度学习模型训练的显存需求，然而随着模型参数规模的发展，业务数据越来越复杂，显存需求快速增长，那么除了这些 trick 之外，还有什么既简单、又高效的显存节省方法吗？本期文章介绍的主角， hfai.nn.CPUOffload，给您提供一条不一样的显存节省之路。

那么 hfai.nn.CPUOffload 为什么可以节省显存？如何使用？它和 PyTorch 原始版本有何不同？本期文章将为大家分享 CPUOffload 设计背后的故事，讲述幻方 AI 设计 CPUOffload 的理念，展现 CPUOffload 的性能和便捷。

概述

训练模型的过程中，显存占用主要包括如下三部分

1. 权重矩阵

2. 前向传播的中间向量

3. 反向传播的梯度矩阵

我们以 ResNet50 网络在输入尺寸为（256，3，224，224）时的计算过程作为示例，分析上述三个部分的显存占用。通过 hfai.utils.profile_memory 工具，我们可以很容易的获得模型各部分的细节，如下图所示（由于网络层数较多，只展示部分网络）：

 每一行为一个 nn.Module 类型网络的详情信息：

• module name：nn.Module 类型网络的名称

• type：网络的类型

• parameter size: 网络参数量

• activation size: 中间变量的 tesnor 大小（不包括参数）

• #calls：被调用的次数