Device

class composer.trainer.devices.device.Device[source]

Bases: composer.core.serializable.Serializable, abc.ABC

Abstract class for a device on which a model runs.

abstract dataloader_to_device(dataloader: DataLoader, prefetch_fn: Optional[Callable[[Union[Tuple[Union[Tensor, Tuple[Tensor, ...], List[Tensor]], Union[Tensor, Tuple[Tensor, ...], List[Tensor]]], List[Tensor], Dict[str, Tensor], Tensor]], Union[Tuple[Union[Tensor, Tuple[Tensor, ...], List[Tensor]], Union[Tensor, Tuple[Tensor, ...], List[Tensor]]], List[Tensor], Dict[str, Tensor], Tensor]]]) → 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.

abstract 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

abstract 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.

optimizer_to_device(optimizer: torch.optim.optimizer.Optimizer) → torch.optim.optimizer.Optimizer[source]

Moves an optimizer’s state onto the device instance’s device.

As a rule, this usually isn’t necessary, since most optimizers lazy initialize their state when .step() is first called, based off of the device of the parameters. The prominent exception to this rule is when we are restoring from a checkpoint.

Parameters

optimizer (Optimizer) – The optimizer to move to the device

Returns

Optimizer – The optimizer on the device

abstract 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]

abstract 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.

abstract tensor_to_device(tensor: Tensor) → 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.