針對單一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的介紹和使用指南已經完全闡明，關於前文留的小問題，還歡迎小夥伴們留言回答喲~~~