This is a simple demo for performing semantic segmentation on the Kitti dataset using Pytorch-Lightning and optimizing the neural network by monitoring and comparing runs with Weights & Biases.

Pytorch-Ligthning includes a logger for W&B that can be called simply with:

from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning import Trainer

wandb_logger = WandbLogger()
trainer = Trainer(logger=wandb_logger)

Refer to the documentation for more details. Hyper-parameters can be defined manually and every run is automatically logged onto Weights & Biases for easier analysis/interpretation of results and how to optimize the architecture.

You can also run sweeps to optimize automatically hyper-parameters.

See full code on Github →


  1. Install dependencies through requirements.txt, Pipfile or manually (Pytorch, Pytorch-Lightning & Wandb)

  2. Log in or sign up for an account -> wandb login

  3. Run python

  4. Visualize and compare your runs through generated link. It'll log your model performance, gradients and any other metrics you choose.


Hyperparameter Optimization with Sweeps

With Sweeps, you can automate hyperparameter optimization and explore the space of possible models.

  1. Run wandb sweep sweep.yaml

  2. Run wandb agent <sweep_id> where <sweep_id> is given by previous command

  3. Visualize and compare the sweep runs.

After running the script a few times, you will be able to compare quickly a large combination of hyperparameters. Feel free to modify the script and define your own hyperparameters.

See full code on Github →


Pytorch-Lightning let us use Pytorch based code and easily add extra features such as distributed computing over several GPU's and machines, half-precision training, gradient accumulation.

In this example, we optimize following hyper-parameters:

It's interesting to see the possible combination of parameters to have a good performance (here defined by low val_loss):

Note: it's important to keep in mind that we limited training to 20 epochs. Deeper networks typically need more time to be trained.

Try Lightning yourself →