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.
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.
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.