create_feature_extractor does not track 'if' in forward

🐛 Describe the bug

It looks like create_feature_extractor does not track if statement in forward, for example with the efficientnet architecture:

import torchvision

efficientnetb0 = torchvision.models.efficientnet_b0()
return_nodes = {f'features.{k}': str(v) for v, k in enumerate([3, 4, 6, 7])}

backbone = torchvision.models.feature_extraction.create_feature_extractor(efficientnetb0, return_nodes=return_nodes)

# print first Sequential module for efficientnetb0
print(efficientnetb0.features[1])
>>> Sequential(
  (0): MBConv(
    (block): Sequential(
      (0): ConvNormActivation(
        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)
        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): SiLU(inplace=True)
      )
      (1): SqueezeExcitation(
        (avgpool): AdaptiveAvgPool2d(output_size=1)
        (fc1): Conv2d(32, 8, kernel_size=(1, 1), stride=(1, 1))
        (fc2): Conv2d(8, 32, kernel_size=(1, 1), stride=(1, 1))
        (activation): SiLU(inplace=True)
        (scale_activation): Sigmoid()
      )
      (2): ConvNormActivation(
        (0): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (stochastic_depth): StochasticDepth(p=0.0, mode=row)
  )
)

# print same module for backbone
>>> (1): Module(
    (0): Module(
      (block): Module(
        (0): Module(
          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)
          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): SiLU(inplace=True)
        )
        (1): Module(
          (avgpool): AdaptiveAvgPool2d(output_size=1)
          (fc1): Conv2d(32, 8, kernel_size=(1, 1), stride=(1, 1))
          (activation): SiLU(inplace=True)
          (fc2): Conv2d(8, 32, kernel_size=(1, 1), stride=(1, 1))
          (scale_activation): Sigmoid()
        )
        (2): Module(
          (0): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
    )
  )

Is the StochasticDepth layer implemented but hidden in the backbone or is it missing (due to the if) ?

Versions

Collecting environment information… PyTorch version: 1.10.0 Is debug build: False CUDA used to build PyTorch: 10.2 ROCM used to build PyTorch: N/A

OS: CentOS Linux release 8.2.2004 (Core) (x86_64) GCC version: (GCC) 8.3.1 20191121 (Red Hat 8.3.1-5) Clang version: Could not collect CMake version: Could not collect Libc version: glibc-2.9

Python version: 3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 19:07:31) [GCC 7.3.0] (64-bit runtime) Python platform: Linux-4.18.0-193.6.3.el8_2.x86_64-x86_64-with-centos-8.2.2004-Core Is CUDA available: False CUDA runtime version: Could not collect GPU models and configuration: No devices found. Nvidia driver version: Could not collect cuDNN version: Could not collect HIP runtime version: N/A MIOpen runtime version: N/A

Versions of relevant libraries: [pip3] numpy==1.19.1 [pip3] torch==1.10.0 [pip3] torchvision==0.11.1 [conda] blas 1.0 mkl
[conda] cudatoolkit 10.2.89 hfd86e86_1
[conda] ffmpeg 4.3 hf484d3e_0 pytorch [conda] mkl 2020.2 256
[conda] mkl-service 2.3.0 py36he8ac12f_0
[conda] mkl_fft 1.3.0 py36h54f3939_0
[conda] mkl_random 1.1.1 py36h0573a6f_0
[conda] numpy 1.19.1 py36hbc911f0_0
[conda] numpy-base 1.19.1 py36hfa32c7d_0
[conda] pytorch 1.10.0 py3.6_cuda10.2_cudnn7.6.5_0 pytorch [conda] pytorch-mutex 1.0 cuda pytorch [conda] torchvision 0.11.1 py36_cu102 pytorch

[datumbox]

@rvandeghen After symbolic tracing, the submodules of the model no longer contain the same information as before. In your example all individual blocks such as MBConv, ConvNormActivation, SqueezeExcitation etc are converted to nn.Module. There is a discussion around this at https://github.com/pytorch/pytorch/issues/66335 that you might find interesting.

@jamesr66a Could you please share your thoughts on whether the above is expected?

[datumbox]

I’ve created an issue to add the layers on the tracer: #4777

Perhaps this will help to reduce some of the confusion. The FX is a relatively new functionality and there is a bit of learning curve at the beginning. Hopefully we will iron things out and make things smoother for end users.