auto_cast

paddle.amp. auto_cast ( enable=True, custom_white_list=None, custom_black_list=None, level='O1', dtype='float16' ) [source]

Create a context which enables auto-mixed-precision(AMP) of operators executed in dynamic graph mode. If enabled, the input data type (float32 or float16) of each operator is decided by autocast algorithm for better performance.

Commonly, it is used together with GradScaler to achieve Auto-Mixed-Precision in imperative mode. It is used together with decorator to achieve Pure fp16 in imperative mode.

Parameters
  • enable (bool, optional) – Enable auto-mixed-precision or not. Default is True.

  • custom_white_list (set|list|tuple, optional) – The custom white_list. It’s the set of ops that support fp16 calculation and are considered numerically-safe and performance-critical. These ops will be converted to fp16.

  • custom_black_list (set|list|tuple, optional) – The custom black_list. The set of ops that support fp16 calculation and are considered numerically-dangerous and whose effects may also be observed in downstream ops. These ops will not be converted to fp16.

  • level (str, optional) – Auto mixed precision level. Accepted values are “O1” and “O2”: O1 represent mixed precision, the input data type of each operator will be casted by white_list and black_list; O2 represent Pure fp16, all operators parameters and input data will be casted to fp16, except operators in black_list, don’t support fp16 kernel and batchnorm. Default is O1(amp)

  • dtype (str, optional) – Whether to use ‘float16’ or ‘bfloat16’. Default is ‘float16’.

Examples

import paddle

conv2d = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
data = paddle.rand([10, 3, 32, 32])

with paddle.amp.auto_cast():
    conv = conv2d(data)
    print(conv.dtype) # FP16

with paddle.amp.auto_cast(enable=False):
    conv = conv2d(data)
    print(conv.dtype) # FP32

with paddle.amp.auto_cast(custom_black_list={'conv2d'}):
    conv = conv2d(data)
    print(conv.dtype) # FP32

a = paddle.rand([2,3])
b = paddle.rand([2,3])
with paddle.amp.auto_cast(custom_white_list={'elementwise_add'}):
    c = a + b
    print(c.dtype) # FP16

with paddle.amp.auto_cast(custom_white_list={'elementwise_add'}, level='O2'):
    d = a + b
    print(d.dtype) # FP16

Used in the guide/tutorials