深度学习训练降低显存指南

Myth丶恋晨 2022-10-18 00:59 303阅读 0赞

# 一、小模块API参数inplace设置为True（省一点点） #

比如：Relu()有一个默认参数inplace，默认设置为False，当设置为True时，计算时的得到的新值不会占用新的空间而是直接覆盖原来的值，进而可以节省一点点内存。

# 二、Apex半精度计算（省一半左右） #

安装方式

git clone https://github.com/NVIDIA/apex
    cd apex
    python3 setup.py install

原理：一款Nvidia开发的基于PyTorch的混合精度训练加速神器，可以用短短三行代码（导包、初始化、反向传播）就能实现不同程度的**混合精度加速**，训练时间直接缩小一半。

from apex import amp
    model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”
    with amp.scale_loss(loss, optimizer) as scaled_loss:
        scaled_loss.backward()

其中只有一个`opt_level`需要用户自行配置（\# 这里是**“欧x”**，不是“零x”）：

>  *  `O0`：纯FP32训练，可以作为accuracy的baseline；
>  *  `O1`：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（GEMM, 卷积）还是FP32（Softmax）进行计算。
>  *  `O2`：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算。
>  *  `O3`：纯FP16训练，很不稳定，但是可以作为speed的baseline；

注：apex不能和contiguous-params一起用，否则无法实现模型参数优化。（因为这俩模型参数指针要打架）

# 三、删无用变量并及时清空显存垃圾（省将近一半） #

这个的用法比较简单，就是先del变量，再cuda垃圾回收。

del answers_types_emb_all, answers_types_len
    torch.cuda.empty_cache()

但这个的使用方式，不能在整个代码中只使用一次torch.cuda.empty\_cache()——否则整体的最大显存不会变化。

正确的使用方法是：应该在整体代码的多个位置，分别del变量后，对应位置之后即使用torch.cuda.empty\_cache()。

这样整体最大显存，就会下降（del的变量整体显存之和 - max的一坨del变量的显存）——分批清空显存，才能真正地用torch.cuda.empty\_cache()降低显存。

# 四、重计算方法（省一半以上） #

用pytorch的checkpoint断点保存机制，将连续的几个函数计算，分为若干段，保存了再计算——用保存加载的时间，换显存计算的空间

原来一个函数的计算方式：

sequence_enc = self.encoder(sequence_emb, sequence_mask, output_all_encoded_layers=False)[-1]

用checkpoint断点保存的方式：