I have trained a yolov5 (Model A), and I want to train another model (Model B) using the backbone of “A”, and a specific head.
Example problem.
I have a dataset of traffic scenes that include different types of vehicles. For example, the same image can have a class(Super category) of “Car” and class (category) of “BMW” or “Mercedes”[Model A], another image can have a class (Super category) of “Motorcycle”, and class of “Honda” or “Yamaha”[Model B].
–
Since the input images and therefore most of the features should be the same, I wish to train one model then freeze it’s feature-extractor, and sort of transfer learn the same model for the categories using the pretrained feature extractor.
During inference, I want to load Model A and do its prediction → Load Model B Head → Swap “Model A” head with “Model B” → Do the prediction of secondary model.
All of this process has the objective of saving computational resources, since the weights of the feature extracting backbone is the same, I only want to save the weights of the head B.
The idea sounds valid assuming you could replace the “head” attribute in the original model with your custom modules. Did you try to use this approach and are you seeing any issues?
Thanks for the fast response, Unfortunately I’m just giving the first steps in deep learning and PyTorch, I have read so much but couldn’t reach a conclusion on a starting point. I think a better example of my problem is shown below.
I have a dataset of traffic scenes that include different brands of cars. For example, the same image can have a ‘BMW’ or a ‘Mercedes’ among other brands. I want to create a model, where each head classifies a ‘car model’ for a specific brand, like the image below.
I kind of want to “future-prof” the model, and train only a specific head when a new car model launches. Also, one of the request was using yolov5.
