megengine.amp#

GradScaler

autocast 模式下执行梯度缩放以防止数据溢出的helper类。

autocast

作为上下文管理器或装饰器来控制amp的自动转换模式的类。

convert_module_format

将 NCHW 模块转换为 NHWC 模块。

convert_tensor_format

将 NCHW 张量转换为 NHWC 张量。

enabled

获取或设置是否启用 amp autocast.

high_prec_dtype

获取或设置amp autocast模式的最高精度数据类型

low_prec_dtype

获取或设置amp autocast模式的最低精度数据类型