How to save and load specific layers of yolov5, to aplly in diferent heads during inference?

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.