I tried the Pytorch integration and it doesn’t work:
from ultralytics import YOLO
import wandb
wandb.init()
# Load a model
model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training)
# Magic
wandb.watch(model, log_freq=100)
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
wandb.log({"loss": loss})
# Use the model
results = model.train(data="coco128.yaml", epochs=3, batch = 2) # train the model
I get the following:
Traceback (most recent call last):
File "train.py", line 10, in <module>
wandb.watch(model, log_freq=100)
File "/home/henry/.local/bin/.virtualenvs/ultralytics/lib/python3.8/site-packages/wandb/sdk/wandb_watch.py", line 71, in watch
raise ValueError(
ValueError: Expected a pytorch model (torch.nn.Module). Received <class 'ultralytics.yolo.engine.model.YOLO'>