every blog every motto:Just live your life cause we don’t live twice.

0. 前言

深度学习加速之自动混合精度

1. 正文

总共两种方法

1.1. 方法一：NVIDIA第三方包

需求：pytorch1.5之前的版本（包括1.5）

使用方法：

from apex import ampmodel,optimizer = amp.initial(model,optimizer,opt_level="O1") #注意是O,不是０with amp.scale_loss(loss,optimizer) as scaled_loss:scaled_loss.backward()取代loss.backward()

其中,opt_level配置如下：

O0:纯FP32训练，可作为accuracy的baseline；

O1:混合精度训练（推荐使用），根据黑白名单自动决定使用FP16(GEMM,卷积）还是FP32（softmax)进行计算。

O2:几乎FP16，混合精度训练，不存在黑白名单，除了bacthnorm，几乎都是用FP16计算；

O3:纯FP16训练，很不稳定，但是可以作为speed的baseline；

1.2 方法二：pytorch1.6自带的torch.cuda.amp

有两个接口

1. autocast

from torch.cuda.amp import autocast as autocastmodel=Net().cuda()optimizer=optim.SGD(model.parameters(),...)for input,target in data:optimizer.zero_grad()with autocast():output=model(input)loss = loss_fn(output,target)loss.backward()optimizer.step()

说明：参考文献2中指出，需要GPU支持Tensor core才可以使用，笔者笔记本GPU为GTX 1050Ti，按理说是不支持的，但实验发现占用显存情况确有明显降低，具体如下：

没有使用：

使用以后：

2. GradScaler

from torch.cuda.amp import autocast as autocastmodel=Net().cuda()optimizer=optim.SGD(model.parameters(),...)scaler = GradScaler() #训练前实例化一个GradScaler对象for epoch in epochs:for input,target in data:optimizer.zero_grad()with autocast():＃前后开启autocastoutput=model(input)loss = loss_fn(output,targt)scaler.scale(loss).backward() #为了梯度放大#scaler.step()首先把梯度值unscale回来，如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。 scaler.step(optimizer)scaler.update() #准备着，看是否要增大scaler

参考文献

[1] /jimchen1218/p/14315008.html

[2] /p/295468659