Track PyTorch Lightning With Fabric And Wandb
⚡ Track PyTorch Lightning with Fabric and Wandb
At Weights & Biases, we love anything that makes training deep learning models easier. That's why we worked with the folks at PyTorch Lightning to integrate our experiment tracking tool directly into the Fabric library of PyTorch Lightning
PyTorch Lightning is a lightweight wrapper for organizing your PyTorch code and easily adding advanced features such as distributed training and 16-bit precision. It retains all the flexibility of PyTorch, in case you need it, but adds some useful abstractions and builds in some best practices.
Pytorch Fabric allows you to scale PyTorch models on distributed machines while maintaining full control of your training loop.
What this notebook covers:
- How to get basic metric logging with the
WandbLogger - How to log media with W&B
The interactive dashboard in W&B will look like this:
## 💡 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.fabric.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. |
Log custom hyperparameters and configurations
Save Data to Weights and Biases Artifacts
This allows us to audit and create direct data lineages to our experiments
Configure our Model and Training
Load our model, datasources, and loggers into PyTorch Fabric
Run training and log test predictions
For every epoch, run a training step and a test step. For each n test batches, we log the batch of test images caption by the prediction and label, and we create a wandb.Table() in which to store test predictions using our custom callback
No additional dependencies outside the Torch modeling you're used to!
Finish our experiment!