megengine.amp

GradScaler([init_scale, growth_factor, …])

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

autocast([enabled, low_prec_dtype, …])

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

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.