Skip to content

PyTorch Lightning ​

Log Lightning metrics to a Rundial run.

Start Rundial and create an API key with the Quickstart. Keep the stack running and set RUNDIAL_API_KEY.

Install the packages:

bash
pip install rundial lightning torch

Add RundialLogger to your training script. This example trains a small model:

python
import torch
from torch.utils.data import DataLoader, TensorDataset
from lightning.pytorch import LightningModule, Trainer
from rundial.integrations import RundialLogger

class Model(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(1, 1)

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = torch.nn.functional.mse_loss(self.layer(x), y)
        self.log("loss", loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.SGD(self.parameters(), lr=0.01)

x = torch.arange(32, dtype=torch.float32).reshape(-1, 1) / 32
loader = DataLoader(TensorDataset(x, 2 * x), batch_size=8)
logger = RundialLogger(workspace="default-workspace", project="default-project", name="lightning")
Trainer(max_epochs=3, logger=logger, log_every_n_steps=1).fit(Model(), loader)

Open the run's Metrics tab to view loss. Checkpoints that Lightning saves are uploaded to the run's Artifacts.

Next steps ​