# Copyright 2021 MosaicML. All Rights Reserved.
"""The base :class:`~composer.trainer.devices.device.Device` class."""
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from contextlib import contextmanager
from typing import Any, Callable, Generator, TypeVar, Union
import torch
import torch.nn
from torch.optim import Optimizer
from composer.core.precision import Precision
from composer.core.serializable import Serializable
__all__ = ["Device", "T_nnModule"]
T_nnModule = TypeVar("T_nnModule", bound=torch.nn.Module)
T_Batch = TypeVar('T_Batch')
[docs]class Device(Serializable, ABC):
"""Abstract class for a device on which a model runs.
Attributes:
dist_backend (str): Distributed backend to use.
Should be ``gloo``, ``mpi``, or ``nccl``.
See `the pytorch docs <https://pytorch.org/docs/stable/distributed.html>`_
for details.
"""
dist_backend: str
[docs] @abstractmethod
def module_to_device(self, module: T_nnModule) -> T_nnModule:
"""Invoked by the :class:`.Trainer` to move a ``module`` onto the device.
Args:
module (torch.nn.Module): The module to move to the device.
Returns:
torch.nn.Module: The module on the device.
"""
pass
[docs] @abstractmethod
def tensor_to_device(self, tensor: torch.Tensor) -> torch.Tensor:
"""Invoked by the :class:`.Trainer` to move a tensor onto a device.
Args:
tensor (Tensor): The tensor to move to the device.
Returns:
Tensor: The tensor on the device.
"""
pass
[docs] def batch_to_device(self, batch: T_Batch) -> T_Batch:
"""Invoked by the :class:`.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.
Args:
batch (Any): The batch to move to the device.
Returns:
Batch: The batch on the device.
"""
def _to_device(x):
if isinstance(x, torch.Tensor):
return self.tensor_to_device(x)
return x
return _map_batch(batch, _to_device)
[docs] def optimizer_to_device(self, optimizer: Optimizer) -> Optimizer:
"""Invoked by the :class:`.Trainer` to move the optimizer's state onto the device.
Args:
optimizer (Optimizer): The optimizer to move to the device
Returns:
Optimizer: The optimizer on the device
"""
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = self.tensor_to_device(v)
return optimizer
[docs] @abstractmethod
@contextmanager
def precision_context(self, precision: Union[str, Precision]) -> Generator[None, None, None]:
"""Precision returns a context manager that uses the specified precision.
Example usage:
.. doctest::
>>> 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)
>>>
Args:
precision (Precision): The desired precision for the device.
Yields:
Generator[None, None, None]: A context for the precision.
"""
pass
def _map_batch(batch: Any, map_fn: Callable) -> Any:
"""Recursively maps a function to all items in a batch.
Args:
batch: Nested lists and dictionaries.
map_fn: A function to invoke on each element.
Returns:
Collections: The result of applying ``map_fn`` on each element of the ``batch``.
The type of ``batch`` is preserved.
"""
if isinstance(batch, Mapping):
return {k: _map_batch(v, map_fn) for k, v in batch.items()}
if isinstance(batch, Sequence) and not isinstance(batch, (str, bytes)):
try:
return type(batch)(_map_batch(x, map_fn) for x in batch) # type: ignore
except TypeError:
return [_map_batch(x, map_fn) for x in batch]
return map_fn(batch)