Tracking your YOLOX Runs with Weights & Biases

A quick walkthrough on how to use the Weights & Biases integration in YOLOX to track your model training and visualise the results!
Manan Goel
Created on April 1|Last edited on April 6
Comment
﻿
IntroductionIn this report, we will run through how to use the Weights & Biases integration in YOLOX to train your own models with just 1 command line argument. 
python tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o [--cache] --logger wandb
This integration will automatically
Log all the metrics to your W&B dashboard
Save the models at every evaluation step
Tags the model with the best average precision
Show you the predicted bounding boxes along with the confidence score!
We'll be training the model on a subset of the COCO dataset and if you want follow along in a colab notebook with executable code, follow the link below.
Link to Colab
﻿
Setup 🖥We being by downloading the YOLOX Github repository and installing the package along with the wandb client. Run the following commands on the terminal.
git clone https://github.com/manangoel99/YOLOX
cd YOLOX
git checkout -b WandbTables; git pull origin WandbTables; pip install -e .
pip install wandb -qqq
wandb login
Downloading the Dataset 💽With YOLOX and wandb all set up, the next step is to download the dataset. For this report, we will work with a toy subset of the huge COCO dataset. The data has been uploaded as a wandb artifact for ease of use and can be downloaded using the following python script.
api = wandb.Api()
artifact = api.artifact("manan-goel/YOLOX-coco/coco128:latest", type="dataset")
artifact_dir = artifact.download(root="/content")
We will now unzip the dataset and place it in the datasets subdirectory of YOLOX.
unzip coco128.zip
mv coco128 YOLOX/datasets/COCO
cd YOLOX
Training 🏋️Using wandb just requires configuring the command line argument --logger wandb. This automatically turns on the wandb logger for your experiment and further arguments can be added
wandb-project: To specify the project in which experiment is being run.
wandb-run: The name of the wandb run
wandb-entity: Entity which is starting the run
wandb-log_checkpoints: True/False to log model checkpoints to the wandb dashboard
wandb-num_eval_images: Number of images from the validation set to be logged to wandb. Predictions corresponding to these can be visualized on the dashboard. No images are logged if the value is 0 and all are logged if the value is -1.
and more!
Train a YOLOX model from scratchThe different YOLOX models can be trained from scratch with the entire process being logged to W&B.
python -m yolox.tools.train -n yolox-nano -d 1 -b 64 --fp16 \
    --logger wandb \
    wandb-project yolox \
    wandb-log_checkpoints True \
    wandb-num_eval_images 3 \
    eval_interval 1 \
    print_interval 1 \
    max_epoch 10
This command trains a yolox-nano model from scratch on the COCO128 subset. The training metrics are logged to the yolox project with model checkpoints being logged at every evaluation step. At every evaluation step, the dashboard also shows 3 images from the validation set along with the predicted bounding box.
﻿
Run: legendary-valley-691
﻿
The table above shows the confidence score for each class and the predicted bounding box on the image with the corresponding label.
P.S. The visualisations here are from the training of yolox-nano on the entire COCO-2014 dataset for 300 epochs
Train on your own datasetYou can also finetune a pre-trained model on a custom dataset. In this case we continue working on the subset of COCO.
Download the dataset (Optional)
Download the pre-trained modelWe have done a lot of work so far to train a YOLOX model from scratch. Now let's use it! The integration automatically logs the checkpoints and labels them so we can just use that from our run mentioned above.
import wandb
api = wandb.Api()
artifact = api.artifact("manan-goel/yolox-nano/run_3ntph3ki_model:best") # Fetches the model tagged as the best during fine-tuning
artifact.download() # Downloads the model
Now that we have the best model downloaded, we just use that for our finetuning needs
python tools/train.py -f exps/example/custom/nano.py -d 1 -b 64 --fp16 -o \
    -c./artifacts/run_3ntph3ki_model:v40/model_ckpt.pth \
    --logger wandb \
    wandb-project yolox-nano-finetune \
    max_epoch 5 \
    print_interval 1 \
    eval_interval 1
ConclusionYOLOX is an amazing tool for object detection and with the Weights & Biases integration now available, the combination makes it a lot easier to train and debug your models for all your object detection tasks. More documentation is available here.
﻿
Add a comment