[BLOCKED] Replace IntermediateLayerGetter with feature_extraction.create_feature_extractor()

🚀 The feature

A new feature extraction utility based on FX was introduced by #4302. It allows us to extract internal outputs of a module in a simpler way and it should be a drop-in replacement of the old IntermediateLayerGetter class.

Unfortunately the create_feature_extractor() util is not currently used anywhere on vision. We should investigate using it to replace some of the old approaches. Here is a list of candidates:

• Replace IntermediateLayerGetter calls in Segmentation models [1, 2]
• Replace IntermediateLayerGetter calls in Detection models [3]
• Rewrite model surgery from the SSDFeatureExtractorVGG [4] and SSDLiteFeatureExtractorMobileNet modules [5].

Extra care needs to be given to ensure that the old weights of pre-trained models can still load and the validation statistics remain the same.

Motivation, pitch

Using the latest FX-based utility will help us future proof our code-base.

cc @datumbox

[datumbox]

Backbones that have been passed through FX won’t have class information in the submodules (see https://github.com/pytorch/pytorch/issues/66335). As a result we might no longer be able to do isinstance() to check for specific blocks. This idiom is often used in detection and segmentation for the weight initialization. So prior replacing the feature extractors we must check that none of the models are affected.

@vfdev-5 I know that you previously replaced segmentation but back then this limitation was not known and it’s not clear if we are affected. I was hoping to take more time prior the next release to check carefully that we didn’t mess up anything before rolling out the change to more places. But if you feel you have the capacity to thoroughly check that none of the models for segmentation or detection are affected, feel free to do it now.

[vfdev-5]

@kingjuno you can draft a PR for 2nd task (“Replace IntermediateLayerGetter calls in Detection models [3]”) and I can review and help with meanwhile.