Optimize Pytorch Lightning Models With Weights & Biases
⚡ Pytorch Lightning models with Weights & Biases
Pytorch Lightning is a lightweight wrapper for organizing your PyTorch code and easily adding advanced features such as distributed training, 16-bit precision or gradient accumulation.
Coupled with the Weights & Biases integration, you can quickly train and monitor models for full traceability and reproducibility with only 2 extra lines of code:
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch import Trainer
wandb_logger = WandbLogger()
trainer = Trainer(logger=wandb_logger)
W&B integration with Pytorch-Lightning can automatically:
- log your configuration parameters
- log your losses and metrics
- log your model
- keep track of your code
- log your system metrics (GPU, CPU, memory, temperature, etc)
📚 Docs
You can find the PyTorch Lightning WandbLogger docs here and the Weights & Biases docs here
🛠️ Installation and set-up
We make sure we're logged into W&B so that our experiments can be associated with our account.
📊 Setting up the dataloader
For the context of this tutorial we use vanilla pytorch dataloaders on the MNIST dataset
🤓 Defining the Model
Tips:
- Call
self.save_hyperparameters()in__init__to automatically log your hyperparameters to W&B - Call self.log in
training_stepandvalidation_stepto log the metrics
The model is now ready!
💾 Save Model Checkpoints
The ModelCheckpoint callback is required along with the WandbLogger argument to log model checkpoints to W&B.
## 💡 Tracking Experiments with WandbLogger
PyTorch Lightning has a WandbLogger to easily log your experiments with Wights & Biases. Just pass it to your Trainer to log to W&B. See the WandbLogger docs for all parameters. Note, to log the metrics to a specific W&B Team, pass your Team name to the entity argument in WandbLogger
lightning.pytorch.loggers.WandbLogger()
| Functionality | Argument/Function | PS |
|---|---|---|
| Logging models | WandbLogger(... ,log_model='all') or WandbLogger(... ,log_model=True) | Log all models if log_model="all" and at end of training if log_model=True |
| Set custom run names | WandbLogger(... ,name='my_run_name') | |
| Organize runs by project | WandbLogger(... ,project='my_project') | |
| Log histograms of gradients and parameters | WandbLogger.watch(model) | WandbLogger.watch(model, log='all') to log parameter histograms |
| Log hyperparameters | Call self.save_hyperparameters() within LightningModule.__init__() | |
| Log custom objects (images, audio, video, molecules…) | Use WandbLogger.log_text, WandbLogger.log_image and WandbLogger.log_table, etc. |
See the WandbLogger docs here for all parameters.
⚙️ Using WandbLogger to log Images, Text and More
Pytorch Lightning is extensible through its callback system. We can create a custom callback to automatically log sample predictions during validation. WandbLogger provides convenient media logging functions:
WandbLogger.log_textfor text dataWandbLogger.log_imagefor imagesWandbLogger.log_tablefor W&B Tables.
An alternate to self.log in the Model class is directly using wandb.log({dict}) or trainer.logger.experiment.log({dict})
In this case we log the first 20 images in the first batch of the validation dataset along with the predicted and ground truth labels.
🏋️ Train Your Model
When we want to close our W&B run, we call wandb.finish() (mainly useful in notebooks, called automatically in scripts).
We can monitor losses, metrics, gradients, parameters and sample predictions as the model trains.
📚 Resources
- Pytorch Lightning and W&B integration documentation contains a few tips for taking most advantage of W&B
- Pytorch Lightning documentation is extremely thorough and full of examples
❓ Questions about W&B
If you have any questions about using W&B to track your model performance and predictions, please contact support@wandb.com