Update Weights & Biases Integration + Add Prediction Visualization for Object Detection

[soumik12345]

Changes made in this PR

Log inference results to Weights & Biases

Fixes #953

Optionally log inference results for object detection in Weights & Biases dashboard using bounding-box image overlays.

This PR adds an optional parameter visualize_on_wandb to ImagesDetectionPrediction.show() which when set to True would log the predictions using the bounding-box image overlays in the Weights & Biases dashboard of the user.

Here's a sample code:

from super_gradients.training import models
from super_gradients.common.object_names import Models

import wandb

# Initialize Weights & Biases run
wandb.init(project="yolo-nas", entity="geekyrakshit")

net = models.get(Models.YOLO_NAS_S, pretrained_weights="coco")
prediction = net.predict([
    "https://cdn.autonomous.ai/static/upload/images/common/upload/20210919/5-Desk-Setup-with-a-Laptop--Monitor-Ideas-for-2021_10f89201411b.jpg",
    "https://p.imgci.com/db/PICTURES/CMS/352000/352019.jpg",
    "https://p.imgci.com/db/PICTURES/CMS/352200/352298.jpg",
    "https://www.encirclephotos.com/wp-content/uploads/Laos-Ban-Pak-Ou-Man-Riding-Asian-Elephant-1440x961.jpg"
])
prediction.show(visualize_on_wandb=True)

This would log the images to be visualized in a media panel in a Weights & Biases dashboard.

👇 Here's how the results would look

Here's a video of how to interact with the bounding-box overlay interface on Weights & Biases 👇

Screen.Recording.2023-05-09.at.6.24.17.PM.mov

Update WandBSGLogger

Fixes #980
Fixes #1008

• WandBSGLogger will now call wandb.init() only if it has not been called earlier.
• Save model checkpoints versioned with Weights & Biases checkpoints.
• Add an additional parameter sync_tensorboard to WandBSGLogger that automatically syncs TensorBoard with the wandb run.

👇 Here's a sample training code that works with the changes made in this PR

👇 Here's a sample Weights & Biases Run demonstrating the features

👇 Tensorboard instance synced with Weights & Biases Run

👇 Lineage of Model Checkpoints

[dagshub]

Join the discussion on DagsHub!

[BenSpex]

+1 would love to see WandB integration

[soumik12345]

Hi @BloodAxe
Can you please take a look at this PR?

[ofrimasad]

Thank you for the code contribution.
The wandb boxes feature looks awesome.

I am not sure I understand the need for Artifact. How is that different from uploading a file to the experiment using wandb.save? Will it be linked to the experiment?

please find my comments regarding the code placement, naming, documentation and some questions inline.

[soumik12345]

I am not sure I understand the need for Artifact. How is that different from uploading a file to the experiment using wandb.save? Will it be linked to the experiment?

Hi @ofrimasad
W&B Artifacts can be thought of as a versioned directory that is either an input of a run or an output of a run. Artifacts can be used to not only store and version your datasets and model checkpoints but logging your models as artifacts enables you to use W&B Model Registry.

Artifacts and runs form a directed graph because a given W&B run can use another run’s output artifact as input. You do not need to define pipelines ahead of time. Weights and Biases will create the DAG for you when you use and log artifacts.

(Sample YOLO-NAS Checkpoint Artifact | Sample YOLO-NAS Checkpoint Lineage)

For more information about exploring an artifacts graph, you can check Explore and traverse artifact graphs.

[ofrimasad]

I see that you still kept the option to init wandb from outside the logger.
In that case please add a warning that will be presented if the run was initialized from outside the logger.
"The wandb run was initialized from outside of the WandbSGLogger. This means that SG cannot control the run ID to which this session will be logged. The current run id is: {run.id}"

Thanks

[ofrimasad]

please use the common logging behavior of SG

[soumik12345]

Hi @ofrimasad is there any more changes we need to make as part of this PR?

[ofrimasad]

LGTM

[ofrimasad]

Thank you for the code contribution @soumik12345. :)
I apologize for the delayed review.

[soumik12345]

Thanks a lot for your reviews @ofrimasad ❤️