针对单一Optimizer使能AMP
对於单一Optimizer,使能AMP只需要以下修改两部分即可:
1. 代码中Optimizer之后添加:
optimizer = # ... some optimizer
from apex import amp
amp_handle = amp.init()
2. 修改代码中反向迭代部分:
# ... do a bunch of stuff to compute a loss
loss.backward()
optimizer.step()
# ...finish the iteration
修改为
# ... same as before
with amp_handle.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
optimizer.step()
# ... same as before
针对多Optimizer使能AMP
多Optimizer使能AMP相较单一Optimizer复杂,需要显性的warp多个Optimizer,并处理多个Loss的反向迭代,例如:
amp_handle = amp.init()
optimizer = amp_handle.wrap_optimizer(optimizer, num_loss=2)
# ...
optimizer.zero_grad()
loss1 = ComputeLoss1(model)
with optimizer.scale_loss(loss1) as scaled_loss:
scaled_loss.backward()
# ...
loss2 = ComputeLoss2(model)
with optimizer.scale_loss(loss2) as scaled_loss:
scaled_loss.backward()
# ...
optimizer.step()
AMP添加自定义函数
在AMP添加自定义函数的支持有两种方法,分别使用Function annotation或者函数显性注册如下所示:
对于自定义函数Fru
from backend import FRUBackend
def fru(input, hidden, weight, bias):
# call to CUDA code
FRUBackend(input, hidden, weight, bias)
使用Function Annotation
@amp.half_function
def fru(input, hidden, weight, bias):
#...
使用显性注册
import backend
amp.register_half_function(backend, 'FRUBackend')
amp.init()
此处的Function Annotation有三种类型分别是:
类似显性注册也有对应三种方式,分别是:
amp.register_half_function(module, function_name)
amp.register_float_function(module, function_name)
amp.register_promote_function(module, function_name)
这里需要注意的是,采用显性注册需要在amp.init()之前。
至此,关于AMP的介绍和使用指南已经完全阐明,关于前文留的小问题,还欢迎小伙伴们留言回答哟~~~