Choosing the neural network model based on the outcome of a primary NN model

1

Hi,

My model’s (ResNet50) weights are frozen except its final Fully Connected layer. I am thinking of this approach: For e.g.

Model 1: FC will be Linear(in_features=2048, out_features=2, bias=True)
Model 2: FC will be Linear(in_features=2048, out_features=10, bias=True)
Model 3: FC will be Linear(in_features=2048, out_features=5, bias=True)

My Model 1 will do binary classification for any provided image. Then, based on its outcome, it will choose Model 2 or Model 3 and completes the task for the same image. In all these cases, only the FC layer is getting changed and learnt. The rest of weights in the architecture are fixed and does not undergo training.

How can I integrate the above points together in pytorch in such a way that if I pass an image to model 1, it follows the flow? If someone has better approach or suggestions for this task, I would be happy to consider. Please share your thoughts.

Thanks in advance!