megengine.amp

GradScaler

A helper class that performs grad scaling to prevent from data overflow in autocast mode.

autocast

A class to control autocast mode for amp as a context manager or a decorator.

convert_module_format

Convert NCHW Module to NHWC Module.

convert_tensor_format

Convert NCHW Tensor to NHWC Tensor.

enabled

Get or set amp autocast mode enabled or not.

high_prec_dtype

Get or set amp autocast mode's higher precision dtype.

low_prec_dtype

Get or set amp autocast mode's lower precision dtype.