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, cells that are not in black_list will perform the pereision conversion. Only one of white_list and black_list should be provided.
Note
After using custom_mixed_precision for precision conversion, it is not supported to use other interfaces for precision conversion again. If interfaces like Model and build_train_network is used to train the converted network, amp_level need to be configured to
O0
to avoid the duplicated accuracy conversion.Primitives for blacklist is not support yet.
- Parameters
- 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, nn >>> # Define the network structure of LeNet5. Refer to >>> # https://gitee.com/mindspore/docs/blob/r2.1/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)