PyTorch Lightning

PyTorch Lightning provides lifecycle hooks that make metric logging seamless. Use the on_train_epoch_end hook to automatically log metrics after each training epoch without cluttering your training loop.

---

Quick Example

import pytorch_lightning as pl
import json

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        # Your model setup...
    
    def training_step(self, batch, batch_idx):
        # Your training logic...
        loss = self.compute_loss(batch)
        self.log('train_loss', loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        # Your validation logic...
        loss = self.compute_loss(batch)
        acc = self.compute_accuracy(batch)
        self.log('val_loss', loss)
        self.log('val_acc', acc)
    
    def on_train_epoch_end(self):
        # Log to Valohai after each epoch
        metrics = self.trainer.callback_metrics
        
        print(json.dumps({
            "epoch": self.current_epoch,
            "train_loss": float(metrics.get('train_loss', 0)),
            "val_loss": float(metrics.get('val_loss', 0)),
            "val_acc": float(metrics.get('val_acc', 0))
        }))

---

Why Use Hooks?

PyTorch Lightning already tracks metrics internally. The hook pattern lets you:

• Access all logged metrics in one place

• Log to Valohai without modifying training logic

• Keep your code clean and framework-idiomatic

---

Complete Working Example

Here's a full training script with Valohai integration:

import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import valohai

class LitMNIST(pl.LightningModule):
    def __init__(self, learning_rate=0.001):
        super().__init__()
        self.save_hyperparameters()
        
        # Simple CNN
        self.conv1 = torch.nn.Conv2d(1, 32, 3, 1)
        self.conv2 = torch.nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = torch.nn.Dropout(0.25)
        self.dropout2 = torch.nn.Dropout(0.5)
        self.fc1 = torch.nn.Linear(9216, 128)
        self.fc2 = torch.nn.Linear(128, 10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        
        # PyTorch Lightning automatically tracks this
        self.log('train_loss', loss, prog_bar=True)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = (preds == y).float().mean()
        
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', acc, prog_bar=True)
        return loss
    
    def on_train_epoch_end(self):
        # Access metrics collected during the epoch
        metrics = self.trainer.callback_metrics
        
        # Log to Valohai
        metadata = {
            "epoch": self.current_epoch
        }
        
        # Log all collected metrics
        if 'train_loss' in metrics:
            metadata["train_loss"] = float(metrics['train_loss'])
        if 'val_loss' in metrics:
            metadata["val_loss"] = float(metrics['val_loss'])
        if 'val_acc' in metrics:
            metadata["val_acc"] = float(metrics['val_acc'])
        
        # Log current learning rate
        current_lr = self.trainer.optimizers[0].param_groups[0]['lr']
        metadata["learning_rate"] = current_lr
        
        print(json.dumps(metadata))
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
        return [optimizer], [scheduler]

# Training script
def main():
    # Prepare data
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_dataset = datasets.MNIST(
        '/valohai/inputs/dataset',
        train=True,
        download=True,
        transform=transform
    )
    val_dataset = datasets.MNIST(
        '/valohai/inputs/dataset',
        train=False,
        transform=transform
    )
    
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=32)
    
    # Initialize model
    model = LitMNIST(learning_rate=0.001)
    
    # Log hyperparameters at the start
    import json
    print(json.dumps({
        "model": "LitMNIST",
        "optimizer": "adam",
        "learning_rate": 0.001,
        "batch_size": 32,
        "scheduler": "StepLR"
    }))
    
    # Train
    trainer = pl.Trainer(
        max_epochs=10,
        accelerator='auto',
        devices=1
    )
    
    trainer.fit(model, train_loader, val_loader)
    
    # Log final test results
    test_loader = DataLoader(val_dataset, batch_size=32)
    test_results = trainer.test(model, test_loader)
    
    with valohai.metadata.logger() as logger:
        logger.log("final_test_loss", test_results[0]['test_loss'])
        logger.log("final_test_acc", test_results[0]['test_acc'])
    
    # Save model
    output_path = valohai.outputs().path('model.ckpt')
    trainer.save_checkpoint(output_path)

if __name__ == '__main__':
    main()

---

valohai.yaml Configuration

Make sure to change the input data and environment to match your own values.

- step:
    name: train-lightning
    image: pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime
    command:
      - pip install pytorch-lightning
      - python train.py
    inputs:
      - name: dataset
        default: dataset://training-data/latest
    parameters:
      - name: learning_rate
        type: float
        default: 0.001
      - name: batch_size
        type: integer
        default: 32
      - name: epochs
        type: integer
        default: 10
    environment: aws-eu-west-1-g4dn-xlarge

---

Alternative: Using Callbacks

If you prefer not to modify your LightningModule, create a custom callback:

import pytorch_lightning as pl
import json
import torch

class ValohaiMetricsCallback(pl.Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        metrics = trainer.callback_metrics
        
        metadata = {
            "epoch": trainer.current_epoch + 1
        }
        
        # Log all metrics tracked by Lightning
        for key, value in metrics.items():
            if isinstance(value, torch.Tensor):
                value = float(value)
            metadata[key] = value
        
        # Log learning rate
        if trainer.optimizers:
            lr = trainer.optimizers[0].param_groups[0]['lr']
            metadata["learning_rate"] = lr
        
        print(json.dumps(metadata))

# Use the callback
trainer = pl.Trainer(
    max_epochs=10,
    callbacks=[ValohaiMetricsCallback()]
)

Benefits of the callback approach:

• No changes to your LightningModule

• Reusable across projects

• Easier to enable/disable Valohai logging

---

Logging Custom Metrics

You can log any metric you compute:

def on_train_epoch_end(self):
    metadata = {
        "epoch": self.current_epoch + 1,
        "train_loss": float(self.trainer.callback_metrics['train_loss'])
    }
    
    # Custom computations
    model_size_mb = sum(p.numel() for p in self.parameters()) * 4 / 1024 / 1024
    metadata["model_size_mb"] = model_size_mb
    
    # Gradient norms
    total_norm = 0
    for p in self.parameters():
        if p.grad is not None:
            total_norm += p.grad.data.norm(2).item() ** 2
    total_norm = total_norm ** 0.5
    metadata["gradient_norm"] = total_norm
    
    print(json.dumps(metadata))

---

Available Hooks

PyTorch Lightning provides many hooks you can use:

Training hooks:

• on_train_epoch_start — Before each epoch

• on_train_epoch_end — After each epoch (most common)

• on_train_batch_end — After each batch (use sparingly)

Validation hooks:

• on_validation_epoch_end — After validation completes

Testing hooks:

• on_test_epoch_end — After testing completes

Example using multiple hooks:

def on_train_epoch_start(self):
    # Log at the start of each epoch
    print(json.dumps({
        "epoch_start": self.current_epoch + 1,
        "learning_rate": self.trainer.optimizers[0].param_groups[0]['lr']
    }))

def on_train_epoch_end(self):
    # Log at the end of each epoch
    print(json.dumps({
        "epoch": self.current_epoch + 1,
        "train_loss": float(self.trainer.callback_metrics['train_loss'])
    }))

def on_validation_epoch_end(self):
    # Log validation results
    print(json.dumps({
        "val_loss": float(self.trainer.callback_metrics['val_loss']),
        "val_acc": float(self.trainer.callback_metrics['val_acc'])
    }))

---

Best Practices

Convert Tensors to Python Types

PyTorch Lightning metrics are often tensors. Convert them before logging:

# ✅ Good: Convert to float
metadata["loss"] = float(metrics['train_loss'])

# ❌ Avoid: Logging tensors directly
metadata["loss"] = metrics['train_loss']  # May cause JSON serialization issues

---

Use Consistent Metric Names

Keep metric names consistent with PyTorch Lightning's conventions:

# ✅ Good: Clear naming
self.log('train_loss', loss)  # In training_step
logger.log("train_loss", float(metrics['train_loss']))  # In hook

# ❌ Avoid: Renaming metrics
self.log('loss', loss)  # In training_step
logger.log("training_loss", float(metrics['loss']))  # In hook

---

Don't Log Every Batch

Logging after every batch creates too much data:

# ✅ Good: Log per epoch
def on_train_epoch_end(self):
    with valohai.metadata.logger() as logger:
        logger.log("epoch", self.current_epoch + 1)
        logger.log("train_loss", float(self.trainer.callback_metrics['train_loss']))

# ❌ Avoid: Log every batch (unless debugging)
def on_train_batch_end(self, outputs, batch, batch_idx):
    with valohai.metadata.logger() as logger:
        logger.log("batch", batch_idx)
        logger.log("loss", float(outputs['loss']))

---

Log Hyperparameters at Start

Log hyperparameters once before training begins:

def main():
    model = LitModel(learning_rate=0.001)
    
    # Log hyperparameters
    import json
    print(json.dumps({
        "model": model.__class__.__name__,
        "learning_rate": model.hparams.learning_rate,
        "optimizer": "adam",
        "batch_size": 32
    }))
    
    trainer = pl.Trainer(max_epochs=10)
    trainer.fit(model, train_loader, val_loader)

---

Common Issues

Metrics Not Available in Hook

Symptom: KeyError when accessing metrics

Cause: Metric not logged in training_step or validation_step

Solution: Make sure you call self.log() for metrics you want to access:

def validation_step(self, batch, batch_idx):
    loss = self.compute_loss(batch)
    acc = self.compute_accuracy(batch)
    
    # Log metrics so they're available in hooks
    self.log('val_loss', loss)
    self.log('val_acc', acc)

def on_train_epoch_end(self):
    metrics = self.trainer.callback_metrics
    
    # Now these are available
    if 'val_loss' in metrics:
        with valohai.metadata.logger() as logger:
            logger.log("val_loss", float(metrics['val_loss']))

---

Tensor Serialization Errors

Symptom: TypeError: Object of type Tensor is not JSON serializable

Solution: Convert tensors to Python types:

# ✅ Good
logger.log("loss", float(metrics['train_loss']))

# ❌ Wrong
logger.log("loss", metrics['train_loss'])  # Still a tensor

---

Example Project

Check out our complete working example on GitHub:

valohai/pytorch-lightning-example

The repository includes:

• Complete training script with Valohai integration

• valohai.yaml configuration

• Example notebooks

• Step-by-step setup instructions

---

Next Steps

• Visualize your metrics in Valohai

• Compare experiments to find the best hyperparameters

• Learn more about PyTorch Lightning callbacks

• Back to Collect Metrics overview