Train YOLOv5 from scratch (not pre-trained)

Hello,
I want to train the YOLOv5 model (YOLOv5 | PyTorch) from scratch (not using the pretrained weights) on my own dataset and classes for a task of Face Mask Detection.

I have seen that in order to train I should load:
model = torch.hub.load(‘ultralytics/yolov5’, ‘yolov5s’, autoshape=False, pretrained=False) # load scratch

However, How do I actually train it? Can I use it as one layer in my model?

Thank you,
Almog

I don’t believe you could use the model as one “layer” in another model, as I would assume that the YOLOv5 model expects image tensor inputs and returns the predictions (not any intermediate activations), so I’m unsure how it should be used as a layer in a larger model.

This train.py script from the repository seems to show the training of the model.

Thank you for you reply @ptrblck !
I have seen theirs but it requires more modifications to my code… Can I use it as my model but train it with my own train.py code?

I have also tried to use both their model and their loss function existed in train.py (class ComputeLoss in loss.py) but I haven’t found documentation on what the input ‘targets’ to this function should be.

What I mean to do is:

for epoch in tqdm(range(train_params.num_epochs)):
    model.train()
    for i, (images, bboxs, labels) in enumerate(train_loader):
        preds = YOLO_model(images)
        loss = compute_loss(preds, targets)
        loss.backward()
  • What should I give in ‘targets’?
  • And how do I separate ‘preds’ to labels and bboxes?

Thank you!

Yes, you should be able to adapt your custom training script using the reference.
I’m not familiar with the repository so don’t know which shapes and objects are expected.
I would probably start by checking the “train custom data” tutorial and in case that doesn’t help you could try to execute their script and check the expected shapes etc.

1 Like