Improved W&B integration

[AyushExel]

New features introduced:

• Number of steps logged is now equal to the number of epochs. See example run page
• Log datasets as interactive dsviz tables. View, query, filter, and analyze dataset directly in your browser. Logged dataset will result in a new config file which can be used to train models directly from dsviz datasets. Example dsviz table
• Set --save_period to log model artifacts after every n epochs. Won't log artifacts if this is not set.
• Set --bbox_interval to set bounding-box image logging interval. If not set explicitly, set opt.bbox_interval = opt.epochs // 10
• Resume training runs directly from models saved as artifacts by passing model name as --resume_from_artifact. Combined with dataset logged artifact, this will allow users to train large models on spot instances on cloud or colab without loosing their progress.
• Log training progress as dsviz table to see the training progress. This feature won't use any local storage and will be enabled only if the dataset being used is a logged artifact. Example training progress table

Try this colab walkthrough for all the new features

@glenn-jocher This PR introduces some refactoring of code. To enable all the above features, I had to move a few things around:

• checkdataset() in train.py is called after initializing wandb_logger. This is done because if the dataset is an artifact link, it needs to be downloaded first( which is done by wandb_logger) before being checked
• Optimizer from the final model are stripped after calling wandb.finish() because it interfered with model logging.
• Instead of passing opt.data in test.py, I'm passing the data_dict because in cases where dataset is an artifact link, it will throw error. data_dict contains the location of the locally downloaded artifact. The only reason to use opt.data was to initialize is_coco boolean which I've done in train.py instead. So, the previous functionally will remain unchanged. If data passed is an str, the code will load the yaml file as before, otherwise treat it as data_dict
  Test the colab and tinker with the code and let me know if you find any bugs or possible improvements.

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Integrated Wandb logging enhancements for YOLOv5 🚀

📊 Key Changes

• Improved Wandb logging with added support for dataset and model artifact logging.
• Simplified and refactored Wandb logging setup in train.py.
• Updated how Wandb logger is initialized and called in test methods.
• Wandb logger now checks and uploads datasets as artifacts upon training initialization.
• Fixed Wandb logging during test time with accurate epoch handling and bounding box debug images.
• Removed certain Wandb login related code and CLI options to streamline the interface.
• The artifact_alias option was added to allow specification of dataset version.
• Added automatic resuming from Wandb artifacts, making restarts more graceful.

🎯 Purpose & Impact

• Enhanced Reproducibility: By logging datasets and models as artifacts, it can be easier to track and reproduce experiments.
• Simplified Configuration: The changes lead to a clearer setup process with less manual configuration.
• Improved Usability: Automatic resumption of training and logging simplifies restarting interrupted sessions for users.
• Streamlined Usage: Removing rarely used CLI options can make the interface more user-friendly.
• Visibility and Debugging: Enhanced logging, particularly for bounding box images, will help users better understand model performance.

[glenn-jocher]

@AyushExel looks good! Looking forward to chatting about it tomorrow. /rebase

[AyushExel]

@glenn-jocher Sounds good. I have rebased this branch to track all the latest changes.

[AyushExel]

@glenn-jocher let me when you'll review this and based on that I'll rebase this branch to master again. I think rebasing it when we're ready to merge will be most efficient as manually solving merge conflicts multiple times doesn't make sense :)

[glenn-jocher]

/rebase

[glenn-jocher]

TODO/Look into list:

• instead of passing --data coco128_wandb.yaml, search for coco128.yaml --> Path('coco128_wandb.yaml').isfile() and use instead
• add --upload-dataset argument to train.py to remove need to call to log_datasets.py
• remove --log-imgs and assume hard coded 16.
• suggested resume command python train.py --resume wandb/project/run or --resume wandb/project/run --epochs 350, with pseudocode:

if type(resume) is str and resume.startswith('wandb/'):
   _, project, run = resume.split('/')
   run_id = find_runid_from_run(run, assert_unique_run=True)

[glenn-jocher]

@AyushExel was thinking for the wandb shortcut we should adopt the GCP naming conventions used by Google's gsutil package. gsutil lets you transfer data between local and Google Cloud storage buckets, they use gs://bucket/... to identify a cloud bucket located at https://storage.googleapis.com/bucket/...

In our case we might use python train.py --resume wandb://project/run. This has the added benefit of completely avoiding confusion with the local wandb/ directory.

[AyushExel]

@glenn-jocher I'll address all these changes and get back to you

[AyushExel]

/rebase

[AyushExel]

@glenn-jocher is '/rebase' command not available for all users?

[AyushExel]

@glenn-jocher. I've addressed all the changes we discussed

1. get rid of multiple wandb imports
2. Save the final stripped epoch always ( as --log_artifacts used to do)
3. set save_json to true
4. Don't reset wandb_logger in the test.py
5. Remove dsviz logging from test.py

[AyushExel]

@glenn-jocher I have made all the changes that we discussed about in the last meeting. Please have a look and run the tests. Let's finish this off tomorrow

[AyushExel]

@glenn-jocher Merged master into the branch and all CI tests are passing. All of the tests in colab are also passing. Let's get this merged :)

[glenn-jocher]

@AyushExel PR is merged! Thank you for your contributions.