mindspore.amp.custom_mixed_precision

mindspore.amp.custom_mixed_precision(network, *, white_list=None, black_list=None)[source]

Custom mixed precision by setting whitelist or blacklist. When the white_list is provided, primitives and cells in white_list will perform the precision conversion. When the black_list is provided, primitives and cells that are not in black_list will perform the pereision conversion. Only one of white_list and black_list should be provided.

Warning

This is an experimental API that is subject to change or deletion.

Note

  • custom_mixed_precision should not be used at the same time as auto_mixed_precision . When both build_train_network and custom_mixed_precision are used, build_train_network need to be called with level=’O0’ before call custom_mixed_precision .

  • Primitives for blacklist is not support yet.

Parameters

  • network (Cell) – Definition of the network.

  • white_list (list[Primitive, Cell], optional) – White list of custom mixed precision. Defaults: None, means white list is not used.

  • black_list (list[Primitive, Cell], optional) – Black list of custom mixed precision. Defaults: None, means black list is not used.

Returns

network (Cell), A network supporting mixed precision.

Raises

  • TypeError – The network type is not Cell.

  • ValueError – Neither white_list nor black_list is provided.

  • ValueError – Both white_list and black_list are provided.

Examples

>>> from mindspore import amp
>>> # Define the network structure of LeNet5. Refer to
>>> # https://gitee.com/mindspore/docs/blob/r2.0/docs/mindspore/code/lenet.py
>>> net = LeNet5()
>>> custom_white_list = amp.get_white_list()
>>> custom_white_list.append(nn.Flatten)
>>> net = amp.custom_mixed_precision(net, white_list=custom_white_list)