Trainer¶
- class lightning.pytorch.trainer.trainer.Trainer(*, accelerator='auto', strategy='auto', devices='auto', num_nodes=1, precision=None, logger=None, callbacks=None, fast_dev_run=False, max_epochs=None, min_epochs=None, max_steps=-1, min_steps=None, max_time=None, limit_train_batches=None, limit_val_batches=None, limit_test_batches=None, limit_predict_batches=None, overfit_batches=0.0, val_check_interval=None, check_val_every_n_epoch=1, num_sanity_val_steps=None, log_every_n_steps=None, enable_checkpointing=None, enable_progress_bar=None, enable_model_summary=None, accumulate_grad_batches=1, gradient_clip_val=None, gradient_clip_algorithm=None, deterministic=None, benchmark=None, inference_mode=True, use_distributed_sampler=True, profiler=None, detect_anomaly=False, barebones=False, plugins=None, sync_batchnorm=False, reload_dataloaders_every_n_epochs=0, default_root_dir=None)[source]¶
Bases:
object
Customize every aspect of training via flags.
- Parameters:
accelerator¶ (
Union
[str
,Accelerator
]) – Supports passing different accelerator types (“cpu”, “gpu”, “tpu”, “ipu”, “hpu”, “mps”, “auto”) as well as custom accelerator instances.strategy¶ (
Union
[str
,Strategy
]) – Supports different training strategies with aliases as well custom strategies. Default:"auto"
.devices¶ (
Union
[List
[int
],str
,int
]) – The devices to use. Can be set to a positive number (int or str), a sequence of device indices (list or str), the value-1
to indicate all available devices should be used, or"auto"
for automatic selection based on the chosen accelerator. Default:"auto"
.num_nodes¶ (
int
) – Number of GPU nodes for distributed training. Default:1
.precision¶ (
Union
[Literal
[64
,32
,16
],Literal
['transformer-engine'
,'transformer-engine-float16'
,'16-true'
,'16-mixed'
,'bf16-true'
,'bf16-mixed'
,'32-true'
,'64-true'
],Literal
['64'
,'32'
,'16'
,'bf16'
],None
]) – Double precision (64, ‘64’ or ‘64-true’), full precision (32, ‘32’ or ‘32-true’), 16bit mixed precision (16, ‘16’, ‘16-mixed’) or bfloat16 mixed precision (‘bf16’, ‘bf16-mixed’). Can be used on CPU, GPU, TPUs, HPUs or IPUs. Default:'32-true'
.logger¶ (
Union
[Logger
,Iterable
[Logger
],bool
,None
]) – Logger (or iterable collection of loggers) for experiment tracking. ATrue
value uses the defaultTensorBoardLogger
if it is installed, otherwiseCSVLogger
.False
will disable logging. If multiple loggers are provided, local files (checkpoints, profiler traces, etc.) are saved in thelog_dir
of the first logger. Default:True
.callbacks¶ (
Union
[List
[Callback
],Callback
,None
]) – Add a callback or list of callbacks. Default:None
.fast_dev_run¶ (
Union
[int
,bool
]) – Runs n if set ton
(int) else 1 if set toTrue
batch(es) of train, val and test to find any bugs (ie: a sort of unit test). Default:False
.max_epochs¶ (
Optional
[int
]) – Stop training once this number of epochs is reached. Disabled by default (None). If both max_epochs and max_steps are not specified, defaults tomax_epochs = 1000
. To enable infinite training, setmax_epochs = -1
.min_epochs¶ (
Optional
[int
]) – Force training for at least these many epochs. Disabled by default (None).max_steps¶ (
int
) – Stop training after this number of steps. Disabled by default (-1). Ifmax_steps = -1
andmax_epochs = None
, will default tomax_epochs = 1000
. To enable infinite training, setmax_epochs
to-1
.min_steps¶ (
Optional
[int
]) – Force training for at least these number of steps. Disabled by default (None
).max_time¶ (
Union
[str
,timedelta
,Dict
[str
,int
],None
]) – Stop training after this amount of time has passed. Disabled by default (None
). The time duration can be specified in the format DD:HH:MM:SS (days, hours, minutes seconds), as adatetime.timedelta
, or a dictionary with keys that will be passed todatetime.timedelta
.limit_train_batches¶ (
Union
[int
,float
,None
]) – How much of training dataset to check (float = fraction, int = num_batches). Default:1.0
.limit_val_batches¶ (
Union
[int
,float
,None
]) – How much of validation dataset to check (float = fraction, int = num_batches). Default:1.0
.limit_test_batches¶ (
Union
[int
,float
,None
]) – How much of test dataset to check (float = fraction, int = num_batches). Default:1.0
.limit_predict_batches¶ (
Union
[int
,float
,None
]) – How much of prediction dataset to check (float = fraction, int = num_batches). Default:1.0
.overfit_batches¶ (
Union
[int
,float
]) – Overfit a fraction of training/validation data (float) or a set number of batches (int). Default:0.0
.val_check_interval¶ (
Union
[int
,float
,None
]) – How often to check the validation set. Pass afloat
in the range [0.0, 1.0] to check after a fraction of the training epoch. Pass anint
to check after a fixed number of training batches. Anint
value can only be higher than the number of training batches whencheck_val_every_n_epoch=None
, which validates after everyN
training batches across epochs or during iteration-based training. Default:1.0
.check_val_every_n_epoch¶ (
Optional
[int
]) – Perform a validation loop every after every N training epochs. IfNone
, validation will be done solely based on the number of training batches, requiringval_check_interval
to be an integer value. Default:1
.num_sanity_val_steps¶ (
Optional
[int
]) – Sanity check runs n validation batches before starting the training routine. Set it to -1 to run all batches in all validation dataloaders. Default:2
.log_every_n_steps¶ (
Optional
[int
]) – How often to log within steps. Default:50
.enable_checkpointing¶ (
Optional
[bool
]) – IfTrue
, enable checkpointing. It will configure a default ModelCheckpoint callback if there is no user-defined ModelCheckpoint incallbacks
. Default:True
.enable_progress_bar¶ (
Optional
[bool
]) – Whether to enable to progress bar by default. Default:True
.enable_model_summary¶ (
Optional
[bool
]) – Whether to enable model summarization by default. Default:True
.accumulate_grad_batches¶ (
int
) – Accumulates gradients over k batches before stepping the optimizer. Default: 1.gradient_clip_val¶ (
Union
[int
,float
,None
]) – The value at which to clip gradients. Passinggradient_clip_val=None
disables gradient clipping. If using Automatic Mixed Precision (AMP), the gradients will be unscaled before. Default:None
.gradient_clip_algorithm¶ (
Optional
[str
]) – The gradient clipping algorithm to use. Passgradient_clip_algorithm="value"
to clip by value, andgradient_clip_algorithm="norm"
to clip by norm. By default it will be set to"norm"
.deterministic¶ (
Union
[bool
,Literal
['warn'
],None
]) – IfTrue
, sets whether PyTorch operations must use deterministic algorithms. Set to"warn"
to use deterministic algorithms whenever possible, throwing warnings on operations that don’t support deterministic mode. If not set, defaults toFalse
. Default:None
.benchmark¶ (
Optional
[bool
]) – The value (True
orFalse
) to settorch.backends.cudnn.benchmark
to. The value fortorch.backends.cudnn.benchmark
set in the current session will be used (False
if not manually set). Ifdeterministic
is set toTrue
, this will default toFalse
. Override to manually set a different value. Default:None
.inference_mode¶ (
bool
) – Whether to usetorch.inference_mode()
ortorch.no_grad()
during evaluation (validate
/test
/predict
).use_distributed_sampler¶ (
bool
) – Whether to wrap the DataLoader’s sampler withtorch.utils.data.DistributedSampler
. If not specified this is toggled automatically for strategies that require it. By default, it will addshuffle=True
for the train sampler andshuffle=False
for validation/test/predict samplers. If you want to disable this logic, you can passFalse
and add your own distributed sampler in the dataloader hooks. IfTrue
and a distributed sampler was already added, Lightning will not replace the existing one. For iterable-style datasets, we don’t do this automatically.profiler¶ (
Union
[Profiler
,str
,None
]) – To profile individual steps during training and assist in identifying bottlenecks. Default:None
.detect_anomaly¶ (
bool
) – Enable anomaly detection for the autograd engine. Default:False
.barebones¶ (
bool
) – Whether to run in “barebones mode”, where all features that may impact raw speed are disabled. This is meant for analyzing the Trainer overhead and is discouraged during regular training runs. The following features are deactivated:enable_checkpointing
,logger
,enable_progress_bar
,log_every_n_steps
,enable_model_summary
,num_sanity_val_steps
,fast_dev_run
,detect_anomaly
,profiler
,log()
,log_dict()
.plugins¶ (
Union
[PrecisionPlugin
,ClusterEnvironment
,CheckpointIO
,LayerSync
,str
,List
[Union
[PrecisionPlugin
,ClusterEnvironment
,CheckpointIO
,LayerSync
,str
]],None
]) – Plugins allow modification of core behavior like ddp and amp, and enable custom lightning plugins. Default:None
.sync_batchnorm¶ (
bool
) – Synchronize batch norm layers between process groups/whole world. Default:False
.reload_dataloaders_every_n_epochs¶ (
int
) – Set to a positive integer to reload dataloaders every n epochs. Default:0
.default_root_dir¶ (
Union
[str
,Path
,None
]) – Default path for logs and weights when no logger/ckpt_callback passed. Default:os.getcwd()
. Can be remote file paths such as s3://mybucket/path or ‘hdfs://path/’
- Raises:
TypeError – If
gradient_clip_val
is not an int or float.MisconfigurationException – If
gradient_clip_algorithm
is invalid. Iftrack_grad_norm
is not a positive number or inf.
- fit(model, train_dataloaders=None, val_dataloaders=None, datamodule=None, ckpt_path=None)[source]¶
Runs the full optimization routine.
- Parameters:
model¶ (
LightningModule
) – Model to fit.train_dataloaders¶ (
Union
[Any
,LightningDataModule
,None
]) – An iterable or collection of iterables specifying training samples. Alternatively, aLightningDataModule
that defines thetrain_dataloader
hook.val_dataloaders¶ (
Optional
[Any
]) – An iterable or collection of iterables specifying validation samples.datamodule¶ (
Optional
[LightningDataModule
]) – ALightningDataModule
that defines thetrain_dataloader
hook.ckpt_path¶ (
Optional
[str
]) – Path/URL of the checkpoint from which training is resumed. Could also be one of two special keywords"last"
and"hpc"
. If there is no checkpoint file at the path, an exception is raised.
- Raises:
TypeError – If
model
is notLightningModule
for torch version less than 2.0.0 and ifmodel
is notLightningModule
ortorch._dynamo.OptimizedModule
for torch versions greater than or equal to 2.0.0 .
For more information about multiple dataloaders, see this section. :rtype:
None
- init_module(empty_init=None)[source]¶
Tensors that you instantiate under this context manager will be created on the device right away and have the right data type depending on the precision setting in the Trainer.
The parameters and tensors get created on the device and with the right data type right away without wasting memory being allocated unnecessarily. The automatic device placement under this context manager is only supported with PyTorch 2.0 and newer.
- Parameters:
empty_init¶ (
Optional
[bool
]) – Whether to initialize the model with empty weights (uninitialized memory). IfNone
, the strategy will decide. Some strategies may not support all options. Set this toTrue
if you are loading a checkpoint into a large model. Requires torch >= 1.13.- Return type:
- predict(model=None, dataloaders=None, datamodule=None, return_predictions=None, ckpt_path=None)[source]¶
Run inference on your data. This will call the model forward function to compute predictions. Useful to perform distributed and batched predictions. Logging is disabled in the predict hooks.
- Parameters:
model¶ (
Optional
[LightningModule
]) – The model to predict with.dataloaders¶ (
Union
[Any
,LightningDataModule
,None
]) – An iterable or collection of iterables specifying predict samples. Alternatively, aLightningDataModule
that defines thepredict_dataloader
hook.datamodule¶ (
Optional
[LightningDataModule
]) – ALightningDataModule
that defines thepredict_dataloader
hook.return_predictions¶ (
Optional
[bool
]) – Whether to return predictions.True
by default except when an accelerator that spawns processes is used (not supported).ckpt_path¶ (
Optional
[str
]) – Either"best"
,"last"
,"hpc"
or path to the checkpoint you wish to predict. IfNone
and the model instance was passed, use the current weights. Otherwise, the best model checkpoint from the previoustrainer.fit
call will be loaded if a checkpoint callback is configured.
For more information about multiple dataloaders, see this section.
- Return type:
- Returns:
Returns a list of dictionaries, one for each provided dataloader containing their respective predictions.
- Raises:
TypeError – If no
model
is passed and there was noLightningModule
passed in the previous run. Ifmodel
passed is not LightningModule or torch._dynamo.OptimizedModule.MisconfigurationException – If both
dataloaders
anddatamodule
are passed. Pass only one of these.RuntimeError – If a compiled
model
is passed and the strategy is not supported.
See Lightning inference section for more.
- print(*args, **kwargs)[source]¶
Print something only on the first process. If running on multiple machines, it will print from the first process in each machine.
Arguments passed to this method are forwarded to the Python built-in
print()
function.- Return type:
- save_checkpoint(filepath, weights_only=False, storage_options=None)[source]¶
Runs routine to create a checkpoint.
This method needs to be called on all processes in case the selected strategy is handling distributed checkpointing.
- Parameters:
- Raises:
AttributeError – If the model is not attached to the Trainer before calling this method.
- Return type:
- test(model=None, dataloaders=None, ckpt_path=None, verbose=True, datamodule=None)[source]¶
Perform one evaluation epoch over the test set. It’s separated from fit to make sure you never run on your test set until you want to.
- Parameters:
model¶ (
Optional
[LightningModule
]) – The model to test.dataloaders¶ (
Union
[Any
,LightningDataModule
,None
]) – An iterable or collection of iterables specifying test samples. Alternatively, aLightningDataModule
that defines thetest_dataloader
hook.ckpt_path¶ (
Optional
[str
]) – Either"best"
,"last"
,"hpc"
or path to the checkpoint you wish to test. IfNone
and the model instance was passed, use the current weights. Otherwise, the best model checkpoint from the previoustrainer.fit
call will be loaded if a checkpoint callback is configured.datamodule¶ (
Optional
[LightningDataModule
]) – ALightningDataModule
that defines thetest_dataloader
hook.
For more information about multiple dataloaders, see this section.
- Return type:
- Returns:
List of dictionaries with metrics logged during the test phase, e.g., in model- or callback hooks like
test_step()
etc. The length of the list corresponds to the number of test dataloaders used.- Raises:
TypeError – If no
model
is passed and there was noLightningModule
passed in the previous run. Ifmodel
passed is not LightningModule or torch._dynamo.OptimizedModule.MisconfigurationException – If both
dataloaders
anddatamodule
are passed. Pass only one of these.RuntimeError – If a compiled
model
is passed and the strategy is not supported.
- validate(model=None, dataloaders=None, ckpt_path=None, verbose=True, datamodule=None)[source]¶
Perform one evaluation epoch over the validation set.
- Parameters:
model¶ (
Optional
[LightningModule
]) – The model to validate.dataloaders¶ (
Union
[Any
,LightningDataModule
,None
]) – An iterable or collection of iterables specifying validation samples. Alternatively, aLightningDataModule
that defines theval_dataloader
hook.ckpt_path¶ (
Optional
[str
]) – Either"best"
,"last"
,"hpc"
or path to the checkpoint you wish to validate. IfNone
and the model instance was passed, use the current weights. Otherwise, the best model checkpoint from the previoustrainer.fit
call will be loaded if a checkpoint callback is configured.datamodule¶ (
Optional
[LightningDataModule
]) – ALightningDataModule
that defines theval_dataloader
hook.
For more information about multiple dataloaders, see this section.
- Return type:
- Returns:
List of dictionaries with metrics logged during the validation phase, e.g., in model- or callback hooks like
validation_step()
etc. The length of the list corresponds to the number of validation dataloaders used.- Raises:
TypeError – If no
model
is passed and there was noLightningModule
passed in the previous run. Ifmodel
passed is not LightningModule or torch._dynamo.OptimizedModule.MisconfigurationException – If both
dataloaders
anddatamodule
are passed. Pass only one of these.RuntimeError – If a compiled
model
is passed and the strategy is not supported.
- property callback_metrics: Dict[str, Tensor]¶
The metrics available to callbacks.
def training_step(self, batch, batch_idx): self.log("a_val", 2.0) callback_metrics = trainer.callback_metrics assert callback_metrics["a_val"] == 2.0
- property checkpoint_callback: Optional[Checkpoint]¶
The first
ModelCheckpoint
callback in the Trainer.callbacks list, orNone
if it doesn’t exist.
- property checkpoint_callbacks: List[Checkpoint]¶
A list of all instances of
ModelCheckpoint
found in the Trainer.callbacks list.
- property ckpt_path: Optional[Union[str, Path]]¶
Set to the path/URL of a checkpoint loaded via
fit()
,validate()
,test()
, orpredict()
.None
otherwise.
- property default_root_dir: str¶
The default location to save artifacts of loggers, checkpoints etc.
It is used as a fallback if logger or checkpoint callback do not define specific save paths.
- property early_stopping_callback: Optional[EarlyStopping]¶
The first
EarlyStopping
callback in the Trainer.callbacks list, orNone
if it doesn’t exist.
- property early_stopping_callbacks: List[EarlyStopping]¶
A list of all instances of
EarlyStopping
found in the Trainer.callbacks list.
- property estimated_stepping_batches: Union[int, float]¶
The estimated number of batches that will
optimizer.step()
during training.This accounts for gradient accumulation and the current trainer configuration. This might sets up your training dataloader if hadn’t been set up already.
def configure_optimizers(self): optimizer = ... stepping_batches = self.trainer.estimated_stepping_batches scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-3, total_steps=stepping_batches) return [optimizer], [scheduler]
- Raises:
MisconfigurationException – If estimated stepping batches cannot be computed due to different accumulate_grad_batches at different epochs.
- property global_step: int¶
The number of optimizer steps taken (does not reset each epoch).
This includes multiple optimizers (if enabled).
- property is_global_zero: bool¶
Whether this process is the global zero in multi-node training.
def training_step(self, batch, batch_idx): if self.trainer.is_global_zero: print("in node 0, accelerator 0")
- property log_dir: Optional[str]¶
The directory for the current experiment. Use this to save images to, etc…
Note
You must call this on all processes. Failing to do so will cause your program to stall forever.
def training_step(self, batch, batch_idx): img = ... save_img(img, self.trainer.log_dir)
- property logged_metrics: Dict[str, Tensor]¶
The metrics sent to the loggers.
This includes metrics logged via
log()
with thelogger
argument set.
- property loggers: List[Logger]¶
The list of
Logger
used.for logger in trainer.loggers: logger.log_metrics({"foo": 1.0})
- property model: Optional[Module]¶
The LightningModule, but possibly wrapped into DataParallel or DistributedDataParallel.
To access the pure LightningModule, use
lightning_module()
instead.
- property num_predict_batches: List[Union[int, float]]¶
The number of prediction batches that will be used during
trainer.predict()
.
- property num_sanity_val_batches: List[Union[int, float]]¶
The number of validation batches that will be used during the sanity-checking part of
trainer.fit()
.
- property num_test_batches: List[Union[int, float]]¶
The number of test batches that will be used during
trainer.test()
.
- property num_training_batches: Union[int, float]¶
The number of training batches that will be used during
trainer.fit()
.
- property num_val_batches: List[Union[int, float]]¶
The number of validation batches that will be used during
trainer.fit()
ortrainer.validate()
.
- property predict_dataloaders: Optional[Any]¶
The prediction dataloader(s) used during
trainer.predict()
.
- property progress_bar_callback: Optional[ProgressBar]¶
An instance of
ProgressBar
found in the Trainer.callbacks list, orNone
if one doesn’t exist.
- property progress_bar_metrics: Dict[str, float]¶
The metrics sent to the progress bar.
This includes metrics logged via
log()
with theprog_bar
argument set.
- property received_sigterm: bool¶
Whether a
signal.SIGTERM
signal was received.For example, this can be checked to exit gracefully.