Since I have converted my model from onnx to PyTorch using the onnx2torch library, I do not have access to model definition and cannot hence modify it. In tf/keras, we are able to build a model by combining the input and output of functional layers, hence in my case, we can find out the name of any intermediate layer and build a model backbone and a new head with that strategy. How can I do the same in PyTorch?
I have tried using model.layer = CustomHead() as @ptrblck has suggested, however, that will only replace one module (layer) with CustomHead. I want to break the entire trained model and attach a new head to it. For reference, following is an image of onnx graph depicting the where I wish to break the model and attach a custom head:
I wish to break the model at Clip (which is ReLU) and replace all the succeeding modules with my CustomHead().