error loading pre-trained weights for fine-tuning YOLONAS

[albertofernandezvillan]

Instantiating a YOLONAS model for fine-tuning:

model = models.get('yolo_nas_l', num_classes=len(dataset_params['classes']), pretrained_weights="coco")

Line above works, but following line (just copy/paste the download weights to the current working directory):

model = models.get('yolo_nas_l', num_classes=len(dataset_params['classes']), pretrained_weights="coco", checkpoint_path="./yolo_nas_l_coco.pth")

Throws an error loading the weights of the head (as my dataset has different number of classes).
I want to load the weights from a file on disk for fine-tuning.

How to cope with this error?

[dagshub]

Join the discussion on DagsHub!

[Louis-Dupont]

Hi @albertofernandezvillan

I listed below 2 ways you can load and fine-tune your checkpoint, depending on your case

1. Standard flow

Step 1: Download the model to fine-tune

model = models.get('yolo_nas_m', num_classes=20, pretrained_weights="coco")
Trainer.train(model=model, training_params=...)
...

Step 2. Load fine-tuned model

model = models.get('yolo_nas_m', checkpoint_path="<path-to-fine-tuned-checkpoint-file>")

Note that:

• This standard flow doesn't suffer from the bug you raised
• No need to set pretrained_weights="coco" in the step2., since you already provide a specific checkpoint and therefore you are not using the default weights pretrained on COCO anymore.
• No need to set num_classes=20 in step 2., since your model was already fine-tuned with 20 classes

I believe this should cover your case, but I am aware this is not exactly what you asked for.

2. Fine-tuning from a local checkpoint

If you want to load the model before fine-tuning it (basically step 1) from a local checkpoint, and this local checkpoint has a different number of heads than what you want, this is the way to do:

model = models.get('yolo_nas_m', num_classes=20, checkpoint_num_classes=80, checkpoint_path=".cache/torch/hub/checkpoints/yolo_nas_m_coco.pth")

Step 2. remains the same.

Hoping this will help you!