create_feature_extractor does not track 'if' in forward

[rvandeghen]

🐛 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

[rvandeghen]

Here is the output of I obtain after tracing the backbone module:

from torch.fx import symbolic_trace

symbolic_traced = symbolic_trace(backbone)
print(symbolic_traced.code)

>>>
torch.fx._symbolic_trace.wrap("torchvision_ops_stochastic_depth_stochastic_depth")

def forward(self, x : torch.Tensor):
    features_0_0 = getattr(getattr(self.features, "0"), "0")(x);  x = None
    features_0_1 = getattr(getattr(self.features, "0"), "1")(features_0_0);  features_0_0 = None
    features_0_2 = getattr(getattr(self.features, "0"), "2")(features_0_1);  features_0_1 = None
    features_1_0_block_0_0 = getattr(getattr(getattr(getattr(self.features, "1"), "0").block, "0"), "0")(features_0_2);  features_0_2 = None
    features_1_0_block_0_1 = getattr(getattr(getattr(getattr(self.features, "1"), "0").block, "0"), "1")(features_1_0_block_0_0);  features_1_0_block_0_0 = None
    features_1_0_block_0_2 = getattr(getattr(getattr(getattr(self.features, "1"), "0").block, "0"), "2")(features_1_0_block_0_1);  features_1_0_block_0_1 = None
    features_1_0_block_1_avgpool = getattr(getattr(getattr(self.features, "1"), "0").block, "1").avgpool(features_1_0_block_0_2)
    features_1_0_block_1_fc1 = getattr(getattr(getattr(self.features, "1"), "0").block, "1").fc1(features_1_0_block_1_avgpool);  features_1_0_block_1_avgpool = None
    features_1_0_block_1_activation = getattr(getattr(getattr(self.features, "1"), "0").block, "1").activation(features_1_0_block_1_fc1);  features_1_0_block_1_fc1 = None
    features_1_0_block_1_fc2 = getattr(getattr(getattr(self.features, "1"), "0").block, "1").fc2(features_1_0_block_1_activation);  features_1_0_block_1_activation = None
    features_1_0_block_1_scale_activation = getattr(getattr(getattr(self.features, "1"), "0").block, "1").scale_activation(features_1_0_block_1_fc2);  features_1_0_block_1_fc2 = None
    mul = features_1_0_block_1_scale_activation * features_1_0_block_0_2;  features_1_0_block_1_scale_activation = features_1_0_block_0_2 = None
    features_1_0_block_2_0 = getattr(getattr(getattr(getattr(self.features, "1"), "0").block, "2"), "0")(mul);  mul = None
    features_1_0_block_2_1 = getattr(getattr(getattr(getattr(self.features, "1"), "0").block, "2"), "1")(features_1_0_block_2_0);  features_1_0_block_2_0 = None
    features_2_0_block_0_0 = getattr(getattr(getattr(getattr(self.features, "2"), "0").block, "0"), "0")(features_1_0_block_2_1);  features_1_0_block_2_1 = None
    features_2_0_block_0_1 = getattr(getattr(getattr(getattr(self.features, "2"), "0").block, "0"), "1")(features_2_0_block_0_0);  features_2_0_block_0_0 = None
    features_2_0_block_0_2 = getattr(getattr(getattr(getattr(self.features, "2"), "0").block, "0"), "2")(features_2_0_block_0_1);  features_2_0_block_0_1 = None
    features_2_0_block_1_0 = getattr(getattr(getattr(getattr(self.features, "2"), "0").block, "1"), "0")(features_2_0_block_0_2);  features_2_0_block_0_2 = None
    features_2_0_block_1_1 = getattr(getattr(getattr(getattr(self.features, "2"), "0").block, "1"), "1")(features_2_0_block_1_0);  features_2_0_block_1_0 = None
    features_2_0_block_1_2 = getattr(getattr(getattr(getattr(self.features, "2"), "0").block, "1"), "2")(features_2_0_block_1_1);  features_2_0_block_1_1 = None
    features_2_0_block_2_avgpool = getattr(getattr(getattr(self.features, "2"), "0").block, "2").avgpool(features_2_0_block_1_2)
    features_2_0_block_2_fc1 = getattr(getattr(getattr(self.features, "2"), "0").block, "2").fc1(features_2_0_block_2_avgpool);  features_2_0_block_2_avgpool = None
    features_2_0_block_2_activation = getattr(getattr(getattr(self.features, "2"), "0").block, "2").activation(features_2_0_block_2_fc1);  features_2_0_block_2_fc1 = None
    features_2_0_block_2_fc2 = getattr(getattr(getattr(self.features, "2"), "0").block, "2").fc2(features_2_0_block_2_activation);  features_2_0_block_2_activation = None
    features_2_0_block_2_scale_activation = getattr(getattr(getattr(self.features, "2"), "0").block, "2").scale_activation(features_2_0_block_2_fc2);  features_2_0_block_2_fc2 = None
    mul_1 = features_2_0_block_2_scale_activation * features_2_0_block_1_2;  features_2_0_block_2_scale_activation = features_2_0_block_1_2 = None
    features_2_0_block_3_0 = getattr(getattr(getattr(getattr(self.features, "2"), "0").block, "3"), "0")(mul_1);  mul_1 = None
    features_2_0_block_3_1 = getattr(getattr(getattr(getattr(self.features, "2"), "0").block, "3"), "1")(features_2_0_block_3_0);  features_2_0_block_3_0 = None
    features_2_1_block_0_0 = getattr(getattr(getattr(getattr(self.features, "2"), "1").block, "0"), "0")(features_2_0_block_3_1)
    features_2_1_block_0_1 = getattr(getattr(getattr(getattr(self.features, "2"), "1").block, "0"), "1")(features_2_1_block_0_0);  features_2_1_block_0_0 = None
    features_2_1_block_0_2 = getattr(getattr(getattr(getattr(self.features, "2"), "1").block, "0"), "2")(features_2_1_block_0_1);  features_2_1_block_0_1 = None
    features_2_1_block_1_0 = getattr(getattr(getattr(getattr(self.features, "2"), "1").block, "1"), "0")(features_2_1_block_0_2);  features_2_1_block_0_2 = None
    features_2_1_block_1_1 = getattr(getattr(getattr(getattr(self.features, "2"), "1").block, "1"), "1")(features_2_1_block_1_0);  features_2_1_block_1_0 = None
    features_2_1_block_1_2 = getattr(getattr(getattr(getattr(self.features, "2"), "1").block, "1"), "2")(features_2_1_block_1_1);  features_2_1_block_1_1 = None
    features_2_1_block_2_avgpool = getattr(getattr(getattr(self.features, "2"), "1").block, "2").avgpool(features_2_1_block_1_2)
    features_2_1_block_2_fc1 = getattr(getattr(getattr(self.features, "2"), "1").block, "2").fc1(features_2_1_block_2_avgpool);  features_2_1_block_2_avgpool = None
    features_2_1_block_2_activation = getattr(getattr(getattr(self.features, "2"), "1").block, "2").activation(features_2_1_block_2_fc1);  features_2_1_block_2_fc1 = None
    features_2_1_block_2_fc2 = getattr(getattr(getattr(self.features, "2"), "1").block, "2").fc2(features_2_1_block_2_activation);  features_2_1_block_2_activation = None
    features_2_1_block_2_scale_activation = getattr(getattr(getattr(self.features, "2"), "1").block, "2").scale_activation(features_2_1_block_2_fc2);  features_2_1_block_2_fc2 = None
    mul_2 = features_2_1_block_2_scale_activation * features_2_1_block_1_2;  features_2_1_block_2_scale_activation = features_2_1_block_1_2 = None
    features_2_1_block_3_0 = getattr(getattr(getattr(getattr(self.features, "2"), "1").block, "3"), "0")(mul_2);  mul_2 = None
    features_2_1_block_3_1 = getattr(getattr(getattr(getattr(self.features, "2"), "1").block, "3"), "1")(features_2_1_block_3_0);  features_2_1_block_3_0 = None
    stochastic_depth = torchvision_ops_stochastic_depth_stochastic_depth(features_2_1_block_3_1, 0.025, 'row', True);  features_2_1_block_3_1 = None
    add = stochastic_depth + features_2_0_block_3_1;  stochastic_depth = features_2_0_block_3_1 = None
   ...
    stochastic_depth_8 = torchvision_ops_stochastic_depth_stochastic_depth(features_6_3_block_3_1, 0.17500000000000002, 'row', True);  features_6_3_block_3_1 = None
    add_8 = stochastic_depth_8 + add_7;  stochastic_depth_8 = add_7 = None
    features_7_0_block_0_0 = getattr(getattr(getattr(getattr(self.features, "7"), "0").block, "0"), "0")(add_8)
    features_7_0_block_0_1 = getattr(getattr(getattr(getattr(self.features, "7"), "0").block, "0"), "1")(features_7_0_block_0_0);  features_7_0_block_0_0 = None
    features_7_0_block_0_2 = getattr(getattr(getattr(getattr(self.features, "7"), "0").block, "0"), "2")(features_7_0_block_0_1);  features_7_0_block_0_1 = None
    features_7_0_block_1_0 = getattr(getattr(getattr(getattr(self.features, "7"), "0").block, "1"), "0")(features_7_0_block_0_2);  features_7_0_block_0_2 = None
    features_7_0_block_1_1 = getattr(getattr(getattr(getattr(self.features, "7"), "0").block, "1"), "1")(features_7_0_block_1_0);  features_7_0_block_1_0 = None
    features_7_0_block_1_2 = getattr(getattr(getattr(getattr(self.features, "7"), "0").block, "1"), "2")(features_7_0_block_1_1);  features_7_0_block_1_1 = None
    features_7_0_block_2_avgpool = getattr(getattr(getattr(self.features, "7"), "0").block, "2").avgpool(features_7_0_block_1_2)
    features_7_0_block_2_fc1 = getattr(getattr(getattr(self.features, "7"), "0").block, "2").fc1(features_7_0_block_2_avgpool);  features_7_0_block_2_avgpool = None
    features_7_0_block_2_activation = getattr(getattr(getattr(self.features, "7"), "0").block, "2").activation(features_7_0_block_2_fc1);  features_7_0_block_2_fc1 = None
    features_7_0_block_2_fc2 = getattr(getattr(getattr(self.features, "7"), "0").block, "2").fc2(features_7_0_block_2_activation);  features_7_0_block_2_activation = None
    features_7_0_block_2_scale_activation = getattr(getattr(getattr(self.features, "7"), "0").block, "2").scale_activation(features_7_0_block_2_fc2);  features_7_0_block_2_fc2 = None
    mul_15 = features_7_0_block_2_scale_activation * features_7_0_block_1_2;  features_7_0_block_2_scale_activation = features_7_0_block_1_2 = None
    features_7_0_block_3_0 = getattr(getattr(getattr(getattr(self.features, "7"), "0").block, "3"), "0")(mul_15);  mul_15 = None
    features_7_0_block_3_1 = getattr(getattr(getattr(getattr(self.features, "7"), "0").block, "3"), "1")(features_7_0_block_3_0);  features_7_0_block_3_0 = None
    return {'0': add_1, '1': add_3, '2': add_8, '3': features_7_0_block_3_1}

From the torch.fx documentation, I don't know how I can follow the example of Static Control Flow.

@datumbox @NicolasHug any idea on that case ?

[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 pytorch/pytorch#66335 that you might find interesting.

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

[rvandeghen]

@datumbox What I meant is that in the MBConv block (converted in Module), I do not have the StochasticDepth layer and I do not understand why it is present in symbolic_trace but not in the "model" itself.

[rvandeghen]

@datumbox @jamesr66a I think the problem is that StochasticDepth is a module implemented in torchvision but not in torch given that I face the same problem with the FrozenBatchNorm2d. Here is the snippet code where I simply change the self.stochastic_depth layer:

class MBConv(nn.Module):
    def __init__(self, cnf: MBConvConfig, stochastic_depth_prob: float, norm_layer: Callable[..., nn.Module],
                 se_layer: Callable[..., nn.Module] = SqueezeExcitation) -> None:
        super().__init__()

        if not (1 <= cnf.stride <= 2):
            raise ValueError('illegal stride value')

        self.use_res_connect = cnf.stride == 1 and cnf.input_channels == cnf.out_channels

        layers: List[nn.Module] = []
        activation_layer = nn.SiLU

        # expand
        expanded_channels = cnf.adjust_channels(cnf.input_channels, cnf.expand_ratio)
        if expanded_channels != cnf.input_channels:
            layers.append(ConvNormActivation(cnf.input_channels, expanded_channels, kernel_size=1,
                                             norm_layer=norm_layer, activation_layer=activation_layer))

        # depthwise
        layers.append(ConvNormActivation(expanded_channels, expanded_channels, kernel_size=cnf.kernel,
                                         stride=cnf.stride, groups=expanded_channels,
                                         norm_layer=norm_layer, activation_layer=activation_layer))

        # squeeze and excitation
        squeeze_channels = max(1, cnf.input_channels // 4)
        layers.append(se_layer(expanded_channels, squeeze_channels, activation=partial(nn.SiLU, inplace=True)))

        # project
        layers.append(ConvNormActivation(expanded_channels, cnf.out_channels, kernel_size=1, norm_layer=norm_layer,
                                         activation_layer=None))

        self.block = nn.Sequential(*layers)
        # Change StochasticDepth with another torch vision layer
        self.stochastic_depth = torchvision.ops.misc.FrozenBatchNorm2d(1,1)#StochasticDepth(stochastic_depth_prob, "row")
        self.out_channels = cnf.out_channels

    def forward(self, input: Tensor) -> Tensor:
        result = self.block(input)
        # Remove if statement
        #if self.use_res_connect:
        result = self.stochastic_depth(result)
        result += input
        return result

mbconv = MBConv(inverted_residual_setting[1], 0.2, nn.BatchNorm2d)
new_mbconv = torchvision.models.feature_extraction.create_feature_extractor(mbconv, return_nodes={'add': 0})
print(mbconv)
>>> MBConv(
  (block): Sequential(
    (0): ConvNormActivation(
      (0): Conv2d(16, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): SiLU(inplace=True)
    )
    (1): ConvNormActivation(
      (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=96, bias=False)
      (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): SiLU(inplace=True)
    )
    (2): SqueezeExcitation(
      (avgpool): AdaptiveAvgPool2d(output_size=1)
      (fc1): Conv2d(96, 4, kernel_size=(1, 1), stride=(1, 1))
      (fc2): Conv2d(4, 96, kernel_size=(1, 1), stride=(1, 1))
      (activation): SiLU(inplace=True)
      (scale_activation): Sigmoid()
    )
    (3): ConvNormActivation(
      (0): Conv2d(96, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (stochastic_depth): FrozenBatchNorm2d(1, eps=1)
)

print(new_mbconv)
>>> MBConv(
  (block): Module(
    (0): Module(
      (0): Conv2d(16, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): SiLU(inplace=True)
    )
    (1): Module(
      (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=96, bias=False)
      (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): SiLU(inplace=True)
    )
    (2): Module(
      (avgpool): AdaptiveAvgPool2d(output_size=1)
      (fc1): Conv2d(96, 4, kernel_size=(1, 1), stride=(1, 1))
      (activation): SiLU(inplace=True)
      (fc2): Conv2d(4, 96, kernel_size=(1, 1), stride=(1, 1))
      (scale_activation): Sigmoid()
    )
    (3): Module(
      (0): Conv2d(96, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
)

[datumbox]

@rvandeghen I believe by design only the pytorch.nn layers are maintained. The rest are flattened by the tracer.

@fmassa Do you think we should update the tracer to exclude all torchvision.ops defined layers?

[fmassa]

@rvandeghen the description from @datumbox is accurate, the layers are in there but they have been flattened.

@datumbox I believe we could update the tracer to make all torchvision.ops to be leaf nodes. That's a good idea

[rvandeghen]

@fmassa Thanks for the precision. I guess I could have trust the output of the symbolic_trace but it was not that clear at first sight. I guess doing your recommended changes may give something more intuitive.
If that's it we can close the issue.

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