Created by: Jie-Fang
Implement custom black white list feature. Users can specify custom ops to use fp16 or fp32 calculation.
Sample code:
optimizer = fluid.optimizer.SGD(learning_rate=0.001)
amp_lists = AutoMixedPrecisionLists(custom_white_list={"conv2d", "matmul"}, custom_black_list={"exp", "sum"})
mp_optimizer = fluid.contrib.mixed_precision.decorate(optimizer=optimizer,
amp_lists=amp_lists,
init_loss_scaling=8.0,
use_dynamic_loss_scaling=True)
mp_optimizer.minimize(cost)
Users can pass custom white list and black list to construct AutoMixedPrecisionLists
, then pass AutoMixedPrecisionLists
object to decorate
function. AutoMixedPrecisionLists
will update default white and black list according to custom lists. If users want to use default lists, they don't need to pass amp_lists, and AMP will use default lists.