Shortcuts

Strategy

class lightning_fabric.strategies.Strategy(accelerator=None, checkpoint_io=None, precision=None)[source]

Bases: abc.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)[source]

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

abstract all_reduce(tensor, group=None, reduce_op='mean')[source]

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]

backward(tensor, module, *args, **kwargs)[source]

Forwards backward-calls to the precision plugin.

Return type

None

abstract barrier(name=None)[source]

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)[source]

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

Return type

Any

abstract broadcast(obj, src=0)[source]

Broadcasts an object to all processes.

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

  • src (int) – source rank

Return type

TypeVar(TBroadcast)

get_module_state_dict(module)[source]

Returns model state.

Return type

Dict[str, Union[Any, Tensor]]

get_optimizer_state(optimizer)[source]

Returns state of an optimizer.

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

Return type

Dict[str, Tensor]

abstract module_to_device(module)[source]

Moves the model to the correct device.

Return type

None

optimizer_step(optimizer, **kwargs)[source]

Performs the actual optimizer step.

Parameters
  • optimizer (Optimizable) – the optimizer performing the step

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

Return type

Any

process_dataloader(dataloader)[source]

Wraps the dataloader if necessary.

Parameters

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

Return type

DataLoader

reduce_boolean_decision(decision, all=True)[source]

Reduce a boolean decision across all processes.

Return type

bool

remove_checkpoint(filepath)[source]

Remove checkpoint filepath from the filesystem.

Parameters

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

Return type

None

save_checkpoint(checkpoint, filepath, storage_options=None)[source]

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_environment()[source]

Setup any processes or distributed connections.

This must be called by the framework at the beginning of every process, before any distributed communication takes place.

Return type

None

setup_module(module)[source]

Performs setup for the model, e.g., by wrapping it by another class.

Return type

Module

setup_module_and_optimizers(module, optimizers)[source]

Set up a model and multiple optimizers together.

The returned objects are expected to be in the same order they were passed in. The default implementation will call setup_module() and setup_optimizer() on the inputs.

Return type

Tuple[Module, List[Optimizer]]

setup_optimizer(optimizer)[source]

Performs setup for the optimizer, e.g., by wrapping it by another class.

Return type

Optimizer

teardown()[source]

This method is called to teardown the training process.

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

Return type

None

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

abstract property root_device: torch.device

Returns the root device.

Return type

device


© Copyright Copyright (c) 2018-2023, Lightning AI et al...

Built with Sphinx using a theme provided by Read the Docs.