Skip to content

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

@chandan-wiai

Description

@chandan-wiai

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?

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions