Apex 一键混合双精度计算

针对单一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的介绍和使用指南已经完全阐明,关于前文留的小问题,还欢迎小伙伴们留言回答哟~~~

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章