How to Integrate YOLOv5 with Weights & Biases for Experiment Tracking

YOLOv5 integrates with Weights & Biases through the WandbLogger class in utils/loggers/wandb/wandb_utils.py, enabling automatic metric logging, model checkpoint artifact storage, and dataset versioning when you add the --log wandb flag to your training command.

The ultralytics/yolov5 repository provides native integration with Weights & Biases (W&B) for comprehensive experiment tracking and model management. By leveraging the built-in WandbLogger class, you can monitor training metrics in real-time, visualize validation predictions, and version your datasets without modifying the core training loop. This guide explains how to activate and configure the integration using the actual source code implementation.

Prerequisites

Before starting, install the W&B Python package and authenticate your account. You only need to log in once per machine.

pip install wandb
wandb login

Alternatively, set the WANDB_API_KEY environment variable to authenticate non-interactively.

Enabling W&B Logging in YOLOv5

When you execute train.py, the script creates a Trainer object that initializes loggers through utils/loggers/__init__.py. At line 124 of this file, the code inspects the --log argument and instantiates a WandbLogger if wandb is present in the list (or if the legacy --log-wandb flag is used). No code changes are required—simply pass the appropriate flag to enable tracking.

Core Logging Capabilities

The WandbLogger class handles five critical lifecycle events during training, all implemented in utils/loggers/wandb/wandb_utils.py.

Run Initialization

Inside WandbLogger.__init__ (lines 64-73), the logger calls wandb.init() with your training configuration (opt), project name, entity, run name, and an optional run ID. The resulting run object is stored as self.wandb_run, making it available for all subsequent logging operations.

Dataset Versioning

If you pass the --upload-dataset flag, the setup_training method uploads your dataset metadata as a W&B artifact. This creates a reproducible link between your training run and the exact dataset version used, including path information and hash values.

Metric Tracking

During each training epoch, the Trainer calls WandbLogger.log() to send scalar metrics (loss, mAP, precision, recall) via self.wandb.log(log_dict). At the end of every epoch, the end_epoch method flushes accumulated metrics to the server with self.wandb.log(self.log_dict), ensuring real-time dashboard updates.

Model Checkpoint Logging

The log_model() method (lines 23-45) creates a wandb.Artifact object that stores your checkpoint files (last.pt and best.pt). These artifacts are tagged with semantic labels including "latest", "epoch X", and "best" when applicable, allowing you to retrieve specific model versions directly from the W&B interface.

Run Finalization

When training completes, the finish_run() method flushes any remaining logs and gracefully closes the connection by calling wandb.run.finish().

Command-Line Usage Examples

Run a complete training job with full W&B tracking, including dataset artifact upload:

python train.py \
  --data coco.yaml \
  --cfg yolov5s.yaml \
  --weights '' \
  --batch-size 16 \
  --epochs 100 \
  --project yolov5-experiments \
  --entity my-wandb-team \
  --name yolov5s_coco_run1 \
  --log wandb \
  --upload-dataset

Resume a previous run using its W&B run ID (useful for recovering from interruptions):

python train.py \
  --data coco.yaml \
  --cfg yolov5s.yaml \
  --weights '' \
  --batch-size 16 \
  --epochs 100 \
  --resume run_1a2b3c4d5e \
  --log wandb

Programmatic Integration (Advanced)

For custom training loops or Jupyter notebooks, instantiate the logger directly and control logging manually:

from utils.loggers.wandb.wandb_utils import WandbLogger
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--project', default='my-yolov5')
parser.add_argument('--entity', default='my-team')
parser.add_argument('--name', default='exp')
parser.add_argument('--upload_dataset', action='store_true')
opt = parser.parse_args()

# Initialize the logger

wandb_logger = WandbLogger(opt)

# Log custom metrics during your loop

wandb_logger.log({'train_loss': 0.123, 'learning_rate': 0.001})
wandb_logger.end_epoch()

# Upload model files

wandb_logger.log_model('path/to/best.pt', 'best')

# Close the run

wandb_logger.finish_run()

Summary

• Install the wandb package and authenticate with wandb login before your first run.
• Enable tracking by adding --log wandb to your train.py command; the logger initializes automatically via utils/loggers/__init__.py.
• Configure project, entity, and run name using the standard --project, --entity, and --name arguments.
• Version your dataset by adding --upload-dataset, which triggers artifact creation in setup_training.
• Resume interrupted runs by passing the W&B run ID to the --resume flag; the logger sets resume="allow" internally.

Frequently Asked Questions

Do I need to modify YOLOv5 source code to use Weights & Biases?

No. The integration is completely automatic once you install the wandb package and add the --log wandb flag. The Trainer class in train.py detects this argument and instantiates the WandbLogger from utils/loggers/wandb/wandb_utils.py without requiring any edits to the repository files.

How do I resume a crashed training run in Weights & Biases?

Pass the original W&B run ID to the --resume argument when restarting training. The WandbLogger.__init__ method detects the ID and sets resume="allow" in the wandb.init() call, restoring the run's context and continuing the logging stream from the previous step.

Can I log custom metrics beyond the default YOLOv5 metrics?

Yes. While the standard integration automatically logs loss and mAP values, you can instantiate WandbLogger directly in a custom script and call wandb_logger.log({'custom_metric': value}) at any point. This is useful for tracking domain-specific metrics or external validation results.

Where are model checkpoints stored in Weights & Biases?

Checkpoints are stored as versioned artifacts via the log_model() method in wandb_utils.py. Each artifact contains the last.pt file (and best.pt when applicable) with tags like "epoch X" and "best", allowing you to download specific model versions directly from the W&B artifacts tab or API using the artifact name and version hash.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}