pytorch_lightning.lite.LightningLite¶
- class pytorch_lightning.lite.LightningLite(accelerator=None, strategy=None, devices=None, num_nodes=1, precision=32, plugins=None, gpus=None, tpu_cores=None)[source]¶
Bases:
abc.ABC
Lite accelerates your PyTorch training or inference code with minimal changes required.
Automatic placement of models and data onto the device.
Automatic support for mixed and double precision (smaller memory footprint).
Seamless switching between hardware (CPU, GPU, TPU) and distributed training strategies (data-parallel training, sharded training, etc.).
Automated spawning of processes, no launch utilities required.
Multi-node support.
- Parameters
accelerator¶ (
Union
[str
,Accelerator
,None
]) – The hardware to run on. Possible choices are:"cpu"
,"gpu"
,"tpu"
,"auto"
.strategy¶ (
Union
[str
,TrainingTypePlugin
,None
]) – Strategy for how to run across multiple devices. Possible choices are:"dp"
,"ddp"
,"ddp_spawn"
,"deepspeed"
,"ddp_sharded"
.devices¶ (
Union
[int
,str
,List
[int
],None
]) – Number of devices to train on (int
), which GPUs to train on (list
orstr
), or"auto"
. The value applies per node.num_nodes¶ (
int
) – Number of GPU nodes for distributed training.precision¶ (
Union
[int
,str
]) – Double precision (64
), full precision (32
), half precision (16
), or bfloat16 precision ("bf16"
).plugins¶ (
Union
[TrainingTypePlugin
,PrecisionPlugin
,ClusterEnvironment
,CheckpointIO
,str
,List
[Union
[TrainingTypePlugin
,PrecisionPlugin
,ClusterEnvironment
,CheckpointIO
,str
]],None
]) – One or several custom pluginsgpus¶ (
Union
[int
,str
,List
[int
],None
]) – Provides the same function as thedevices
argument but impliesaccelerator="gpu"
.tpu_cores¶ (
Union
[int
,str
,List
[int
],None
]) – Provides the same function as thedevices
argument but impliesaccelerator="tpu"
.
- __init__(accelerator=None, strategy=None, devices=None, num_nodes=1, precision=32, plugins=None, gpus=None, tpu_cores=None)[source]¶
Initialize self. See help(type(self)) for accurate signature.
Methods
__init__
([accelerator, strategy, devices, …])Initialize self.
all_gather
(data[, group, sync_grads])Gather tensors or collections of tensors from multiple processes.
autocast
()A context manager to automatically convert operations for the chosen precision.
backward
(tensor, *args[, model])Replaces
loss.backward()
in your training loop.barrier
([name])Wait for all processes to enter this call.
broadcast
(obj[, src])- rtype
load
(filepath)Load a checkpoint from a file.
print
(*args, **kwargs)Print something only on the first process.
run
()All the code inside this run method gets accelerated by Lite.
save
(content, filepath)Save checkpoint contents to a file.
seed_everything
([seed, workers])Helper function to seed everything without explicitly importing Lightning.
setup
(model, *optimizers[, move_to_device])Setup a model and its optimizers for accelerated training.
setup_dataloaders
(*dataloaders[, …])Setup one or multiple dataloaders for accelerated training.
to_device
(obj)Move a
torch.nn.Module
or a collection of tensors to the current device, if it is not already on that device.Attributes
The current device this process runs on.
The global index of the current process across all devices and nodes.
Wether this rank is rank zero.
The index of the current process among the processes running on the local node.
The index of the current node.
The total number of processes running across all devices and nodes.
- all_gather(data, group=None, sync_grads=False)[source]¶
Gather tensors or collections of tensors from multiple processes.
- Parameters
data¶ (
Union
[Tensor
,Dict
,List
,Tuple
]) – int, float, tensor of shape (batch, …), or a (possibly nested) collection thereof.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 the all_gather operation
- Return type
- Returns
A tensor of shape (world_size, batch, …), or if the input was a collection the output will also be a collection with tensors of this shape.
- autocast()[source]¶
A context manager to automatically convert operations for the chosen precision.
Use this only if the forward method of your model does not cover all operations you wish to run with the chosen precision setting.
- backward(tensor, *args, model=None, **kwargs)[source]¶
Replaces
loss.backward()
in your training loop. Handles precision and automatically for you.- Parameters
tensor¶ (
Tensor
) – The tensor (loss) to back-propagate gradients from.*args¶ (
Any
) – Optional positional arguments passed to the underlying backward function.model¶ (
Optional
[_LiteModule
]) – Optional model instance for plugins that require the model for backward().**kwargs¶ (
Any
) – Optional named keyword arguments passed to the underlying backward function.
- Return type
Note
When using
strategy="deepspeed"
and multiple models were setup, it is required to pass in the model as argument here.
- barrier(name=None)[source]¶
Wait for all processes to enter this call. Use this to synchronize all parallel processes, but only if necessary, otherwise the overhead of synchronization will cause your program to slow down.
Example:
if self.global_rank == 0: # let process 0 download the dataset dataset.download_files() # let all processes wait before reading the dataset self.barrier() # now all processes can read the files and start training
- Return type
- load(filepath)[source]¶
Load a checkpoint from a file.
How and which processes load gets determined by the strategy
- print(*args, **kwargs)[source]¶
Print something only on the first process.
Arguments passed to this method are forwarded to the Python built-in
print()
function.- Return type
- abstract run()[source]¶
All the code inside this run method gets accelerated by Lite.
You can pass arbitrary arguments to this function when overriding it.
- Return type
- save(content, filepath)[source]¶
Save checkpoint contents to a file.
How and which processes save gets determined by the strategy. For example, the ddp strategy saves checkpoints only on process 0.
- static seed_everything(seed=None, workers=None)[source]¶
Helper function to seed everything without explicitly importing Lightning.
See
pytorch_lightning.seed_everything()
for more details.- Return type
- setup(model, *optimizers, move_to_device=True)[source]¶
Setup a model and its optimizers for accelerated training.
- Parameters
- Return type
- Returns
The tuple of the wrapped model and list of optimizers, in the same order they were passed in.
- setup_dataloaders(*dataloaders, replace_sampler=True, move_to_device=True)[source]¶
Setup one or multiple dataloaders for accelerated training. If you need different settings for each dataloader, call this method individually for each one.
- Parameters
*dataloaders¶ (
DataLoader
) – A single dataloader or a sequence of dataloaders.replace_sampler¶ (
bool
) – If setTrue
(default), automatically wraps or replaces the sampler on the dataloader(s) for distributed training. If you have a custom sampler defined, set this to this argument toFalse
.move_to_device¶ (
bool
) – If setTrue
(default), moves the data returned by the dataloader(s) automatially to the correct device. Set this toFalse
and alternatively useto_device()
manually on the returned data.
- Return type
- Returns
The wrapped dataloaders, in the same order they were passed in.
- to_device(obj)[source]¶
Move a
torch.nn.Module
or a collection of tensors to the current device, if it is not already on that device.
- property device: torch.device¶
The current device this process runs on.
Use this to create tensors directly on the device if needed.
- property global_rank: int¶
The global index of the current process across all devices and nodes.
- property is_global_zero: bool¶
Wether this rank is rank zero.
- property local_rank: int¶
The index of the current process among the processes running on the local node.
- property node_rank: int¶
The index of the current node.
- property world_size: int¶
The total number of processes running across all devices and nodes.