I wrote a minimal example of a training step for yolo11n (or yolo8n) model. It does not propose any targets (no binding boxes, no classes, no labels). Hence, I cannot do backpropagation. I’ve tried decreasing model.conf, but it had no effect. Same thing happens on my custom dataset.
What should I tweak to get some predicted targets, so that the backprop (hence, training) can actually start?
from os import environ
import torch
from ultralytics import YOLO
from torch.optim import Adam
def dummy_train_step():
images = torch.rand(2, 3, 640, 640).cuda()
gt_boxes = [torch.tensor([[0, 0.5, 0.5, 0.2, 0.2]]).cuda(), torch.tensor([[0, 0.7, 0.7, 0.3, 0.3]]).cuda()]
model = YOLO("yolov8n.pt")
model.args['data'] = None
model.conf = 0.0000001 # or 0.1, or 0.00000000000001
optimizer = Adam(model.parameters(), lr=0.001)
for step in range(100):
print(step)
predictions = model(images)
print(f"Predictions: {predictions}")
pred_boxes = predictions[0].boxes.xywh
pred_scores = predictions[0].boxes.conf
pred_labels = predictions[0].boxes.cls
preds_flat = torch.cat([pred_boxes, pred_scores.unsqueeze(-1), pred_labels.unsqueeze(-1)], dim=-1)
if preds_flat.size(0) == 0:
print("No detections in the current batch.")
continue
targets_flat = torch.cat(gt_boxes, dim=0)
print('Shape of targets_flat: ', targets_flat.shape)
loss = model.loss(preds_flat, targets_flat)
print('Loss: ', loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
if __name__ == "__main__":
environ['OMP_NUM_THREADS'] = '1'
dummy_train_step()
orig_shape: (640, 640)
path: ‘image1.jpg’
probs: None
save_dir: ‘runs\detect\predict’
speed: {‘preprocess’: 0.0, ‘inference’: 5.001425743103027, ‘postprocess’: 1.2627840042114258}]
No detections in the current batch.