DeviceGPU

class composer.trainer.devices.device_gpu.DeviceGPU(prefetch_in_cuda_stream: bool)[source]

Bases: composer.trainer.devices.device.Device

An extension of Device for GPUs.

Parameters: prefetch_in_cuda_stream (bool) – True to asyncrhonously prefetch samples with a CUDA stream during dataloading and False otherwise.

dataloader_to_device(dataloader: DataLoader, prefetch_fn: Optional[TPrefetchFn]) → DataLoader[source]

Wraps a Dataloader and ensures all returned batches are on the correct device.

This function is responsible for executing prefetch_fn, if provided, on each batch before it is yielded. The prefetch_fn can be executed in the background, if the device supports it.

Parameters

dataloader (DataLoader) – The dataloader to wrap.
prefetch_fn (Optional[TPrefetchFn]) – A function that takes a batch and returns a batch. It should perform any on-device preprocessing of a batch. (e.g. on a GPU device, this function can be used for gpu transformations.)

Returns

DataLoader – The wrapped dataloader, which yields batches that
have been moved to the device and have been processed through
the prefetch_fn.

property ddp_backend: str

DDP backend to use.

Should return gloo, mpi, or nccl. See the pytorch docs for details.

Returns: str – gloo, mpi, or nccl

load_state_dict(state: composer.core.types.StateDict) → None[source]

Restores the state of the object.

Parameters: state (StateDict) – The state of the object, as previously returned by state_dict()

module_to_device(module: composer.trainer.devices.device.T_nnModule) → composer.trainer.devices.device.T_nnModule[source]

Moves a module onto the device instance’s device.

Parameters: module (T_nnModule) – The module to move to the device
Returns: T_nnModule – The module on the device.

precision_context(precision: Union[str, Precision]) → Generator[None, None, None][source]

Precision returns a context manager that uses the specified precision.

Example usage:

with device.precision(Precision.AMP):
    forward_pass_with_amp()

Parameters: precision (Precision) – [description]
Yields: Generator[None, None, None] – [description]

prepare(state: State) → None[source]

Used for device initialization.

Invoked by the trainer at the beginning of the training loop. It should not modify the state.

Parameters: state (State) – The global state variable.

state_dict() → composer.core.types.StateDict[source]

Returns a dictionary representing the internal state.

The returned dictionary must be pickale-able via torch.save().

Returns: StateDict – The state of the object

tensor_to_device(tensor: composer.core.types.Tensor) → composer.core.types.Tensor[source]

Moves a tensor onto the device instance’s device.

Parameters: tensor (T_nnModule) – The tensor to move to the device
Returns: Tensor – The tensor on the device.