composer.trainer.devices.device

The base Device class.

Classes

Device

Abstract class for a device on which a model runs.

Attributes

• T_nnModule

class composer.trainer.devices.device.Device

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

Abstract class for a device on which a model runs.

dist_backend

Distributed backend to use. Should be gloo, mpi, or nccl. See the pytorch docs for details.

Type

str

batch_to_device(batch)

Invoked by the Trainer to move the batch onto the device.

Parameters

batch (Batch) – The batch to move to the device.

Returns

Batch – The batch on the device.

abstract module_to_device(module)

Invoked by the Trainer to move a module onto the device.

Parameters

module (Module) – The module to move to the device.

Returns

torch.nn.Module – The module on the device.

optimizer_to_device(optimizer)

Invoked by the Trainer to move the optimizer’s state onto the device.

Parameters

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

Returns

Optimizer – The optimizer on the device

abstract precision_context(precision)

Precision returns a context manager that uses the specified precision.

Example usage:

>>> from composer.core.types import Precision
>>> from composer.trainer.devices import DeviceCPU
>>>
>>> device = DeviceCPU()
>>> for batch in train_dataloader:
...     with device.precision_context(Precision.FP32):
...         outputs = model.forward(batch)
...
...     with device.precision_context(Precision.FP32):
...         loss = model.loss(outputs, batch)
>>>

Parameters

precision (Precision) – The desired precision for the device.

Yields

Generator[None, None, None] – A context for the precision.

abstract tensor_to_device(tensor)

Invoked by the Trainer to move a tensor onto a device.

Parameters

tensor (Tensor) – The tensor to move to the device.

Returns

Tensor – The tensor on the device.