Accelerator

class pytorch_lightning.accelerators.Accelerator(precision_plugin, training_type_plugin)

Bases: object

The Accelerator Base Class. An Accelerator is meant to deal with one type of Hardware.

Currently there are accelerators for:

• CPU

• GPU

• TPU

• IPU

Each Accelerator gets two plugins upon initialization: One to handle differences from the training routine and one to handle different precisions.

Parameters

• precision_plugin (PrecisionPlugin) – the plugin to handle precision-specific parts

• training_type_plugin (TrainingTypePlugin) – the plugin to handle different training routines

all_gather(tensor, group=None, sync_grads=False)

Function to gather a tensor from several distributed processes.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.all_gather directly.

Parameters

• tensor (Tensor) – tensor of shape (batch, …)

• group (Optional[Any]) – the process group to gather results from. Defaults to all processes (world)

• sync_grads (bool) – flag that allows users to synchronize gradients for all_gather op

Return type

Tensor

Returns

A tensor of shape (world_size, batch, …)

abstract static auto_device_count()

Get the devices when set to auto.

Return type

int

backward(closure_loss, *args, **kwargs)

Forwards backward-calls to the precision plugin.

Parameters

closure_loss (Tensor) – a tensor holding the loss value to backpropagate

Return type

Tensor

barrier(name=None)

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.barrier directly.

Return type

None

batch_to_device(batch, device=None, dataloader_idx=0)

Moves the batch to the correct device. The returned batch is of the same type as the input batch, just having all tensors on the correct device.

Parameters

• batch (Any) – The batch of samples to move to the correct device

• device (Optional[device]) – The target device

• dataloader_idx (int) – The index of the dataloader to which the batch belongs.

Return type

Any

broadcast(obj, src=0)

Broadcasts an object to all processes, such that the src object is broadcast to all other ranks if needed.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.broadcast directly.

Parameters

• obj (object) – Object to broadcast to all process, usually a tensor or collection of tensors.

• src (int) – The source rank of which the object will be broadcast from

Return type

object

connect(model)

Transfers ownership of the model to this plugin.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_train_batch_start directly.

Return type

None

dispatch(trainer)

Hook to do something before the training/evaluation/prediction starts.

Return type

None

get_device_stats(device)

Gets stats for a given device.

Parameters

device (Union[str, device]) – device for which to get stats

Return type

Dict[str, Any]

Returns

Dictionary of device stats

lightning_module_state_dict()

Returns state of model.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.lightning_module_state_dict directly.

Allows for syncing/collating model state from processes in custom plugins.

Return type

Dict[str, Union[Any, Tensor]]

model_sharded_context()

Provide hook to create modules in a distributed aware context. This is useful for when we’d like to.

shard the model instantly - useful for extremely large models. Can save memory and initialization time.

Return type

Generator[None, None, None]

Returns

Model parallel context.

on_predict_end()

Called when predict ends.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_predict_end directly.

Return type

None

on_predict_start()

Called when predict begins.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_predict_start directly.

Return type

None

on_test_end()

Called when test end.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_test_end directly.

Return type

None

on_test_start()

Called when test begins.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_test_start directly.

Return type

None

on_train_batch_start(batch, batch_idx, dataloader_idx=0)

Called in the training loop before anything happens for that batch.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_train_batch_start directly.

Return type

None

on_train_end()

Called when train ends.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_train_end directly.

Return type

None

on_train_start()

Called when train begins.

Return type

None

on_validation_end()

Called when validation ends.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_validation_end directly.

Return type

None

on_validation_start()

Called when validation begins.

See deprecation warning below.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.on_validation_start directly.

Return type

None

optimizer_state(optimizer)

Returns state of an optimizer.

Allows for syncing/collating optimizer state from processes in custom plugins.

Return type

Dict[str, Tensor]

optimizer_step(optimizer, opt_idx, closure, model=None, **kwargs)

performs the actual optimizer step.

Parameters

• optimizer (Optimizer) – the optimizer performing the step

• opt_idx (int) – index of the current optimizer

• closure (Callable[[], Any]) – closure calculating the loss value

• model (Union[LightningModule, Module, None]) – reference to the model, optionally defining optimizer step related hooks

• **kwargs – Any extra arguments to optimizer.step

Return type

None

optimizer_zero_grad(current_epoch, batch_idx, optimizer, opt_idx)

Zeros all model parameter’s gradients.

Return type

None

post_dispatch(trainer)

Hook to do something after the training/evaluation/prediction starts.

Return type

None

post_training_step()

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.post_training_step directly.

Return type

None

pre_dispatch(trainer)

Hook to do something before the training/evaluation/prediction starts.

Return type

None

predict_step(step_kwargs)

The actual predict step.

See predict_step() for more details

Return type

Union[Tensor, Dict[str, Any]]

process_dataloader(dataloader)

Wraps the dataloader if necessary.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.process_dataloader directly.

Parameters

dataloader (Union[Iterable, DataLoader]) – iterable. Ideally of type: torch.utils.data.DataLoader

Return type

Union[Iterable, DataLoader]

save_checkpoint(checkpoint, filepath)

Save model/training states as a checkpoint file through state-dump and file-write.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.save_checkpoint directly.

Parameters

• checkpoint (Dict[str, Any]) – dict containing model and trainer state

• filepath (Union[str, Path]) – write-target file’s path

Return type

None

setup(trainer)

Setup plugins for the trainer fit and creates optimizers.

Parameters

trainer (Trainer) – the trainer instance

Return type

None

setup_environment()

Setup any processes or distributed connections.

This is called before the LightningModule/DataModule setup hook which allows the user to access the accelerator environment before setup is complete.

Return type

None

setup_optimizers(trainer)

Creates optimizers and schedulers.

Parameters

trainer (Trainer) – the Trainer, these optimizers should be connected to

Return type

None

setup_precision_plugin()

Attaches the precision plugin to the accelerator.

Return type

None

setup_training_type_plugin()

Attaches the training type plugin to the accelerator.

Return type

None

start_evaluating(trainer)

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.start_evaluating directly.

Return type

None

start_predicting(trainer)

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.start_predicting directly.

Return type

None

start_training(trainer)

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.start_training directly.

Return type

None

teardown()

This method is called to teardown the training process.

It is the right place to release memory and free other resources.

Return type

None

test_step(step_kwargs)

The actual test step.

See test_step() for more details

Return type

Union[Tensor, Dict[str, Any], None]

test_step_end(output)

A hook to do something at the end of the test step.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.test_step_end directly.

Parameters

output (Union[Tensor, Dict[str, Any], None]) – the output of the test step

Return type

Union[Tensor, Dict[str, Any], None]

training_step(step_kwargs)

The actual training step.

See training_step() for more details

Return type

Union[Tensor, Dict[str, Any]]

training_step_end(output)

A hook to do something at the end of the training step.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.training_step_end directly.

Parameters

output (Union[Tensor, Dict[str, Any]]) – the output of the training step

Return type

Union[Tensor, Dict[str, Any]]

validation_step(step_kwargs)

The actual validation step.

See validation_step() for more details

Return type

Union[Tensor, Dict[str, Any], None]

validation_step_end(output)

A hook to do something at the end of the validation step.

Deprecated since version v1.5: This method is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.validation_step_end directly.

Parameters

output (Union[Tensor, Dict[str, Any], None]) – the output of the validation step

Return type

Union[Tensor, Dict[str, Any], None]

property lightning_module: pytorch_lightning.core.lightning.LightningModule

Returns the pure LightningModule.

To get the potentially wrapped model use Accelerator.model

Return type

LightningModule

property model: torch.nn.modules.module.Module

Returns the model.

This can also be a wrapped LightningModule. For retrieving the pure LightningModule use Accelerator.lightning_module

Return type

Module

property restore_checkpoint_after_pre_dispatch: bool

Override to delay restoring from checkpoint till after pre-dispatch. This is useful when the plugin requires all the setup hooks to run before loading checkpoint.

Deprecated since version v1.5: This property is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.restore_checkpoint_after_pre_dispatch directly.

Return type

bool

Returns

If true, restore checkpoint after pre_dispatch.

property results: Any

The results of the last run will be cached within the training type plugin.

Deprecated since version v1.5: This property is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.results directly.

In distributed training, we make sure to transfer the results to the appropriate master process.

Return type

Any

property root_device: torch.device

Returns the root device.

Return type

device

property setup_optimizers_in_pre_dispatch: bool

Override to delay setting optimizers and schedulers till after dispatch. This is useful when the TrainingTypePlugin requires operating on the wrapped accelerator model. However this may break certain precision plugins such as APEX which require optimizers to be set.

Deprecated since version v1.5: This property is deprecated in v1.5 and will be removed in v1.6. Please call training_type_plugin.setup_optimizers_in_pre_dispatch directly.

Return type

bool

Returns

If True, delay setup optimizers until pre_dispatch, else call within setup.