转载请注明出处
个人博客: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
。可见效果还是好的。不过计算速度方面没有统计。