[BLOCKED] Replace IntermediateLayerGetter with feature_extraction.create_feature_extractor()

[datumbox]

🚀 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

[vfdev-5]

@datumbox looks like it is one-line change. I can do that for segm models. As FX is beta etc, do you think we should keep an option to use IntermediateLayerGetter instead of create_feature_extractor ?

[datumbox]

Thanks @vfdev-5, it might require a bit more work for the 3rd bullet point but overall it should be straightforward. It's not a high priority task and potentially a good intro task for someone who onboards on the code-base. :)

Concerning the FX status, this is a valid concern. I'm in discussions with the FX team to assess better. My current understanding is that the API is not expected to change dramatically and because the use of the create_feature_extractor method will be within the public model constructors, it can be seen as an implementation detail not visible to the users of the methods. So even if it changes, we can easily fix without BC headaches.

[vfdev-5]

OK, I'll send only a PR for segmentation models and leave the rest for others :)

My concern on BC was from pytorch side and torch fx functionality (being still in beta), I was imagining if something breaks in their API etc, torchvision models can be directly impacted without a way to go back and use IntermediateLayerGetter. Maybe, we can add use_fe to _segm and _segm_lraspp_mobilenetv3 with default as True (for example) and thus it still gives a way to get previous behaviour... If you think we can go without this flag, I'm OK as well.

[datumbox]

I see. Given the create_feature_extractor is public, I think we are already exposed to the same risk. If you prefer you can delay picking this up until we have a definitive answer.

[kingjuno]

@datumbox can I work on this issue?

[datumbox]

@kingjuno There are a few limitations of using the FX-based approach (losing class info) which means replacing the SSD ones might not work. I'm currently OOO, but I can look into if this still makes sense once I'm back.

[kingjuno]

ok

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

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

[kingjuno]

Backbones that have been passed through FX won't have class information in the submodules (see pytorch/pytorch#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 Should I proceed with this issue?