--- title: LightningModule description: >- When you convert to use Lightning, the code IS NOT abstracted - just organized. canonical: 'https://lightning.ai/docs/pytorch/stable/core-api/lightning_module' --- # LightningModule {#lightningmodule} A `LightningModule` organizes your PyTorch code into 6 sections: - Initialization (`__init__` and `setup`). - Train Loop (`training_step`) - Validation Loop (`validation_step`) - Test Loop (`test_step`) - Prediction Loop (`predict_step`) - Optimizers and LR Schedulers (`configure_optimizers`) When you convert to use Lightning, the code IS NOT abstracted - just organized. All the other code that's not in the `LightningModule` has been automated for you by the `Trainer`. ```python net = MyLightningModuleNet() trainer = Trainer() trainer.fit(net) ``` There are no `.cuda()` or `.to(device)` calls required. Lightning does these for you. ```python # don't do in Lightning x = torch.Tensor(2, 3) x = x.cuda() x = x.to(device) # do this instead x = x # leave it alone! # or to init a new tensor new_x = torch.Tensor(2, 3) new_x = new_x.to(x) ``` When running under a distributed strategy, Lightning handles the distributed sampler for you by default. ```python # Don't do in Lightning... data = MNIST(...) sampler = DistributedSampler(data) DataLoader(data, sampler=sampler) # do this instead data = MNIST(...) DataLoader(data) ``` A `LightningModule` is a `torch.nn.Module` but with added functionality. Use it as such! ```python net = Net.load_from_checkpoint(PATH) net.freeze() out = net(x) ``` Thus, to use Lightning, you just need to organize your code which takes about 30 minutes, (and let's be real, you probably should do anyway). ## Starter Example {#starter-example} Here are the only required methods. ```python import lightning as L import torch from lightning.pytorch.demos import Transformer class LightningTransformer(L.LightningModule): def __init__(self, vocab_size): super().__init__() self.model = Transformer(vocab_size=vocab_size) def forward(self, inputs, target): return self.model(inputs, target) def training_step(self, batch, batch_idx): inputs, target = batch output = self(inputs, target) loss = torch.nn.functional.nll_loss(output, target.view(-1)) return loss def configure_optimizers(self): return torch.optim.SGD(self.model.parameters(), lr=0.1) ``` Which you can train by doing: ```python from lightning.pytorch.demos import WikiText2 from torch.utils.data import DataLoader dataset = WikiText2() dataloader = DataLoader(dataset) model = LightningTransformer(vocab_size=dataset.vocab_size) trainer = L.Trainer(fast_dev_run=100) trainer.fit(model=model, train_dataloaders=dataloader) ``` The LightningModule has many convenient methods, but the core ones you need to know about are: | Name | Description | | --- | --- | | `__init__` and `setup` | Define initialization here | | `forward` | To run data through your model only (separate from `training_step`) | | `training_step` | the complete training step | | `validation_step` | the complete validation step | | `test_step` | the complete test step | | `predict_step` | the complete prediction step | | `configure_optimizers` | define optimizers and LR schedulers | ## Training {#training} ### Training Loop {#training-loop} To activate the training loop, override the `training_step` method. ```python class LightningTransformer(L.LightningModule): def __init__(self, vocab_size): super().__init__() self.model = Transformer(vocab_size=vocab_size) def training_step(self, batch, batch_idx): inputs, target = batch output = self.model(inputs, target) loss = torch.nn.functional.nll_loss(output, target.view(-1)) return loss ``` Under the hood, Lightning does the following (pseudocode): ```python # enable gradient calculation torch.set_grad_enabled(True) for batch_idx, batch in enumerate(train_dataloader): loss = training_step(batch, batch_idx) # clear gradients optimizer.zero_grad() # backward loss.backward() # update parameters optimizer.step() ``` ### Train Epoch-level Metrics {#train-epoch-level-metrics} If you want to calculate epoch-level metrics and log them, use `log`. ```python def training_step(self, batch, batch_idx): inputs, target = batch output = self.model(inputs, target) loss = torch.nn.functional.nll_loss(output, target.view(-1)) # logs metrics for each training_step, # and the average across the epoch, to the progress bar and logger self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True) return loss ``` The `log` method automatically reduces the requested metrics across a complete epoch and devices. Here's the pseudocode of what it does under the hood: ```python outs = [] for batch_idx, batch in enumerate(train_dataloader): # forward loss = training_step(batch, batch_idx) outs.append(loss.detach()) # clear gradients optimizer.zero_grad() # backward loss.backward() # update parameters optimizer.step() # note: in reality, we do this incrementally, instead of keeping all outputs in memory epoch_metric = torch.mean(torch.stack(outs)) ``` ### Train Epoch-level Operations {#train-epoch-level-operations} In the case that you need to make use of all the outputs from each `training_step`, override the `on_train_epoch_end` method. ```python class LightningTransformer(L.LightningModule): def __init__(self, vocab_size): super().__init__() self.model = Transformer(vocab_size=vocab_size) self.training_step_outputs = [] def training_step(self, batch, batch_idx): inputs, target = batch output = self.model(inputs, target) loss = torch.nn.functional.nll_loss(output, target.view(-1)) preds = ... self.training_step_outputs.append(preds) return loss def on_train_epoch_end(self): all_preds = torch.stack(self.training_step_outputs) # do something with all preds ... self.training_step_outputs.clear() # free memory ``` ## Validation {#validation} ### Validation Loop {#validation-loop} To activate the validation loop while training, override the `validation_step` method. ```python class LightningTransformer(L.LightningModule): def validation_step(self, batch, batch_idx): inputs, target = batch output = self.model(inputs, target) loss = F.cross_entropy(y_hat, y) self.log("val_loss", loss) ``` Under the hood, Lightning does the following (pseudocode): ```python # ... for batch_idx, batch in enumerate(train_dataloader): loss = model.training_step(batch, batch_idx) loss.backward() # ... if validate_at_some_point: # disable grads + batchnorm + dropout torch.set_grad_enabled(False) model.eval() # ----------------- VAL LOOP --------------- for val_batch_idx, val_batch in enumerate(val_dataloader): val_out = model.validation_step(val_batch, val_batch_idx) # ----------------- VAL LOOP --------------- # enable grads + batchnorm + dropout torch.set_grad_enabled(True) model.train() ``` You can also run just the validation loop on your validation dataloaders by overriding `validation_step` and calling `validate`. ```python model = LightningTransformer(vocab_size=dataset.vocab_size) trainer = L.Trainer() trainer.validate(model) ``` > **Note** > > It is recommended to validate on single device to ensure each sample/batch gets evaluated exactly once. > This is helpful to make sure benchmarking for research papers is done the right way. Otherwise, in a > multi-device setting, samples could occur duplicated when `DistributedSampler` > is used, for eg. with `strategy="ddp"`. It replicates some samples on some devices to make sure all devices have > same batch size in case of uneven inputs. ### Validation Epoch-level Metrics {#validation-epoch-level-metrics} In the case that you need to make use of all the outputs from each `validation_step`, override the `on_validation_epoch_end` method. Note that this method is called before `on_train_epoch_end`. ```python class LightningTransformer(L.LightningModule): def __init__(self, vocab_size): super().__init__() self.model = Transformer(vocab_size=vocab_size) self.validation_step_outputs = [] def validation_step(self, batch, batch_idx): x, y = batch inputs, target = batch output = self.model(inputs, target) loss = torch.nn.functional.nll_loss(output, target.view(-1)) pred = ... self.validation_step_outputs.append(pred) return pred def on_validation_epoch_end(self): all_preds = torch.stack(self.validation_step_outputs) # do something with all preds ... self.validation_step_outputs.clear() # free memory ``` ## Testing {#testing} ### Test Loop {#test-loop} The process for enabling a test loop is the same as the process for enabling a validation loop. Please refer to the section above for details. For this you need to override the `test_step` method. The only difference is that the test loop is only called when `test` is used. ```python model = LightningTransformer(vocab_size=dataset.vocab_size) dataloader = DataLoader(dataset) trainer = L.Trainer() trainer.fit(model=model, train_dataloaders=dataloader) # automatically loads the best weights for you trainer.test(model) ``` There are two ways to call `test()`: ```python # call after training trainer = L.Trainer() trainer.fit(model=model, train_dataloaders=dataloader) # automatically auto-loads the best weights from the previous run trainer.test(dataloaders=test_dataloaders) # or call with pretrained model model = LightningTransformer.load_from_checkpoint(PATH) dataset = WikiText2() test_dataloader = DataLoader(dataset) trainer = L.Trainer() trainer.test(model, dataloaders=test_dataloader) ``` > **Note** > > `WikiText2` is used in a manner that does not create a train, test, val split. This is done for illustrative purposes only. > A proper split can be created in `lightning.pytorch.core.LightningModule.setup` or `lightning.pytorch.core.LightningDataModule.setup`. > **Note** > > It is recommended to validate on single device to ensure each sample/batch gets evaluated exactly once. > This is helpful to make sure benchmarking for research papers is done the right way. Otherwise, in a > multi-device setting, samples could occur duplicated when `DistributedSampler` > is used, for eg. with `strategy="ddp"`. It replicates some samples on some devices to make sure all devices have > same batch size in case of uneven inputs. ## Inference {#inference} ### Prediction Loop {#prediction-loop} By default, the `predict_step` method runs the `forward` method. In order to customize this behaviour, simply override the `predict_step` method. For the example let's override `predict_step`: ```python class LightningTransformer(L.LightningModule): def __init__(self, vocab_size): super().__init__() self.model = Transformer(vocab_size=vocab_size) def predict_step(self, batch): inputs, target = batch return self.model(inputs, target) ``` Under the hood, Lightning does the following (pseudocode): ```python # disable grads + batchnorm + dropout torch.set_grad_enabled(False) model.eval() all_preds = [] for batch_idx, batch in enumerate(predict_dataloader): pred = model.predict_step(batch, batch_idx) all_preds.append(pred) ``` There are two ways to call `predict()`: ```python # call after training trainer = L.Trainer() trainer.fit(model=model, train_dataloaders=dataloader) # automatically auto-loads the best weights from the previous run predictions = trainer.predict(dataloaders=predict_dataloader) # or call with pretrained model model = LightningTransformer.load_from_checkpoint(PATH) dataset = WikiText2() test_dataloader = DataLoader(dataset) trainer = L.Trainer() predictions = trainer.predict(model, dataloaders=test_dataloader) ``` ### Inference in Research {#inference-in-research} If you want to perform inference with the system, you can add a `forward` method to the LightningModule. > **Note** > > When using forward, you are responsible to call `eval` and use the `no_grad` context manager. ```python class LightningTransformer(L.LightningModule): def __init__(self, vocab_size): super().__init__() self.model = Transformer(vocab_size=vocab_size) def forward(self, batch): inputs, target = batch return self.model(inputs, target) def training_step(self, batch, batch_idx): inputs, target = batch output = self.model(inputs, target) loss = torch.nn.functional.nll_loss(output, target.view(-1)) return loss def configure_optimizers(self): return torch.optim.SGD(self.model.parameters(), lr=0.1) model = LightningTransformer(vocab_size=dataset.vocab_size) model.eval() with torch.no_grad(): batch = dataloader.dataset[0] pred = model(batch) ``` The advantage of adding a forward is that in complex systems, you can do a much more involved inference procedure, such as text generation: ```python class Seq2Seq(L.LightningModule): def forward(self, x): embeddings = self(x) hidden_states = self.encoder(embeddings) for h in hidden_states: # decode ... return decoded ``` In the case where you want to scale your inference, you should be using `predict_step`. ```python class Autoencoder(L.LightningModule): def forward(self, x): return self.decoder(x) def predict_step(self, batch, batch_idx, dataloader_idx=0): # this calls forward return self(batch) data_module = ... model = Autoencoder() trainer = Trainer(accelerator="gpu", devices=2) trainer.predict(model, data_module) ``` ### Inference in Production {#inference-in-production} For cases like production, you might want to iterate different models inside a LightningModule. ```python from torchmetrics.functional import accuracy class ClassificationTask(L.LightningModule): def __init__(self, model): super().__init__() self.model = model def training_step(self, batch, batch_idx): x, y = batch y_hat = self.model(x) loss = F.cross_entropy(y_hat, y) return loss def validation_step(self, batch, batch_idx): loss, acc = self._shared_eval_step(batch, batch_idx) metrics = {"val_acc": acc, "val_loss": loss} self.log_dict(metrics) return metrics def test_step(self, batch, batch_idx): loss, acc = self._shared_eval_step(batch, batch_idx) metrics = {"test_acc": acc, "test_loss": loss} self.log_dict(metrics) return metrics def _shared_eval_step(self, batch, batch_idx): x, y = batch y_hat = self.model(x) loss = F.cross_entropy(y_hat, y) acc = accuracy(y_hat, y) return loss, acc def predict_step(self, batch, batch_idx, dataloader_idx=0): x, y = batch y_hat = self.model(x) return y_hat def configure_optimizers(self): return torch.optim.Adam(self.model.parameters(), lr=0.02) ``` Then pass in any arbitrary model to be fit with this task ```python for model in [resnet50(), vgg16(), BidirectionalRNN()]: task = ClassificationTask(model) trainer = Trainer(accelerator="gpu", devices=2) trainer.fit(task, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader) ``` Tasks can be arbitrarily complex such as implementing GAN training, self-supervised or even RL. ```python class GANTask(L.LightningModule): def __init__(self, generator, discriminator): super().__init__() self.generator = generator self.discriminator = discriminator ... ``` When used like this, the model can be separated from the Task and thus used in production without needing to keep it in a `LightningModule`. The following example shows how you can run inference in the Python runtime: ```python task = ClassificationTask(model) trainer = Trainer(accelerator="gpu", devices=2) trainer.fit(task, train_dataloader, val_dataloader) trainer.save_checkpoint("best_model.ckpt") # use model after training or load weights and drop into the production system model = ClassificationTask.load_from_checkpoint("best_model.ckpt") x = ... model.eval() with torch.no_grad(): y_hat = model(x) ``` Check out [Inference in Production](/docs/pytorch/2.6.6/reference/deploy/production.md) guide to learn about the possible ways to perform inference in production. ## Save Hyperparameters {#save-hyperparameters} Often times we train many versions of a model. You might share that model or come back to it a few months later at which point it is very useful to know how that model was trained (i.e.: what learning rate, neural network, etc...). Lightning has a standardized way of saving the information for you in checkpoints and YAML files. The goal here is to improve readability and reproducibility. ### save_hyperparameters {#save-hyperparameters} Use `save_hyperparameters` within your `LightningModule`'s `__init__` method. It will enable Lightning to store all the provided arguments under the `self.hparams` attribute. These hyperparameters will also be stored within the model checkpoint, which simplifies model re-instantiation after training. ```python class LitMNIST(L.LightningModule): def __init__(self, layer_1_dim=128, learning_rate=1e-2): super().__init__() # call this to save (layer_1_dim=128, learning_rate=1e-4) to the checkpoint self.save_hyperparameters() # equivalent self.save_hyperparameters("layer_1_dim", "learning_rate") # Now possible to access layer_1_dim from hparams self.hparams.layer_1_dim ``` In addition, loggers that support it will automatically log the contents of `self.hparams`. ### Excluding hyperparameters {#excluding-hyperparameters} By default, every parameter of the `__init__` method will be considered a hyperparameter to the LightningModule. However, sometimes some parameters need to be excluded from saving, for example when they are not serializable. Those parameters should be provided back when reloading the LightningModule. In this case, exclude them explicitly: ```python class LitMNIST(L.LightningModule): def __init__(self, loss_fx, generator_network, layer_1_dim=128): super().__init__() self.layer_1_dim = layer_1_dim self.loss_fx = loss_fx # call this to save only (layer_1_dim=128) to the checkpoint self.save_hyperparameters("layer_1_dim") # equivalent self.save_hyperparameters(ignore=["loss_fx", "generator_network"]) ``` ### load_from_checkpoint {#load-from-checkpoint} LightningModules that have hyperparameters automatically saved with `save_hyperparameters` can conveniently be loaded and instantiated directly from a checkpoint with `load_from_checkpoint`: ```python # to load specify the other args model = LitMNIST.load_from_checkpoint(PATH, loss_fx=torch.nn.SomeOtherLoss, generator_network=MyGenerator()) ``` If parameters were excluded, they need to be provided at the time of loading: ```python # the excluded parameters were `loss_fx` and `generator_network` model = LitMNIST.load_from_checkpoint(PATH, loss_fx=torch.nn.SomeOtherLoss, generator_network=MyGenerator()) ``` ## Child Modules {#child-modules} Research projects tend to test different approaches to the same dataset. This is very easy to do in Lightning with inheritance. For example, imagine we now want to train an `AutoEncoder` to use as a feature extractor for images. The only things that change in the `LitAutoEncoder` model are the init, forward, training, validation and test step. ```python class Encoder(torch.nn.Module): ... class Decoder(torch.nn.Module): ... class AutoEncoder(torch.nn.Module): def __init__(self): super().__init__() self.encoder = Encoder() self.decoder = Decoder() def forward(self, x): return self.decoder(self.encoder(x)) class LitAutoEncoder(LightningModule): def __init__(self, auto_encoder): super().__init__() self.auto_encoder = auto_encoder self.metric = torch.nn.MSELoss() def forward(self, x): return self.auto_encoder.encoder(x) def training_step(self, batch, batch_idx): x, _ = batch x_hat = self.auto_encoder(x) loss = self.metric(x, x_hat) return loss def validation_step(self, batch, batch_idx): self._shared_eval(batch, batch_idx, "val") def test_step(self, batch, batch_idx): self._shared_eval(batch, batch_idx, "test") def _shared_eval(self, batch, batch_idx, prefix): x, _ = batch x_hat = self.auto_encoder(x) loss = self.metric(x, x_hat) self.log(f"{prefix}_loss", loss) ``` and we can train this using the `Trainer`: ```python auto_encoder = AutoEncoder() lightning_module = LitAutoEncoder(auto_encoder) trainer = Trainer() trainer.fit(lightning_module, train_dataloader, val_dataloader) ``` And remember that the forward method should define the practical use of a `LightningModule`. In this case, we want to use the `LitAutoEncoder` to extract image representations: ```python some_images = torch.Tensor(32, 1, 28, 28) representations = lightning_module(some_images) ``` ## LightningModule API {#lightningmodule-api} ### Methods {#methods} #### all_gather {#all-gather} ``lightning.pytorch.core.module.LightningModule.all_gather`` #### configure_callbacks {#configure-callbacks} ``lightning.pytorch.core.module.LightningModule.configure_callbacks`` #### configure_optimizers {#configure-optimizers} ``lightning.pytorch.core.module.LightningModule.configure_optimizers`` #### forward {#forward} ``lightning.pytorch.core.module.LightningModule.forward`` #### freeze {#freeze} ``lightning.pytorch.core.module.LightningModule.freeze`` #### log {#log} ``lightning.pytorch.core.module.LightningModule.log`` #### log_dict {#log-dict} ``lightning.pytorch.core.module.LightningModule.log_dict`` #### lr_schedulers {#lr-schedulers} ``lightning.pytorch.core.module.LightningModule.lr_schedulers`` #### manual_backward {#manual-backward} ``lightning.pytorch.core.module.LightningModule.manual_backward`` #### optimizers {#optimizers} ``lightning.pytorch.core.module.LightningModule.optimizers`` #### print {#print} ``lightning.pytorch.core.module.LightningModule.print`` #### predict_step {#predict-step} ``lightning.pytorch.core.module.LightningModule.predict_step`` #### save_hyperparameters {#save-hyperparameters} ``lightning.pytorch.core.module.LightningModule.save_hyperparameters`` #### toggle_optimizer {#toggle-optimizer} ``lightning.pytorch.core.module.LightningModule.toggle_optimizer`` #### test_step {#test-step} ``lightning.pytorch.core.module.LightningModule.test_step`` #### to_onnx {#to-onnx} ``lightning.pytorch.core.module.LightningModule.to_onnx`` #### to_torchscript {#to-torchscript} ``lightning.pytorch.core.module.LightningModule.to_torchscript`` #### training_step {#training-step} ``lightning.pytorch.core.module.LightningModule.training_step`` #### unfreeze {#unfreeze} ``lightning.pytorch.core.module.LightningModule.unfreeze`` #### untoggle_optimizer {#untoggle-optimizer} ``lightning.pytorch.core.module.LightningModule.untoggle_optimizer`` #### validation_step {#validation-step} ``lightning.pytorch.core.module.LightningModule.validation_step`` ### Properties {#properties} These are properties available in a LightningModule. #### current_epoch {#current-epoch} The number of epochs run. ```python def training_step(self, batch, batch_idx): if self.current_epoch == 0: ... ``` #### device {#device} The device the module is on. Use it to keep your code device agnostic. ```python def training_step(self, batch, batch_idx): z = torch.rand(2, 3, device=self.device) ``` #### global_rank {#global-rank} The `global_rank` is the index of the current process across all nodes and devices. Lightning will perform some operations such as logging, weight checkpointing only when `global_rank=0`. You usually do not need to use this property, but it is useful to know how to access it if needed. ```python def training_step(self, batch, batch_idx): if self.global_rank == 0: # do something only once across all the nodes ... ``` #### global_step {#global-step} The number of optimizer steps taken (does not reset each epoch). This includes multiple optimizers (if enabled). ```python def training_step(self, batch, batch_idx): self.logger.experiment.log_image(..., step=self.global_step) ``` #### hparams {#hparams} The arguments passed through `LightningModule.__init__()` and saved by calling `save_hyperparameters` could be accessed by the `hparams` attribute. ```python def __init__(self, learning_rate): self.save_hyperparameters() def configure_optimizers(self): return Adam(self.parameters(), lr=self.hparams.learning_rate) ``` #### logger {#logger} The current logger being used (tensorboard or other supported logger) ```python def training_step(self, batch, batch_idx): # the generic logger (same no matter if tensorboard or other supported logger) self.logger # the particular logger tensorboard_logger = self.logger.experiment ``` #### loggers {#loggers} The list of loggers currently being used by the Trainer. ```python def training_step(self, batch, batch_idx): # List of Logger objects loggers = self.loggers for logger in loggers: logger.log_metrics({"foo": 1.0}) ``` #### local_rank {#local-rank} The `local_rank` is the index of the current process across all the devices for the current node. You usually do not need to use this property, but it is useful to know how to access it if needed. For example, if using 10 machines (or nodes), the GPU at index 0 on each machine has local_rank = 0. ```python def training_step(self, batch, batch_idx): if self.local_rank == 0: # do something only once across each node ... ``` #### precision {#precision} The type of precision used: ```python def training_step(self, batch, batch_idx): if self.precision == "16-true": ... ``` #### trainer {#trainer} Pointer to the trainer ```python def training_step(self, batch, batch_idx): max_steps = self.trainer.max_steps any_flag = self.trainer.any_flag ``` #### prepare_data_per_node {#prepare-data-per-node} If set to `True` will call `prepare_data()` on LOCAL_RANK=0 for every node. If set to `False` will only call from NODE_RANK=0, LOCAL_RANK=0. ```python class LitModel(LightningModule): def __init__(self): super().__init__() self.prepare_data_per_node = True ``` #### automatic_optimization {#automatic-optimization} When set to `False`, Lightning does not automate the optimization process. This means you are responsible for handling your optimizers. However, we do take care of precision and any accelerators used. See [manual optimization](/docs/pytorch/2.6.6/reference/common/optimization.md#manual-optimization) for details. ```python def __init__(self): self.automatic_optimization = False def training_step(self, batch, batch_idx): opt = self.optimizers(use_pl_optimizer=True) loss = ... opt.zero_grad() self.manual_backward(loss) opt.step() ``` Manual optimization is most useful for research topics like reinforcement learning, sparse coding, and GAN research. It is required when you are using 2+ optimizers because with automatic optimization, you can only use one optimizer. ```python def __init__(self): self.automatic_optimization = False def training_step(self, batch, batch_idx): # access your optimizers with use_pl_optimizer=False. Default is True opt_a, opt_b = self.optimizers(use_pl_optimizer=True) gen_loss = ... opt_a.zero_grad() self.manual_backward(gen_loss) opt_a.step() disc_loss = ... opt_b.zero_grad() self.manual_backward(disc_loss) opt_b.step() ``` #### example_input_array {#example-input-array} Set and access example_input_array, which basically represents a single batch. ```python def __init__(self): self.example_input_array = ... self.generator = ... def on_train_epoch_end(self): # generate some images using the example_input_array gen_images = self.generator(self.example_input_array) ``` ### Hooks {#hooks} This is the pseudocode to describe the structure of `fit`. The inputs and outputs of each function are not represented for simplicity. Please check each function's API reference for more information. ```python # runs on every device: devices can be GPUs, TPUs, ... def fit(self): configure_callbacks() if local_rank == 0: prepare_data() setup("fit") configure_model() configure_optimizers() on_fit_start() # the sanity check runs here on_train_start() for epoch in epochs: fit_loop() on_train_end() on_fit_end() teardown("fit") def fit_loop(): torch.set_grad_enabled(True) on_train_epoch_start() for batch_idx, batch in enumerate(train_dataloader()): on_train_batch_start() on_before_batch_transfer() transfer_batch_to_device() on_after_batch_transfer() out = training_step() on_before_zero_grad() optimizer_zero_grad() on_before_backward() backward() on_after_backward() on_before_optimizer_step() configure_gradient_clipping() optimizer_step() on_train_batch_end(out, batch, batch_idx) if should_check_val: val_loop() on_train_epoch_end() def val_loop(): on_validation_model_eval() # calls `model.eval()` torch.set_grad_enabled(False) on_validation_start() on_validation_epoch_start() for batch_idx, batch in enumerate(val_dataloader()): on_validation_batch_start(batch, batch_idx) batch = on_before_batch_transfer(batch) batch = transfer_batch_to_device(batch) batch = on_after_batch_transfer(batch) out = validation_step(batch, batch_idx) on_validation_batch_end(out, batch, batch_idx) on_validation_epoch_end() on_validation_end() # set up for train on_validation_model_train() # calls `model.train()` torch.set_grad_enabled(True) ``` #### backward {#backward} ``lightning.pytorch.core.module.LightningModule.backward`` #### on_before_backward {#on-before-backward} ``lightning.pytorch.core.module.LightningModule.on_before_backward`` #### on_after_backward {#on-after-backward} ``lightning.pytorch.core.module.LightningModule.on_after_backward`` #### on_before_zero_grad {#on-before-zero-grad} ``lightning.pytorch.core.module.LightningModule.on_before_zero_grad`` #### on_fit_start {#on-fit-start} ``lightning.pytorch.core.module.LightningModule.on_fit_start`` #### on_fit_end {#on-fit-end} ``lightning.pytorch.core.module.LightningModule.on_fit_end`` #### on_load_checkpoint {#on-load-checkpoint} ``lightning.pytorch.core.module.LightningModule.on_load_checkpoint`` #### on_save_checkpoint {#on-save-checkpoint} ``lightning.pytorch.core.module.LightningModule.on_save_checkpoint`` #### load_from_checkpoint {#load-from-checkpoint} ``lightning.pytorch.core.module.LightningModule.load_from_checkpoint`` #### on_train_start {#on-train-start} ``lightning.pytorch.core.module.LightningModule.on_train_start`` #### on_train_end {#on-train-end} ``lightning.pytorch.core.module.LightningModule.on_train_end`` #### on_validation_start {#on-validation-start} ``lightning.pytorch.core.module.LightningModule.on_validation_start`` #### on_validation_end {#on-validation-end} ``lightning.pytorch.core.module.LightningModule.on_validation_end`` #### on_test_batch_start {#on-test-batch-start} ``lightning.pytorch.core.module.LightningModule.on_test_batch_start`` #### on_test_batch_end {#on-test-batch-end} ``lightning.pytorch.core.module.LightningModule.on_test_batch_end`` #### on_test_epoch_start {#on-test-epoch-start} ``lightning.pytorch.core.module.LightningModule.on_test_epoch_start`` #### on_test_epoch_end {#on-test-epoch-end} ``lightning.pytorch.core.module.LightningModule.on_test_epoch_end`` #### on_test_start {#on-test-start} ``lightning.pytorch.core.module.LightningModule.on_test_start`` #### on_test_end {#on-test-end} ``lightning.pytorch.core.module.LightningModule.on_test_end`` #### on_predict_batch_start {#on-predict-batch-start} ``lightning.pytorch.core.module.LightningModule.on_predict_batch_start`` #### on_predict_batch_end {#on-predict-batch-end} ``lightning.pytorch.core.module.LightningModule.on_predict_batch_end`` #### on_predict_epoch_start {#on-predict-epoch-start} ``lightning.pytorch.core.module.LightningModule.on_predict_epoch_start`` #### on_predict_epoch_end {#on-predict-epoch-end} ``lightning.pytorch.core.module.LightningModule.on_predict_epoch_end`` #### on_predict_start {#on-predict-start} ``lightning.pytorch.core.module.LightningModule.on_predict_start`` #### on_predict_end {#on-predict-end} ``lightning.pytorch.core.module.LightningModule.on_predict_end`` #### on_train_batch_start {#on-train-batch-start} ``lightning.pytorch.core.module.LightningModule.on_train_batch_start`` #### on_train_batch_end {#on-train-batch-end} ``lightning.pytorch.core.module.LightningModule.on_train_batch_end`` #### on_train_epoch_start {#on-train-epoch-start} ``lightning.pytorch.core.module.LightningModule.on_train_epoch_start`` #### on_train_epoch_end {#on-train-epoch-end} ``lightning.pytorch.core.module.LightningModule.on_train_epoch_end`` #### on_validation_batch_start {#on-validation-batch-start} ``lightning.pytorch.core.module.LightningModule.on_validation_batch_start`` #### on_validation_batch_end {#on-validation-batch-end} ``lightning.pytorch.core.module.LightningModule.on_validation_batch_end`` #### on_validation_epoch_start {#on-validation-epoch-start} ``lightning.pytorch.core.module.LightningModule.on_validation_epoch_start`` #### on_validation_epoch_end {#on-validation-epoch-end} ``lightning.pytorch.core.module.LightningModule.on_validation_epoch_end`` #### configure_model {#configure-model} ``lightning.pytorch.core.module.LightningModule.configure_model`` #### on_validation_model_eval {#on-validation-model-eval} ``lightning.pytorch.core.module.LightningModule.on_validation_model_eval`` #### on_validation_model_train {#on-validation-model-train} ``lightning.pytorch.core.module.LightningModule.on_validation_model_train`` #### on_test_model_eval {#on-test-model-eval} ``lightning.pytorch.core.module.LightningModule.on_test_model_eval`` #### on_test_model_train {#on-test-model-train} ``lightning.pytorch.core.module.LightningModule.on_test_model_train`` #### on_before_optimizer_step {#on-before-optimizer-step} ``lightning.pytorch.core.module.LightningModule.on_before_optimizer_step`` #### configure_gradient_clipping {#configure-gradient-clipping} ``lightning.pytorch.core.module.LightningModule.configure_gradient_clipping`` #### optimizer_step {#optimizer-step} ``lightning.pytorch.core.module.LightningModule.optimizer_step`` #### optimizer_zero_grad {#optimizer-zero-grad} ``lightning.pytorch.core.module.LightningModule.optimizer_zero_grad`` #### prepare_data {#prepare-data} ``lightning.pytorch.core.module.LightningModule.prepare_data`` #### setup {#setup} ``lightning.pytorch.core.module.LightningModule.setup`` #### teardown {#teardown} ``lightning.pytorch.core.module.LightningModule.teardown`` #### train_dataloader {#train-dataloader} ``lightning.pytorch.core.module.LightningModule.train_dataloader`` #### val_dataloader {#val-dataloader} ``lightning.pytorch.core.module.LightningModule.val_dataloader`` #### test_dataloader {#test-dataloader} ``lightning.pytorch.core.module.LightningModule.test_dataloader`` #### predict_dataloader {#predict-dataloader} ``lightning.pytorch.core.module.LightningModule.predict_dataloader`` #### transfer_batch_to_device {#transfer-batch-to-device} ``lightning.pytorch.core.module.LightningModule.transfer_batch_to_device`` #### on_before_batch_transfer {#on-before-batch-transfer} ``lightning.pytorch.core.module.LightningModule.on_before_batch_transfer`` #### on_after_batch_transfer {#on-after-batch-transfer} ``lightning.pytorch.core.module.LightningModule.on_after_batch_transfer`` --- **Navigation** - Parent: [Core API](/docs/pytorch/2.6.6/core-api.md)