pytorch减少精度提升速度、减少GPU内存

转载请注明出处
个人博客:https://maxusun.github.io/

提高运算速度，减少GPU占用内存的方法有很多，其中减少精度的方法对大多数实验结果都没有太大影响。在这里使用Apex比使用pytorch的half靠谱多了.

Apex

Apex具体的介绍可以看这篇文章：基于Apex的混合精度加速

安装

如果直接使用pip install apex安装，可能在运行的时候得到下面的错误：

TypeError: Class advice impossible in Python3. Use the @Implementer class decorator instead

这时候你需要使用下面的方式安装：

pip uninstall apex
git clone https://www.github.com/nvidia/apex
cd apex
python setup.py install

使用方法

使用方法只有简单的三行代码：

from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”

# 下面两句替换源代码里面的loss.backward()
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

多网络，bn同步，多卡等情况：

from apex import amp
(net,net_fc),optimizer = amp.initialize([net,net_fc], optimizer, opt_level= "O1")#这里多个net就用列表
with amp.scale_loss(loss_mean, optimizer) as scaled_loss:
scaled_loss.backward()
#bn同步
from apex.parallel import convert_syncbn_model
net = convert_syncbn_model(net)
net = net.to(device)
#多卡
from apex.parallel import DistributedDataParallel as DDP
net, opt = amp.initialize(net, opt, opt_level="O1")  # 字母小写o,不是零。
net = DDP(net, delay_allreduce=True)  # 必须在initialze之后

opt_level参数设置

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

结果

没有使用前，在1080上使用GPU的内存大小为12G，用了之后降到7.8G。可见效果还是好的。不过计算速度方面没有统计。