autocast

class tensorplay.autocast(device_type: str, dtype: Any | None = None, enabled: bool = True, cache_enabled: bool | None = None)[source]

Instances of autocast serve as context managers or decorators that allow regions of your script to run in mixed precision.

In these regions, ops run in an op-specific dtype chosen by autocast to improve performance while maintaining accuracy.

When entering an autocast-enabled region, Tensors may be any type. You should not call half() or bfloat16() on your model(s) or inputs when using autocasting.

autocast should wrap only the forward pass(es) of your network, including the loss computation(s). Backward passes under autocast are not recommended. Backward ops run in the same type that autocast used for corresponding forward ops.

Example for CUDA Devices:

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # Enables autocasting for the forward pass (model + loss)
    with tensorplay.autocast(device_type="cuda"):
        output = model(input)
        loss = loss_fn(output, target)

    # Exits the context manager before backward()
    loss.backward()
    optimizer.step()

autocast can also be used as a decorator, e.g., on the forward method of your model:

class AutocastModel(nn.Module):
    ...

    @tensorplay.autocast(device_type="cuda")
    def forward(self, input): ...

Floating-point Tensors produced in an autocast-enabled region may be float16. After returning to an autocast-disabled region, using them with floating-point Tensors of different dtypes may cause type mismatch errors. If so, cast the Tensor(s) produced in the autocast region back to float32 (or other dtype if desired).

autocast(enabled=False) subregions can be nested in autocast-enabled regions. Locally disabling autocast can be useful, for example, if you want to force a subregion to run in a particular dtype.

The autocast state is thread-local. If you want it enabled in a new thread, the context manager or decorator must be invoked in that thread.

Parameters:

• device_type (str, required) – Device type to use. Possible values are: ‘cuda’ and ‘cpu’. The type is the same as the type attribute of a tensorplay.device. Thus, you may obtain the device type of a tensor using Tensor.device.type.

• enabled (bool, optional) – Whether autocasting should be enabled in the region. Default: True

• dtype (tensorplay.dtype, optional) – Data type for ops run in autocast. It uses the default value (tensorplay.float16 for CUDA and tensorplay.bfloat16 for CPU), given by get_autocast_dtype(), if dtype is None. Default: None

• cache_enabled (bool, optional) – Whether the weight cache inside autocast should be enabled. Default: True