def backward(tensor: torch.Tensor) -> None:
"""Computes the gradient of the specified tensor w.r.t. graph leaves.
Args:
tensor (torch.Tensor): Tensor of which the derivative will be computed.
"""
get_accelerator(_TorchAccelerator).backward(tensor)
def prepare_data_loader(
data_loader: torch.utils.data.DataLoader,
add_dist_sampler: bool = True,
move_to_device: bool = True,
auto_transfer: bool = True,
) -> torch.utils.data.DataLoader:
"""Prepares DataLoader for distributed execution.
This allows you to use the same exact code regardless of number of
workers or the device type being used (CPU, GPU).
Args:
data_loader (torch.utils.data.DataLoader): The DataLoader to
prepare.
add_dist_sampler: Whether to add a DistributedSampler to
the provided DataLoader.
move_to_device: If set, automatically move the data
returned by the data loader to the correct device.
auto_transfer: If set and device is GPU, another CUDA stream
is created to automatically copy data from host (CPU) memory
to device (GPU) memory (the default CUDA stream still runs the
training procedure). If device is CPU, it will be disabled
regardless of the setting. This configuration will be ignored
if ``move_to_device`` is False.
"""
return get_accelerator(_TorchAccelerator).prepare_data_loader(
data_loader,
add_dist_sampler=add_dist_sampler,
move_to_device=move_to_device,
auto_transfer=auto_transfer,
)
def prepare_model(
model: torch.nn.Module,
move_to_device: bool = True,
wrap_ddp: bool = True,
ddp_kwargs: Optional[Dict[str, Any]] = None,
) -> torch.nn.Module:
"""Prepares the model for distributed execution.
This allows you to use the same exact code regardless of number of
workers or the device type being used (CPU, GPU).
Args:
model (torch.nn.Module): A torch model to prepare.
move_to_device: Whether to move the model to the correct
device. If set to False, the model needs to manually be moved
to the correct device.
wrap_ddp: Whether to wrap models in
``DistributedDataParallel``.
ddp_kwargs (Dict[str, Any]): Args to pass into
``DistributedDataParallel`` initialization if ``wrap_ddp`` is
set to True.
"""
return get_accelerator(_TorchAccelerator).prepare_model(
model,
move_to_device=move_to_device,
wrap_ddp=wrap_ddp,
ddp_kwargs=ddp_kwargs,
)
def enable_reproducibility(seed: int = 0) -> None:
"""Limits sources of nondeterministic behavior.
This function:
* Seeds PyTorch, Python, and NumPy.
* Disables CUDA convolution benchmarking.
* Configures PyTorch to use determinstic algorithms.
* Seeds workers spawned for multi-process data loading.
Args:
seed: The number to seed libraries and data workers with.
.. warning:: ``train.torch.enable_reproducibility()`` can't guarantee
completely reproducible results across executions. To learn more, read
the `PyTorch notes on randomness
<https://pytorch.org/docs/stable/notes/randomness.html>`_.
"""
get_accelerator(_TorchAccelerator).enable_reproducibility(seed)
def prepare_optimizer(
optimizer: torch.optim.Optimizer) -> torch.optim.Optimizer:
"""Wraps optimizer to support automatic mixed precision.
Args:
optimizer (torch.optim.Optimizer): The DataLoader to prepare.
Returns:
A wrapped optimizer.
"""
return get_accelerator(_TorchAccelerator).prepare_optimizer(optimizer)
def test_get_accelerator_raises_error_outside_session():
with pytest.raises(SessionMisuseError):
get_accelerator(FakeAccelerator)
def test_get_accelerator_constructs_default_accelerator(session):
assert isinstance(get_accelerator(FakeAccelerator), FakeAccelerator)
def test_set_and_get_accelerator(session):
accelerator = FakeAccelerator()
set_accelerator(accelerator)
assert get_accelerator(FakeAccelerator) is accelerator
def get_device() -> torch.device:
"""Gets the correct torch device to use for training."""
return get_accelerator(_TorchAccelerator).get_device()
