search

zhiqwang / yolort

yolort is a runtime stack for yolov5 on specialized accelerators such as tensorrt, libtorch, onnxruntime, tvm and ncnn.
https://zhiqwang.com/yolort
GNU General Public License v3.0
708 stars 153 forks source link

Loading pre-trained model is not supported for num_classes != 80 #475

Open chandan-wiai opened 1 year ago

chandan-wiai commented 1 year ago

https://github.com/zhiqwang/yolov5-rt-stack/blob/b7cb695beacec273ea97cc0e3732797580ef37b5/yolort/models/yolo.py#L263

Starting from a pre-trained model on a custom dataset would help in faster convergence and better model performance. But currently when we try to use a pretrained model with num_classes other than 80, it fails and we have train the model from scratch instead. One possible solution of this could be keeping strict=False while loading state dictionary in line 263. model.load_state_dict(state_dict, strict=False)

Can this be implemented?

zhiqwang commented 1 year ago

Hi @chandan-wiai ,

Starting from a pre-trained model on a custom dataset would help in faster convergence and better model performance.

Yep, This would be a very useful feature, but unfortunately the training mechanism is not fully developed yet. See #59 and #60 for more details. Maybe we need to do more preparation, the timing is not very clear to me at the moment, my focus is now moving to some other projects.