composer.trainer.devices.device#
The base Device
class.
Classes
Abstract class for a device on which a model runs. |
Attributes
T_nnModule
- class composer.trainer.devices.device.Device[source]#
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
, ornccl
. See the pytorch docs for details.- Type
- batch_to_device(batch)[source]#
Invoked by the
Trainer
move all tensors items in a batch to device. Supports nested sequences and mappings of tensors. Ignores non-tensor items. Preserves sequence and types when possible, otherwise converts sequences to lists. Converts mappings to dictionaries.- Parameters
batch (Any) โ The batch to move to the device.
- Returns
Batch โ The batch on the device.
- abstract module_to_device(module)[source]#
Invoked by the
Trainer
to move amodule
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)[source]#
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)[source]#
Precision returns a context manager that uses the specified precision.
Example usage:
>>> from composer.core.precision 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.