API Reference

accelerators

• Accelerator — The Accelerator base class for Lightning PyTorch.
• CPUAccelerator — Accelerator for CPU devices.
• CUDAAccelerator — Accelerator for NVIDIA CUDA devices.
• XLAAccelerator — Accelerator for XLA devices, normally TPUs.

callbacks

• BackboneFinetuning — Finetune a backbone model based on a learning rate user-defined scheduling.
• BaseFinetuning — This class implements the base logic for writing your own Finetuning Callback.
• BasePredictionWriter — Base class to implement how the predictions should be stored.
• BatchSizeFinder — Finds the largest batch size supported by a given model before encountering an out of memory (OOM) error.
• Callback — Abstract base class used to build new callbacks.
• DeviceStatsMonitor — Automatically monitors and logs device stats during training, validation and testing stage. DeviceStatsMonitor is a special callback as it requires a logger to passed as argument to the Trainer.
• EarlyStopping — Monitor a metric and stop training when it stops improving.
• GradientAccumulationScheduler — Change gradient accumulation factor according to scheduling.
• LambdaCallback — Create a simple callback on the fly using lambda functions.
• LearningRateFinder — The LearningRateFinder callback enables the user to do a range test of good initial learning rates, to reduce the amount of guesswork in picking a good starting learning rate.
• LearningRateMonitor — Automatically monitor and logs learning rate for learning rate schedulers during training.
• ModelCheckpoint — Save the model after every epoch by monitoring a quantity. Every logged metrics are passed to the Logger for the version it gets saved in the same directory as the checkpoint.
• ModelPruning — Model pruning Callback, using PyTorch’s prune utilities. This callback is responsible of pruning networks parameters during training.
• ModelSummary — Generates a summary of all layers in a LightningModule.
• OnExceptionCheckpoint — Used to save a checkpoint on exception.
• ProgressBar — The base class for progress bars in Lightning. It is a Callback that keeps track of the batch progress in the Trainer. You should implement your highly custom progress bars with this as the base class
• RichModelSummary — Generates a summary of all layers in a LightningModule with rich text formatting.
• RichProgressBar — Create a progress bar with rich text formatting.
• StochasticWeightAveraging — Implements the Stochastic Weight Averaging (SWA) Callback to average a model.
• SpikeDetection
• ThroughputMonitor — Computes and logs throughput with the Throughput
• Timer — The Timer callback tracks the time spent in the training, validation, and test loops and interrupts the Trainer if the given time limit for the training loop is reached.
• TQDMProgressBar — This is the default progress bar used by Lightning. It prints to stdout using the tqdm package and shows up to four different bars:
• WeightAveraging — A callback that updates an averaged model for Stochastic Weight Averaging (SWA) or Exponential Moving Average (EMA) after each training step.

cli

• LightningCLI — Implementation of a configurable command line tool for pytorch-lightning.
• LightningArgumentParser — Extension of jsonargparse’s ArgumentParser for pytorch-lightning.
• SaveConfigCallback — Saves a LightningCLI config to the log_dir when training starts.

core

• CheckpointHooks — Hooks to be used with Checkpointing.
• DataHooks — Hooks to be used for data related stuff.
• ModelHooks — Hooks to be used in LightningModule.
• LightningDataModule — A DataModule standardizes the training, val, test splits, data preparation and transforms. The main advantage is consistent data splits, data preparation and transforms across models.
• LightningModule
• HyperparametersMixin
• LightningOptimizer — This class is used to wrap the user optimizers and handle properly the backward and optimizer_step logic across accelerators, AMP, accumulate_grad_batches.

loggers

• logger — Dummy logger for internal use.
• litlogger — Logger that enables remote experiment tracking, logging, and artifact management on lightning.ai.
• comet — Track your parameters, metrics, source code and more using Comet.
• csv_logs — Log to local file system in yaml and CSV format.
• mlflow — Log using MLflow.
• tensorboard — Log to local or remote file system in TensorBoard format.
• wandb — Log using Weights and Biases.

plugins

precision

• DeepSpeedPrecision — Precision plugin for DeepSpeed integration.
• DoublePrecision — Plugin for training with double (torch.float64) precision.
• HalfPrecision — Plugin for training with half precision.
• FSDPPrecision — Precision plugin for training with Fully Sharded Data Parallel (FSDP).
• MixedPrecision — Plugin for Automatic Mixed Precision (AMP) training with torch.autocast.
• Precision — Base class for all plugins handling the precision-specific parts of the training.
• XLAPrecision — Plugin for training with XLA.
• TransformerEnginePrecision — Plugin for training with fp8 precision via nvidia’s Transformer Engine.
• BitsandbytesPrecision — Plugin for quantizing weights with bitsandbytes.

environments

• ClusterEnvironment — Specification of a cluster environment.
• KubeflowEnvironment — Environment for distributed training using the PyTorchJob operator from Kubeflow.
• LightningEnvironment — The default environment used by Lightning for a single node or free cluster (not managed).
• LSFEnvironment — An environment for running on clusters managed by the LSF resource manager.
• MPIEnvironment — An environment for running on clusters with processes created through MPI.
• SLURMEnvironment — Cluster environment for training on a cluster managed by SLURM.
• TorchElasticEnvironment — Environment for fault-tolerant and elastic training with torchelastic
• XLAEnvironment — Cluster environment for training on a TPU Pod with the PyTorch/XLA library.

io

• AsyncCheckpointIO — AsyncCheckpointIO enables saving the checkpoints asynchronously in a thread.
• CheckpointIO — Interface to save/load checkpoints as they are saved through the Strategy.
• TorchCheckpointIO — CheckpointIO that utilizes torch.save() and torch.load() to save and load checkpoints respectively, common for most use cases.
• XLACheckpointIO — CheckpointIO that utilizes xm.save to save checkpoints for TPU training strategies.

others

• LayerSync — Abstract base class for creating plugins that wrap layers of a model with synchronization logic for multiprocessing.
• TorchSyncBatchNorm — A plugin that wraps all batch normalization layers of a model with synchronization logic for multiprocessing.

profiler

• AdvancedProfiler — This profiler uses Python’s cProfiler to record more detailed information about time spent in each function call recorded during a given action.
• PassThroughProfiler — This class should be used when you don’t want the (small) overhead of profiling.
• Profiler — If you wish to write a custom profiler, you should inherit from this class.
• PyTorchProfiler — This profiler uses PyTorch’s Autograd Profiler and lets you inspect the cost of different operators inside your model - both on the CPU and GPU.
• SimpleProfiler — This profiler simply records the duration of actions (in seconds) and reports the mean duration of each action and the total time spent over the entire training run.
• XLAProfiler — XLA Profiler will help you debug and optimize training workload performance for your models using Cloud TPU performance tools.

trainer

• Trainer — Customize every aspect of training via flags.

strategies

• DDPStrategy — Strategy for multi-process single-device training on one or multiple nodes.
• DeepSpeedStrategy — Provides capabilities to run training using the DeepSpeed library, with training optimizations for large billion parameter models. For more information: DeepSpeed.
• FSDPStrategy — Strategy for Fully Sharded Data Parallel provided by torch.distributed.
• ModelParallelStrategy — Enables user-defined parallelism applied to a model.
• ParallelStrategy — Strategy for training with multiple processes in parallel.
• SingleDeviceStrategy — Strategy that handles communication on a single device.
• SingleDeviceXLAStrategy — Strategy for training on a single XLA device.
• Strategy — Base class for all strategies that change the behaviour of the training, validation and test- loop.
• XLAStrategy — Strategy for training multiple TPU devices using the torch_xla.distributed.xla_multiprocessing.spawn() method.

tuner

• Tuner — Tuner class to tune your model.

utilities

• combined_loader — Combines different iterables under specific sampling modes.
• data — Unpack a batch to find a torch.Tensor.
• deepspeed — Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict file that can be loaded with torch.load(file) + load_state_dict() and used for training without DeepSpeed. It gets copied into
• memory — Garbage collection Torch (CUDA) memory.
• model_summary
• parsing — Extended dictionary accessible with dot notation.
• rank_zero
• seed — A context manager that resets the global random state on exit to what it was before entering.
• warnings