LightningModule — PyTorch Lightning 2.5.1.post0 documentation (original) (raw)

class lightning.pytorch.core.LightningModule(*args, **kwargs)[source]

Bases: _DeviceDtypeModuleMixin, HyperparametersMixin, ModelHooks, DataHooks, CheckpointHooks, Module

all_gather(data, group=None, sync_grads=False)[source]

Gather tensors or collections of tensors from multiple processes.

This method needs to be called on all processes and the tensors need to have the same shape across all processes, otherwise your program will stall forever.

Parameters:

Return type:

Union[Tensor, dict, list, tuple]

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. For the special case where world_size is 1, no additional dimension is added to the tensor(s).

backward(loss, *args, **kwargs)[source]

Called to perform backward on the loss returned in training_step(). Override this hook with your own implementation if you need to.

Parameters:

loss (Tensor) – The loss tensor returned by training_step(). If gradient accumulation is used, the loss here holds the normalized value (scaled by 1 / accumulation steps).

Return type:

None

Example:

def backward(self, loss): loss.backward()

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)[source]

Handles gradient clipping internally.

Note

Parameters:

Return type:

None

configure_callbacks()[source]

Configure model-specific callbacks. When the model gets attached, e.g., when .fit() or .test() gets called, the list or a callback returned here will be merged with the list of callbacks passed to the Trainer’scallbacks argument. If a callback returned here has the same type as one or several callbacks already present in the Trainer’s callbacks list, it will take priority and replace them. In addition, Lightning will make sure ModelCheckpoint callbacks run last.

Return type:

Union[Sequence[Callback], Callback]

Returns:

A callback or a list of callbacks which will extend the list of callbacks in the Trainer.

Example:

def configure_callbacks(self): early_stop = EarlyStopping(monitor="val_acc", mode="max") checkpoint = ModelCheckpoint(monitor="val_loss") return [early_stop, checkpoint]

configure_gradient_clipping(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)[source]

Perform gradient clipping for the optimizer parameters. Called before optimizer_step().

Parameters:

Return type:

None

Example:

def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm): # Implement your own custom logic to clip gradients # You can call self.clip_gradients with your settings: self.clip_gradients( optimizer, gradient_clip_val=gradient_clip_val, gradient_clip_algorithm=gradient_clip_algorithm )

configure_optimizers()[source]

Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you’d need one. But in the case of GANs or similar you might have multiple. Optimization with multiple optimizers only works in the manual optimization mode.

Return type:

Union[Optimizer, Sequence[Optimizer], tuple[Sequence[Optimizer], Sequence[Union[LRScheduler, ReduceLROnPlateau, LRSchedulerConfig]]], OptimizerConfig, OptimizerLRSchedulerConfig, Sequence[OptimizerConfig], Sequence[OptimizerLRSchedulerConfig], None]

Returns:

Any of these 6 options.

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = { # REQUIRED: The scheduler instance "scheduler": lr_scheduler, # The unit of the scheduler's step size, could also be 'step'. # 'epoch' updates the scheduler on epoch end whereas 'step' # updates it after a optimizer update. "interval": "epoch", # How many epochs/steps should pass between calls to # scheduler.step(). 1 corresponds to updating the learning # rate after every epoch/step. "frequency": 1, # Metric to monitor for schedulers like ReduceLROnPlateau "monitor": "val_loss", # If set to True, will enforce that the value specified 'monitor' # is available when the scheduler is updated, thus stopping # training if not found. If set to False, it will only produce a warning "strict": True, # If using the LearningRateMonitor callback to monitor the # learning rate progress, this keyword can be used to specify # a custom logged name "name": None, }

When there are schedulers in which the .step() method is conditioned on a value, such as thetorch.optim.lr_scheduler.ReduceLROnPlateau scheduler, Lightning requires that thelr_scheduler_config contains the keyword "monitor" set to the metric name that the scheduler should be conditioned on.

The ReduceLROnPlateau scheduler requires a monitor

def configure_optimizers(self): optimizer = Adam(...) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": ReduceLROnPlateau(optimizer, ...), "monitor": "metric_to_track", "frequency": "indicates how often the metric is updated", # If "monitor" references validation metrics, then "frequency" should be set to a # multiple of "trainer.check_val_every_n_epoch". }, }

In the case of two optimizers, only one using the ReduceLROnPlateau scheduler

def configure_optimizers(self): optimizer1 = Adam(...) optimizer2 = SGD(...) scheduler1 = ReduceLROnPlateau(optimizer1, ...) scheduler2 = LambdaLR(optimizer2, ...) return ( { "optimizer": optimizer1, "lr_scheduler": { "scheduler": scheduler1, "monitor": "metric_to_track", }, }, {"optimizer": optimizer2, "lr_scheduler": scheduler2}, )

Metrics can be made available to monitor by simply logging it usingself.log('metric_to_track', metric_val) in your LightningModule.

Note

Some things to know:

forward(*args, **kwargs)[source]

Same as torch.nn.Module.forward().

Parameters:

Return type:

Any

Returns:

Your model’s output

freeze()[source]

Freeze all params for inference.

Example:

model = MyLightningModule(...) model.freeze()

Return type:

None

load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs)[source]

Primary way of loading a model from a checkpoint. When Lightning saves a checkpoint it stores the arguments passed to __init__ in the checkpoint under "hyper_parameters".

Any arguments specified through **kwargs will override args stored in "hyper_parameters".

Parameters:

Return type:

Self

Returns:

LightningModule instance with loaded weights and hyperparameters (if available).

Note

load_from_checkpoint is a class method. You should use your LightningModule class to call it instead of the LightningModule instance, or aTypeError will be raised.

Note

To ensure all layers can be loaded from the checkpoint, this function will callconfigure_model() directly after instantiating the model if this hook is overridden in your LightningModule. However, note that load_from_checkpoint does not support loading sharded checkpoints, and you may run out of memory if the model is too large. In this case, consider loading through the Trainer via .fit(ckpt_path=...).

Example:

load weights without mapping ...

model = MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')

or load weights mapping all weights from GPU 1 to GPU 0 ...

map_location = {'cuda:1':'cuda:0'} model = MyLightningModule.load_from_checkpoint( 'path/to/checkpoint.ckpt', map_location=map_location )

or load weights and hyperparameters from separate files.

model = MyLightningModule.load_from_checkpoint( 'path/to/checkpoint.ckpt', hparams_file='/path/to/hparams_file.yaml' )

override some of the params with new values

model = MyLightningModule.load_from_checkpoint( PATH, num_layers=128, pretrained_ckpt_path=NEW_PATH, )

predict

pretrained_model.eval() pretrained_model.freeze() y_hat = pretrained_model(x)

log(name, value, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, metric_attribute=None, rank_zero_only=False)[source]

Log a key, value pair.

Example:

self.log('train_loss', loss)

The default behavior per hook is documented here: Automatic Logging.

Parameters:

Return type:

None

log_dict(dictionary, prog_bar=False, logger=None, on_step=None, on_epoch=None, reduce_fx='mean', enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True, batch_size=None, rank_zero_only=False)[source]

Log a dictionary of values at once.

Example:

values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n} self.log_dict(values)

Parameters:

Return type:

None

lr_scheduler_step(scheduler, metric)[source]

Override this method to adjust the default way the Trainer calls each scheduler. By default, Lightning calls step() and as shown in the example for each scheduler based on its interval.

Parameters:

Return type:

None

Examples:

DEFAULT

def lr_scheduler_step(self, scheduler, metric): if metric is None: scheduler.step() else: scheduler.step(metric)

Alternative way to update schedulers if it requires an epoch value

def lr_scheduler_step(self, scheduler, metric): scheduler.step(epoch=self.current_epoch)

lr_schedulers()[source]

Returns the learning rate scheduler(s) that are being used during training. Useful for manual optimization.

Return type:

Union[None, list[Union[LRScheduler, ReduceLROnPlateau]], LRScheduler, ReduceLROnPlateau]

Returns:

A single scheduler, or a list of schedulers in case multiple ones are present, or None if no schedulers were returned in configure_optimizers().

manual_backward(loss, *args, **kwargs)[source]

Call this directly from your training_step() when doing optimizations manually. By using this, Lightning can ensure that all the proper scaling gets applied when using mixed precision.

See manual optimization for more examples.

Example:

def training_step(...): opt = self.optimizers() loss = ... opt.zero_grad() # automatically applies scaling, etc... self.manual_backward(loss) opt.step()

Parameters:

Return type:

None

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)[source]

Override this method to adjust the default way the Trainer calls the optimizer.

By default, Lightning calls step() and zero_grad() as shown in the example. This method (and zero_grad()) won’t be called during the accumulation phase whenTrainer(accumulate_grad_batches != 1). Overriding this hook has no benefit with manual optimization.

Parameters:

Return type:

None

Examples:

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure): # Add your custom logic to run directly before optimizer.step()

optimizer.step(closure=optimizer_closure)

# Add your custom logic to run directly after `optimizer.step()`

optimizer_zero_grad(epoch, batch_idx, optimizer)[source]

Override this method to change the default behaviour of optimizer.zero_grad().

Parameters:

Return type:

None

Examples:

DEFAULT

def optimizer_zero_grad(self, epoch, batch_idx, optimizer): optimizer.zero_grad()

Set gradients to None instead of zero to improve performance (not required on torch>=2.0.0).

def optimizer_zero_grad(self, epoch, batch_idx, optimizer): optimizer.zero_grad(set_to_none=True)

See torch.optim.Optimizer.zero_grad() for the explanation of the above example.

optimizers(use_pl_optimizer=True)[source]

Returns the optimizer(s) that are being used during training. Useful for manual optimization.

Parameters:

use_pl_optimizer (bool) – If True, will wrap the optimizer(s) in aLightningOptimizer for automatic handling of precision, profiling, and counting of step calls for proper logging and checkpointing. It specifically wraps thestep method and custom optimizers that don’t have this method are not supported.

Return type:

Union[Optimizer, LightningOptimizer, _FabricOptimizer, list[Optimizer], list[LightningOptimizer], list[_FabricOptimizer]]

Returns:

A single optimizer, or a list of optimizers in case multiple ones are present.

predict_step(*args, **kwargs)[source]

Step function called during predict(). By default, it callsforward(). Override to add any processing logic.

The predict_step() is used to scale inference on multi-devices.

To prevent an OOM error, it is possible to use BasePredictionWritercallback to write the predictions to disk or database after each batch or on epoch end.

The BasePredictionWriter should be used while using a spawn based accelerator. This happens for Trainer(strategy="ddp_spawn")or training on 8 TPU cores with Trainer(accelerator="tpu", devices=8) as predictions won’t be returned.

Parameters:

Return type:

Any

Returns:

Predicted output (optional).

Example

class MyModel(LightningModule):

def predict_step(self, batch, batch_idx, dataloader_idx=0):
    return self(batch)

dm = ... model = MyModel() trainer = Trainer(accelerator="gpu", devices=2) predictions = trainer.predict(model, dm)

print(*args, **kwargs)[source]

Prints only from process 0. Use this in any distributed mode to log only once.

Parameters:

Return type:

None

Example:

def forward(self, x): self.print(x, 'in forward')

test_step(*args, **kwargs)[source]

Operates on a single batch of data from the test set. In this step you’d normally generate examples or calculate anything of interest such as accuracy.

Parameters:

Return type:

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

Returns:

if you have one test dataloader:

def test_step(self, batch, batch_idx): ...

if you have multiple test dataloaders:

def test_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

CASE 1: A single test dataset

def test_step(self, batch, batch_idx): x, y = batch

# implement your own
out = self(x)
loss = self.loss(out, y)

# log 6 example images
# or generated text... or whatever
sample_imgs = x[:6]
grid = torchvision.utils.make_grid(sample_imgs)
self.logger.experiment.add_image('example_images', grid, 0)

# calculate acc
labels_hat = torch.argmax(out, dim=1)
test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

# log the outputs!
self.log_dict({'test_loss': loss, 'test_acc': test_acc})

If you pass in multiple test dataloaders, test_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

CASE 2: multiple test dataloaders

def test_step(self, batch, batch_idx, dataloader_idx=0): # dataloader_idx tells you which dataset this is. ...

Note

If you don’t need to test you don’t need to implement this method.

Note

When the test_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of the test epoch, the model goes back to training mode and gradients are enabled.

to_onnx(file_path, input_sample=None, **kwargs)[source]

Saves the model in ONNX format.

Parameters:

Return type:

None

Example:

class SimpleModel(LightningModule): def init(self): super().init() self.l1 = torch.nn.Linear(in_features=64, out_features=4)

def forward(self, x):
    return torch.relu(self.l1(x.view(x.size(0), -1)

model = SimpleModel() input_sample = torch.randn(1, 64) model.to_onnx("export.onnx", input_sample, export_params=True)

to_torchscript(file_path=None, method='script', example_inputs=None, **kwargs)[source]

By default compiles the whole model to a ScriptModule. If you want to use tracing, please provided the argument method='trace' and make sure that either the example_inputs argument is provided, or the model has example_input_array set. If you would like to customize the modules that are scripted you should override this method. In case you want to return multiple modules, we recommend using a dictionary.

Parameters:

Note

Example:

class SimpleModel(LightningModule): def init(self): super().init() self.l1 = torch.nn.Linear(in_features=64, out_features=4)

def forward(self, x):
    return torch.relu(self.l1(x.view(x.size(0), -1)))

model = SimpleModel() model.to_torchscript(file_path="model.pt")

torch.jit.save(model.to_torchscript( file_path="model_trace.pt", method='trace', example_inputs=torch.randn(1, 64)) )

Return type:

Union[ScriptModule, dict[str, ScriptModule]]

Returns:

This LightningModule as a torchscript, regardless of whether file_path is defined or not.

toggle_optimizer(optimizer)[source]

Makes sure only the gradients of the current optimizer’s parameters are calculated in the training step to prevent dangling gradients in multiple-optimizer setup.

It works with untoggle_optimizer() to make sure param_requires_grad_state is properly reset.

Parameters:

optimizer (Union[Optimizer, LightningOptimizer]) – The optimizer to toggle.

Return type:

None

training_step(*args, **kwargs)[source]

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Parameters:

Return type:

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

Returns:

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx): x, y, z = batch out = self.encoder(x) loss = self.loss(out, x) return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def init(self): super().init() self.automatic_optimization = False

Multiple optimizers (e.g.: GANs)

def training_step(self, batch, batch_idx): opt1, opt2 = self.optimizers()

# do training_step with encoder
...
opt1.step()
# do training_step with decoder
...
opt2.step()

Note

When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

unfreeze()[source]

Unfreeze all parameters for training.

model = MyLightningModule(...) model.unfreeze()

Return type:

None

untoggle_optimizer(optimizer)[source]

Resets the state of required gradients that were toggled with toggle_optimizer().

Parameters:

optimizer (Union[Optimizer, LightningOptimizer]) – The optimizer to untoggle.

Return type:

None

validation_step(*args, **kwargs)[source]

Operates on a single batch of data from the validation set. In this step you’d might generate examples or calculate anything of interest like accuracy.

Parameters:

Return type:

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

Returns:

if you have one val dataloader:

def validation_step(self, batch, batch_idx): ...

if you have multiple val dataloaders:

def validation_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

CASE 1: A single validation dataset

def validation_step(self, batch, batch_idx): x, y = batch

# implement your own
out = self(x)
loss = self.loss(out, y)

# log 6 example images
# or generated text... or whatever
sample_imgs = x[:6]
grid = torchvision.utils.make_grid(sample_imgs)
self.logger.experiment.add_image('example_images', grid, 0)

# calculate acc
labels_hat = torch.argmax(out, dim=1)
val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

# log the outputs!
self.log_dict({'val_loss': loss, 'val_acc': val_acc})

If you pass in multiple val dataloaders, validation_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

CASE 2: multiple validation dataloaders

def validation_step(self, batch, batch_idx, dataloader_idx=0): # dataloader_idx tells you which dataset this is. ...

Note

If you don’t need to validate you don’t need to implement this method.

Note

When the validation_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, the model goes back to training mode and gradients are enabled.

property automatic_optimization_: bool_

If set to False you are responsible for calling .backward(), .step(), .zero_grad().

property current_epoch_: int_

The current epoch in the Trainer, or 0 if not attached.

property device_mesh_: Optional[DeviceMesh]_

Strategies like ModelParallelStrategy will create a device mesh that can be accessed in theconfigure_model() hook to parallelize the LightningModule.

property example_input_array_: Optional[Union[Tensor, tuple, dict]]_

The example input array is a specification of what the module can consume in the forward() method. The return type is interpreted as follows:

property global_rank_: int_

The index of the current process across all nodes and devices.

property global_step_: int_

Total training batches seen across all epochs.

If no Trainer is attached, this propery is 0.

property local_rank_: int_

The index of the current process within a single node.

property logger_: Optional[Union[Logger, Logger]]_

Reference to the logger object in the Trainer.

property loggers_: Union[list[lightning.pytorch.loggers.logger.Logger], list[lightning.fabric.loggers.logger.Logger]]_

Reference to the list of loggers in the Trainer.

property on_gpu_: bool_

Returns True if this model is currently located on a GPU.

Useful to set flags around the LightningModule for different CPU vs GPU behavior.

property strict_loading_: bool_

Determines how Lightning loads this model using .load_state_dict(…, strict=model.strict_loading).