Error in Fusing modules in 2-input CNN

quantization
1

Hello,

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?

  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: (<class 'torch.nn.modules.conv.Conv2d'>, <class 'torch.nn.modules.conv.Conv2d'>, <class 'torch.nn.modules.activation.ReLU'>)

The architecture is as below:

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)
  )
)
2

I think fusing conv, conv, relu isn’t supported yet. The only layers you can currently fuse are:

    conv, bn
    conv, bn, relu
    conv, relu
    linear, relu
    bn, relu

From: Fuse Modules Recipe — PyTorch Tutorials 1.10.0+cu102 documentation

Maybe try fusing just the last conv and relu?

3

Thanks Karthik,

One more question,

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?

For example,

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)

Any idea?

4

The fusing is based on the layer names. So in your case that would be encoder_1', ‘encoder_2’ etc. However, I don’t think PyTorch supports fusing those layers. And I’m not sure you can even fuse the conv2d and bn layers since they are actually part of PartialConvLayer.

So if you really wanted to fuse then you may have to write that code yourself.

5

Hello!

Think I figured it out. Looks like I did correctly fused some operations?

    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)

This results in,

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)
  )
)
1 Like
6

Hi, yes looks like it did fuse at least some of the layers. Hope that works for you!