混合精度加速GPU计算:apex

混合精度加速GPU计算:apex

在DistributedDataParallel的基础上,我们可以使用apex包来进行加速.
只需要简单的两步即可:
一,在model = torch.nn.parallel.DistributedDataParallel(model)之前通过apex来初始化model以及optimizer

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

二,为了缓解单精度计算更容易溢出,以及其产生的舍入误差问题,通过apex对training中的loss进行scale,即用scaled_loss.backward()来代替loss.backward()

with amp.scale_loss(loss, optimizer) as scaled_loss:
	scaled_loss.backward()