Shortcuts

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.

Pytorch 1.5.0 provide torch.cuda.amp.autocast for running in mixed precision , and update it to torch.autocast in 1.10.0. Both interfaces have different arguments, and torch.autocast support running with cpu additionally.

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.5.0. If pytorch version <= 1.10.0 and cuda is not available, it will raise an error with enabled=True, since torch.cuda.amp.autocast only support cuda mode.

Examples

>>> # case1: 1.10 > Pytorch version >= 1.5.0
>>> with autocast():
>>>    # run in mixed precision context
>>>    pass
>>> with autocast(device_type='cpu')::
>>>    # raise error, torch.cuda.amp.autocast only support cuda mode.
>>>    pass
>>> # case2: 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.

Read the Docs v: v0.8.3
Versions
latest
stable
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.