Weights & Biases ¶

LaminDB can be integrated with W&B to track the training process and associate datasets & parameters with models.

# pip install lamindb torchvision lightning wandb
!lamin init --storage ./lamin-mlops
!wandb login

import lightning as pl
import lamindb as ln
from lamindb.integrations import lightning as ll
import wandb

from torch import utils
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from autoencoder import LitAutoEncoder

# define model run parameters, features, and labels so that validation passes later on
MODEL_CONFIG = {"hidden_size": 32, "bottleneck_size": 16, "batch_size": 32}
hyperparameter = ln.Feature(name="Autoencoder hyperparameter", is_type=True).save()
hyperparams = ln.Feature.from_dict(MODEL_CONFIG, type=hyperparameter).save()

metrics_to_annotate = ["train_loss", "val_loss", "current_epoch"]
for metric in metrics_to_annotate:
    dtype = int if metric == "current_epoch" else float
    ln.Feature(name=metric, dtype=dtype).save()

# create all Wandb related features like 'wandb_run_id'
ln.examples.wandb.save_wandb_features()

# create all lightning integration features like 'score'
ll.save_lightning_features()

# track this notebook/script run so that all checkpoint artifacts are associated with the source code
ln.track(params=MODEL_CONFIG, project=ln.Project(name="Wandb tutorial").save())

Define a model¶

We use a basic PyTorch Lightning autoencoder as an example model.

Query & download the MNIST dataset¶

We saved the MNIST dataset in a curation notebook which now shows up in the Artifact registry:

ln.Artifact.filter(kind="dataset").to_dataframe()

Let’s get the dataset:

mnist_af = ln.Artifact.get(key="testdata/mnist")
mnist_af

And download it to a local cache:

path = mnist_af.cache()
path

Create a PyTorch-compatible dataset:

mnist_dataset = MNIST(path.as_posix(), transform=ToTensor())
mnist_dataset

Monitor training with wandb¶

Train our example model and track the training progress with wandb.

from lightning.pytorch.loggers import WandbLogger

# create the data loader
train_dataset = MNIST(root="./data", train=True, download=True, transform=ToTensor())
val_dataset = MNIST(root="./data", train=False, download=True, transform=ToTensor())
train_loader = utils.data.DataLoader(train_dataset, batch_size=32)
val_loader = utils.data.DataLoader(val_dataset, batch_size=32)

# init model
autoencoder = LitAutoEncoder(
    MODEL_CONFIG["hidden_size"], MODEL_CONFIG["bottleneck_size"]
)

# initialize the logger
wandb_logger = WandbLogger(project="lamin")

# add batch size to the wandb config
wandb_logger.experiment.config["batch_size"] = MODEL_CONFIG["batch_size"]

# Create a LaminDB LightningCallback which also (optionally) annotates checkpoints by desired metrics
wandb_logger.experiment.id
lamindb_callback = ll.Checkpoint(
    dirpath=f"testmodels/wandb/{wandb_logger.experiment.id}",
    features={
        "run": {
            "wandb_run_id": wandb_logger.experiment.id,
            "wandb_run_name": wandb_logger.experiment.name,
        },
        "artifact": {
            **{metric: None for metric in metrics_to_annotate}
        },  # auto-populated through callback
    },
)

# train model
trainer = pl.Trainer(
    limit_train_batches=3,
    max_epochs=5,
    logger=wandb_logger,
    callbacks=[lamindb_callback],
)
trainer.fit(
    model=autoencoder, train_dataloaders=train_loader, val_dataloaders=val_loader
)

Show code cell output Hide code cell output

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

┏━━━┳━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name    ┃ Type       ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ encoder │ Sequential │ 25.6 K │ train │     0 │
│ 1 │ decoder │ Sequential │ 26.4 K │ train │     0 │
└───┴─────────┴────────────┴────────┴───────┴───────┘

Trainable params: 52.1 K                                                                                           
Non-trainable params: 0                                                                                            
Total params: 52.1 K                                                                                               
Total estimated model params size (MB): 0                                                                          
Modules in train mode: 8                                                                                           
Modules in eval mode: 0                                                                                            
Total FLOPs: 0

/opt/hostedtoolcache/Python/3.13.11/x64/lib/python3.13/site-packages/lightning/pytorch/utilities/_pytree.py:21: 
`isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` 
instead.

/opt/hostedtoolcache/Python/3.13.11/x64/lib/python3.13/site-packages/lightning/pytorch/trainer/connectors/data_conn
ector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the 
value of the `num_workers` argument` to `num_workers=3` in the `DataLoader` to improve performance.

/opt/hostedtoolcache/Python/3.13.11/x64/lib/python3.13/site-packages/lightning/pytorch/trainer/connectors/data_conn
ector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the 
value of the `num_workers` argument` to `num_workers=3` in the `DataLoader` to improve performance.

/opt/hostedtoolcache/Python/3.13.11/x64/lib/python3.13/site-packages/lightning/pytorch/loops/fit_loop.py:317: The 
number of training batches (3) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower 
value for log_every_n_steps if you want to see logs for the training epoch.

! calling anonymously, will miss private instances

`Trainer.fit` stopped: `max_epochs=5` reached.

wandb_logger.experiment.name

wandb.finish()

W&B and LaminDB user interfaces together¶

W&B and LaminDB runs:

Both W&B and LaminDB capture any runs together with run parameters.

W&B experiment overview	LaminHub run overview

W&B run details and LaminDB artifact details:

W&B and LaminDB complement each other. Whereas W&B is excellent at capturing metrics over time, LaminDB excells at capturing lineage of input & output data and training checkpoints.

W&B run view	LaminHub run view

Both frameworks display output artifacts that were generated during the run. LaminDB further captures input artifacts, their origin and the associated source code.

W&B artifact view	LaminHub artifact view

All checkpoints are automatically annotated by the specified training metrics and W&B run ID & name to keep both frameworks in sync:

last_checkpoint_af = (
    ln.Artifact.filter(is_best_model=True)
    .filter(suffix__endswith="ckpt", is_latest=True)
    .last()
)
last_checkpoint_af.describe()

To reuse the checkpoint later:

last_checkpoint_af.cache()

last_checkpoint_af.view_lineage()

Features associated with a whole training run are annotated on a run level:

ln.context.run.features

ln.finish()