Error in Fusing modules in 2-input CNN

dalseeroh: <p>Hello,</p>
<p>I recently wrote a UNet-like model which receives 2 inputs(input, mask) of same shape and returns 1 output. I am trying to fuse convolutions/batchnorm/Relu like below but it returns error. Is there any way to fuse the input/mask conv layers at the same time?</p>
<pre><code class="lang-auto">  for m in model.modules():
        if type(m) == PartialConvLayer:
            torch.quantization.fuse_modules(m, ["input_conv", "mask_conv", "activation"], inplace=True)
            print(type(m))


AssertionError: did not find fuser method for: (&lt;class 'torch.nn.modules.conv.Conv2d'&gt;, &lt;class 'torch.nn.modules.conv.Conv2d'&gt;, &lt;class 'torch.nn.modules.activation.ReLU'&gt;) 
</code></pre>
<p>The architecture is as below:</p>
<pre><code class="lang-auto">PartialConvUNet(
  (encoder_1): PartialConvLayer(
    (input_conv): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    (mask_conv): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    (activation): ReLU()
  )
  (encoder_2): PartialConvLayer(
    (input_conv): Conv2d(64, 128, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), bias=False)
    (mask_conv): Conv2d(64, 128, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), bias=False)
    (batch_normalization): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): ReLU()
  )
  (encoder_3): PartialConvLayer(
    (input_conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): ReLU()
  )
  (encoder_4): PartialConvLayer(
    (input_conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): ReLU()
  )
  (encoder_5): PartialConvLayer(
    (input_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): ReLU()
  )
  (encoder_6): PartialConvLayer(
    (input_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): ReLU()
  )
  (encoder_7): PartialConvLayer(
    (input_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): ReLU()
  )
  (decoder_5): PartialConvLayer(
    (input_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_6): PartialConvLayer(
    (input_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_7): PartialConvLayer(
    (input_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_4): PartialConvLayer(
    (input_conv): Conv2d(768, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(768, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_3): PartialConvLayer(
    (input_conv): Conv2d(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_2): PartialConvLayer(
    (input_conv): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (mask_conv): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_1): PartialConvLayer(
    (input_conv): Conv2d(67, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (mask_conv): Conv2d(67, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  )
)
</code></pre>

Karthik_Ganesan: <p>I think fusing <code>conv, conv, relu</code> isn’t supported yet. The only layers you can currently fuse are:</p>
<pre><code class="lang-auto">    conv, bn
    conv, bn, relu
    conv, relu
    linear, relu
    bn, relu
</code></pre>
<p>From: <a href="https://pytorch.org/tutorials/recipes/fuse.html" class="inline-onebox" rel="noopener nofollow ugc">Fuse Modules Recipe — PyTorch Tutorials 1.10.0+cu102 documentation</a></p>
<p>Maybe try fusing just the last conv and relu?</p>

dalseeroh: <p>Thanks Karthik,</p>
<p>One more question,</p>
<p>If I would like to fuse input_conv + batch_norm + activation at encoder layers, and mask_conv + batch_norm at decoder layers, how can I distinguish those layers(encoder_1, encoder_2, …, decoder_1, decoder_2,…) when iterating the loop?</p>
<p>For example,</p>
<pre><code class="lang-auto">for m in model.modules():
        if m.methods_for_layers_name == "encoder_1":
            torch.quantization.fuse_modules(m, ["input_conv", "batch_normalization", "activation"], inplace=True)
        elif m.methods_for_layers_name == "decoder_1":
            torch.quantization.fuse_modules(m, ["mask_conv", "batch_normalization"], inplace=True)
</code></pre>
<p>Any idea?</p>

Karthik_Ganesan: <p>The fusing is based on the layer names. So in your case that would be <code>encoder_1', </code>‘encoder_2’ etc. However, I don’t think PyTorch supports fusing those layers. And I’m not sure you can even fuse the <code>conv2d</code> and <code>bn</code> layers since they are actually part of <code>PartialConvLayer</code>.</p>
<p>So if you really wanted to fuse then you may have to write that code yourself.</p>

dalseeroh: <p>Hello!</p>
<p>Think I figured it out. Looks like I did correctly fused some operations?</p>
<pre><code class="lang-auto">    for name, module in model.named_modules():
        if type(module) == PartialConvLayer:

            # encoder_1 fusion 
            if "encoder_1" in name:
                torch.quantization.fuse_modules(module, [['input_conv', 'activation']], inplace=True)

            # encoder_2 ~ encoder_7 fusion
            elif "enc" in name:
                torch.quantization.fuse_modules(module, [['input_conv', 'batch_normalization', 'activation']], inplace=True)

            # decoder_2 ~ decoder_7 fusion
            elif "decoder_1" not in name:
                torch.quantization.fuse_modules(module, [['input_conv', 'batch_normalization']], inplace=True)

    print("Fusion completed")
    print(model)
</code></pre>
<p>This results in,</p>
<pre><code class="lang-auto">Fusion completed
PartialConvUNet(
  (encoder_1): PartialConvLayer(
    (input_conv): ConvReLU2d(
      (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
      (1): ReLU()
    )
    (mask_conv): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    (activation): Identity()
  )
  (encoder_2): PartialConvLayer(
    (input_conv): ConvBnReLU2d(
      (0): Conv2d(64, 128, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), bias=False)
      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
    )
    (mask_conv): Conv2d(64, 128, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), bias=False)
    (batch_normalization): Identity()
    (activation): Identity()
  )
  (encoder_3): PartialConvLayer(
    (input_conv): ConvBnReLU2d(
      (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
    )
    (mask_conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): Identity()
  )
  (encoder_4): PartialConvLayer(
    (input_conv): ConvBnReLU2d(
      (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
    )
    (mask_conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): Identity()
  )
  (encoder_5): PartialConvLayer(
    (input_conv): ConvBnReLU2d(
      (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
    )
    (mask_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): Identity()
  )
  (encoder_6): PartialConvLayer(
    (input_conv): ConvBnReLU2d(
      (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
    )
    (mask_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): Identity()
  )
  (encoder_7): PartialConvLayer(
    (input_conv): ConvBnReLU2d(
      (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
    )
    (mask_conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): Identity()
  )
  (decoder_7): PartialConvLayer(
    (input_conv): ConvBn2d(
      (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (mask_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_6): PartialConvLayer(
    (input_conv): ConvBn2d(
      (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (mask_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_5): PartialConvLayer(
    (input_conv): ConvBn2d(
      (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (mask_conv): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_4): PartialConvLayer(
    (input_conv): ConvBn2d(
      (0): Conv2d(768, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (mask_conv): Conv2d(768, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_3): PartialConvLayer(
    (input_conv): ConvBn2d(
      (0): Conv2d(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (mask_conv): Conv2d(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_2): PartialConvLayer(
    (input_conv): ConvBn2d(
      (0): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (mask_conv): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (batch_normalization): Identity()
    (activation): LeakyReLU(negative_slope=0.2)
  )
  (decoder_1): PartialConvLayer(
    (input_conv): Conv2d(67, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (mask_conv): Conv2d(67, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  )
)
</code></pre>

Karthik_Ganesan: <p>Hi, yes looks like it did fuse at least some of the layers. Hope that works for you!</p>