🔥 Fine-tuning YOLO-NAS using Super-Gradients and Weights & Biases 🐝
This notebook demonstrates how to fine-tune YOLO-NAS on a custom dataset using the Super-Graidents library and performing experiment-tracking, logging and versioning model checkpoints, and viusalizing your detection datasets and prediction results during training.
Installing Dependencies
We install Super-Graidents and Weights & Biases.
Experiment Setup
First, we initialize a W&B run using wandb.init.
Next, we will initialize our trainer which will be in charge of the whole workflow, including the likes of training, evaluation, saving checkpoints, visualization of results, etc.
Setting up the Dataset
Next, we fetch a subset of the BDD100K dataset hosted on Weights & Biases as a dataset artifact. Hosting the dataset as an artifact not only enables us to maintain different versions of our datasets, but also enables us to keep track of the runs that produced it or the runs that are using it via the lineage panel.
Next, we set up the dataset parameters, and create the dataloaders for training, validation and testing using the coco_detection_yolo_format_train and the coco_detection_yolo_format_val functions from Super-Graidents, that would automatically create the dataset loading, pre-processing and augmentation pipelines.
We can visualize our datasets using the plot_detection_dataset_on_wandb function on our Weights & Biases dashboard. The datasets are logged into a Weights & Biases Table, in which the images can be visualized overlayed with an interactive overlays for computer vision tasks.
Here's how the datasets look on Weights & Biases 👇
Performing Transfer Learning
For performing transfer learning, we first define the YOLO_NAS_S model as the pre-trained backbone from Super-Graidents. You can check the Super-Gradients Model Zoo for all the available pre-trained models for object detection.
Experiment Tracking with Weights & Biases
Now, we define the training parameters for the experiment. In order to perform experiment tracking with Weights & Biases, in the training parameters, we set the value of sg_logger to wandb_sg_logger. You can also set the value of additional parameters for the wandb_sg_logger by defining them under sg_logger_params like
train_params = {
...
"sg_logger": "wandb_sg_logger",
"sg_logger_params": {
"save_checkpoints_remote": True,
"save_tensorboard_remote": True,
"save_logs_remote": True,
"save_checkpoint_as_artifact": True,
},
}
Here's a complete list of params for the wandb_sg_logger:
| Parameter | Description | Default Value |
|---|---|---|
| experiment_name | Name used for logging and loading purposes | |
| storage_location | If set to 's3' (i.e. s3://my-bucket) saves the Checkpoints in AWS S3 otherwise saves the Checkpoints Locally | |
| resumed | If true, then old tensorboard files will NOT be deleted when tb_files_user_prompt=True | |
| training_params | training_params for the experiment | |
| checkpoints_dir_path | Local root directory path where all experiment logging directories will reside. | |
| tb_files_user_prompt | Asks user for Tensorboard deletion prompt. | |
| launch_tensorboard | Whether to launch a TensorBoard process. | False |
| tensorboard_port | Specific port number for the tensorboard to use when launched (when set to None, some free port number will be used) | False |
| save_checkpoints_remote | Saves checkpoints in s3. | True |
| save_tensorboard_remote | Saves tensorboard in s3. | True |
| save_logs_remote | Saves log files in s3. | True |
| save_code | Save current code to wandb | False |
| save_checkpoint_as_artifact | Save model checkpoint using Weights & Biases Artifact. Note that setting this option to True would save model checkpoints every epoch as a versioned artifact, which will result in use of increased storage usage on Weights & Biases. | False |
The Weights & Biases integration for Super Gradients also come with the callback WandBDetectionValidationPredictionLoggerCallback for logging object detection predictions to Weights & Biases during training. In order to log object detection predictions to Weights & Biases, we can include this callback in the training parameters:
train_params = {
...
"sg_logger": "wandb_sg_logger",
"sg_logger_params": {
"save_checkpoints_remote": True,
"save_tensorboard_remote": True,
"save_logs_remote": True,
"save_checkpoint_as_artifact": True,
},
"phase_callbacks": [
WandBDetectionValidationPredictionLoggerCallback(class_names=labels),
]
}
Now, we perform the training.
Here's how the validation predictions look at the end of the training 👇
