I could’nt find a a-z example of how to integrate Sweep with a yolo classifier so I had to stitch a few pieces of code to get something going so I may have missed something.
This is my current train.py code:
import argparse
from ultralytics import YOLO
# Parse command line arguments
parser = argparse.ArgumentParser()
parser.add_argument('--dropout', type=float, default=0.2, help='dropout rate')
parser.add_argument('--lr0', type=float, default=0.001, help='learning rate')
parser.add_argument('--scale', type=float, default=0.50, help='scaling the image')
parser.add_argument('--erasing', type=float, default=0.50, help='Erasing part of the image')
parser.add_argument('--epochs', type=int, default=40, help='Epochs count')
parser.add_argument('--batch', type=int, default=16, help='Batch size')
args = parser.parse_args()
data_path = '/home/ubuntu/data-ff/'
model = YOLO('yolov8m-cls') # build a new model from YAML
results = model.train(data=data_path, epochs=args.epochs, imgsz=640, batch=args.batch, device=0, dropout=args.dropout, lr0=args.lr0, hsv_h=0.0, hsv_s=0.0, hsv_v=0.0, degrees=0.0, translate=0.0, scale=args.scale, erasing=args.erasing)
and my sweep.yaml:
program: train.py
method: bayes
metric:
goal: maximize
name: mAP50
parameters:
lr0:
distribution: log_uniform_values
min: 1e-5
max: 1e-2
dropout:
values: [0.0, 0.1, 0.2, 0.3]
scale:
values: [0.1, 0.2, 0.3, 0.4, 0.5]
erasing:
values: [0.1, 0.2, 0.3, 0.4, 0.5]
batch:
values: [8, 16, 32]
If I call this with the wand agent command, so I have to also add the add_wandb_callback call in my train.py script or everything is handled by wandb sweep? How does WandB knows where to get the mAP50 to optimize for? I think I am missing something!
Thanks!
