Skip to main content
App version: 3.4.14

PyTorch Lightning

PyTorch Lightning is a lightweight PyTorch wrapper for high-performance AI research. Use the Neptune integration to:

  • Monitor the training live.
  • Log hyperparameters, model checkpoint paths, metrics, and more.
  • Visualize the logged metadata in the Neptune app.

Note that this integration isn't yet included in the official PyTorch Lightning release and it's built on top of pytorch-lightning==2.5.1.

Before your start​

  • Configure your Neptune API token and project. For details, see Get started.

  • Install Neptune:

    pip install -U neptune-scale

  • Install the custom version of PyTorch Lightning from the source:

    pip install git+https://github.com/neptune-ai/pytorch-lightning.git

Use NeptuneScaleLogger​

PyTorch Lightning has a unified way of logging metadata by using loggers. To log metadata to Neptune, add NeptuneScaleLogger to your PyTorch Lightning script:

  1. Create a NeptuneScaleLogger instance:

    from lightning.pytorch.loggers.neptune import NeptuneScaleLogger


    neptune_logger = NeptuneScaleLogger(
        project="workspace-name/project-name",  # Your Neptune project path
        experiment_name="lightning-experiment",
        log_model_checkpoints=True,  # Logs model checkpoint paths to Neptune
        prefix="training",  # Defines the namespace for logging metadata
    )

    See also the Run init parameters that you can pass to NeptuneScaleLogger and customize its behavior.

    For help on finding your project path, see Get started.

  2. Set up your LightningModule and use the regular logging methods such as self.log():

    from lightning import LightningModule


    class LitModel(LightningModule):
        def training_step(self, batch):
            loss = ...
            self.log("train/batch/loss", loss)

            acc = ...
            self.log("train/batch/acc", acc)

    Metrics are logged as nested dictionary-like structures. For details, see Namespaces and attributes.

  3. Pass NeptuneScaleLogger to the trainer:

    trainer = Trainer(
        logger=neptune_logger,
        ...
    )

  4. Train the model and log metadata to Neptune:

    model = LitModel()
    datamodule = LightningDataModule()

    trainer.fit(model, datamodule)

    Neptune automatically stops tracking once the script has completed execution.

Log metadata to Neptune​

With the Neptune integration, you can combine the default PyTorch Lightning logging methods with the Neptune methods and capabilities.

Access Neptune run directly​

In addition to the default logging methods, you can use the Neptune Run methods to log more metadata types. You can also use the Neptune methods outside of the training loop.

To access the Neptune run directly, use the neptune_logger.run reference:

neptune_logger = NeptuneScaleLogger(...)

class LitModel(LightningModule):
    def any_lightning_module_function_or_hook(self):
        # Log a file to Neptune
        self.neptune_logger.run.assign_files({"dataset/data_sample": "sample_data.csv"})

You can also use the neptune_logger.run reference to log outside the prefix passed to NeptuneScaleLogger.

Add tags​

Use tags to identify and organize your runs. Add tags with the neptune_logger.run reference:

neptune_logger.run.add_tags(["notebook", "lightning"])

Log hyperparameters​

To log hyperparameters, use the standard log_hyperparams() method from the PyTorch Lightning logger:

PARAMS = {
    "batch_size": 64,
    "lr": 0.07,
    ...
}

neptune_logger.log_hyperparams(PARAMS)

Hyperparameters are logged under the <prefix>/hyperparams namespace. To keep the hierarchical structure of the parameters, use the Neptune run directly:

neptune_logger.run.log_configs(
    data={
        "data/batch_size": 64,
        "model/optimizer/name": "adam",
        "model/optimizer/lr": 0.07,
        "model/optimizer/decay_factor": 0.97,
        "model/tokenizer/name": "bert-base-uncased",
    },
)

Log model checkpoint paths​

If you have the ModelCheckpoint callback configured, you can log the paths to model checkpoints. Note that:

  • All checkpoint paths are logged. The save_last and save_top_k parameters aren't supported.
  • Neptune logger doesn't upload the checkpoints to Neptune.

To disable logging of checkpoint paths, set the log_model_checkpoints init parameter to False:

neptune_logger = NeptuneScaleLogger(log_model_checkpoints=False)

Log model summary​

To log a text file with the model summary, use:

model = LitModel()

neptune_logger.log_model_summary(model=model, max_depth=-1)

Analyze the logged metadata​

To analyze the logged metadata in the app, follow the link in the console output. For details, see Explore run metadata in the Neptune app.