Strategy

class lightning.pytorch.strategies.Strategy(accelerator=None, checkpoint_io=None, precision_plugin=None)

Bases: ABC

Base class for all strategies that change the behaviour of the training, validation and test- loop.

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

Perform an all_gather on all processes.

Parameters

• tensor (Tensor) – the tensor to all_gather

• group (Optional[Any]) – the process group to gather results from

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

Return type

Tensor

backward(closure_loss, optimizer, *args, **kwargs)

Forwards backward-calls to the precision plugin.

Parameters

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

• optimizer (Optional[Optimizer]) – An optional optimizer that gets passed down to the precision plugin’s backward

• *args – Positional arguments that get passed down to the precision plugin’s backward, intended as arguments for the actual function that performs the backward, like backward().

• **kwargs – Keyword arguments for the same purpose as *args.

Return type

Tensor

abstract barrier(name=None)

Synchronizes all processes which blocks processes until the whole group enters this function.

Parameters

name (Optional[str]) – an optional name to pass into barrier.

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

abstract broadcast(obj, src=0)

Broadcasts an object to all processes.

Parameters

• obj (TypeVar(TBroadcast)) – the object to broadcast

• src (int) – source rank

Return type

TypeVar(TBroadcast)

connect(model)

Called by the accelerator to connect the accelerator and the model with this plugin.

Return type

None

lightning_module_state_dict()

Returns model state.

Return type

Dict[str, Any]

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, which is useful for extremely large models which can save memory and initialization time.

Returns: Model parallel context.

Return type

Generator

abstract model_to_device()

Moves the model to the correct device.

Return type

None

on_exception(exception)

Called when the trainer execution is interrupted by an exception.

Return type

None

on_predict_end()

Called when predict ends.

Return type

None

on_predict_start()

Called when predict begins.

Return type

None

on_test_end()

Called when test end.

Return type

None

on_test_start()

Called when test begins.

Return type

None

on_train_batch_start(batch, batch_idx)

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

Return type

None

on_train_end()

Called when train ends.

Return type

None

on_train_start()

Called when train begins.

Return type

None

on_validation_end()

Called when validation ends.

Return type

None

on_validation_start()

Called when validation begins.

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, closure, model=None, **kwargs)

Performs the actual optimizer step.

Parameters

• optimizer (Optimizer) – the optimizer performing the step

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

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

• **kwargs – Keyword arguments to optimizer.step

Return type

Any

post_backward(closure_loss)

Run after precision plugin executes backward.

Return type

None

pre_backward(closure_loss)

Run before precision plugin executes backward.

Return type

None

predict_step(*args, **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.

Parameters

dataloader (object) – iterable. Ideally of type: torch.utils.data.DataLoader

Return type

object

abstract reduce(tensor, group=None, reduce_op='mean')

Reduces the given tensor (e.g. across GPUs/processes).

Parameters

• tensor (Union[Tensor, Any]) – the tensor to sync and reduce

• group (Optional[Any]) – the process group to reduce

• reduce_op (Union[ReduceOp, str, None]) – the reduction operation. Defaults to ‘mean’. Can also be a string ‘sum’ or ReduceOp.

Return type

Union[Tensor, Any]

reduce_boolean_decision(decision, all=True)

Reduce a boolean decision across all processes.

Return type

bool

remove_checkpoint(filepath)

Remove checkpoint filepath from the filesystem.

Parameters

filepath (Union[str, Path]) – Path to checkpoint

Return type

None

save_checkpoint(checkpoint, filepath, storage_options=None)

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

Parameters

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

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

• storage_options (Optional[Any]) – parameter for how to save to storage, passed to CheckpointIO plugin

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

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(*args, **kwargs)

The actual test step.

See test_step() for more details

Return type

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

training_step(*args, **kwargs)

The actual training step.

See training_step() for more details

Return type

Union[Tensor, Dict[str, Any]]

validation_step(*args, **kwargs)

The actual validation step.

See validation_step() for more details

Return type

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

property handles_gradient_accumulation: bool

Whether the plugin handles gradient accumulation internally.

Return type

bool

abstract property is_global_zero: bool

Whether the current process is the rank zero process not only on the local node, but for all nodes.

Return type

bool

property lightning_module: Optional[LightningModule]

Returns the pure LightningModule without potential wrappers.

Return type

Optional[LightningModule]

property lightning_restore_optimizer: bool

Override to disable Lightning restoring optimizers/schedulers.

This is useful for plugins which manage restoring optimizers/schedulers.

Return type

bool

property model: Optional[Module]

Returns the potentially wrapped LightningModule.

Return type

Optional[Module]

property restore_checkpoint_after_setup: bool

Override to delay restoring from checkpoint till after the setup phase has completed. This is useful when the strategy requires all the setup hooks to run before loading checkpoint.

Return type

bool

Returns

If True, restore checkpoint after strategy setup.

abstract property root_device: device

Returns the root device.

Return type

device