mmengine.runner.autocast¶

mmengine.runner.autocast(device_type=None, dtype=None, enabled=True, cache_enabled=None)[source]¶

A wrapper of torch.autocast and toch.cuda.amp.autocast.

This function provides a unified interface by wrapping torch.autocast and torch.cuda.amp.autocast, which resolves the compatibility issues that torch.cuda.amp.autocast does not support running mixed precision with cpu, and both contexts have different arguments. We suggest users using this function in the code to achieve maximized compatibility of different PyTorch versions.

Note

autocast requires pytorch version >= 1.10.0. If pytorch version <= 1.10.0 it will raise an error.

Examples

>>> # case1: Pytorch version >= 1.10.0
>>> with autocast():
>>>    # default cuda mixed precision context
>>>    pass
>>> with autocast(device_type='cpu'):
>>>    # cpu mixed precision context
>>>    pass
>>> with autocast(
>>>     device_type='cuda', enabled=True, cache_enabled=True):
>>>    # enable precision context with more specific arguments.
>>>    pass

Parameters:

device_type (str, required) – Whether to use ‘cuda’ or ‘cpu’ device.
enabled (bool) – Whether autocasting should be enabled in the region. Defaults to True
dtype (torch_dtype, optional) – Whether to use torch.float16 or torch.bfloat16.
cache_enabled (bool, optional) – Whether the weight cache inside autocast should be enabled.