Number of filters ConvOffset2D layer has some compatibility issues

[saradhimpardha]

#construct deformable cnn model for mnist dataset
def deform_cnn(input_tensor = None, input_shape = None, classes = 10, trainable = True):
if K.image_data_format() == 'channels_first':
input_shape = (1, img_rows, img_cols)
bn_axis = 1
else:
input_shape = (img_rows, img_cols, 1)
bn_axis = 3

if input_tensor is None:
    inputs = Input(shape = input_shape)
else:
    if not K.is_keras_tensor(input_tensor):
        inputs = Input(tensor=input_tensor, shape=input_shape)
    else:
        inputs = input_tensor

#Conv_1 layer
x = Conv2D(32, (3, 3), padding = 'same', name = 'conv_1', trainable = trainable)(inputs)
x = BatchNormalization(axis = bn_axis)(x)
x = Activation('relu')(x)

#Conv_2 layer
x_offset = ConvOffset2D(32, name = 'conv_2_offset')(x)
x = Conv2D(64, (3, 3), strides = (2, 2), padding = 'same', name = 'conv_2', trainable = trainable)(x_offset)
x = BatchNormalization(axis = bn_axis)(x)
x = Activation('relu')(x)

#Conv_3 layer
x_offset = ConvOffset2D(64, name = 'conv_3_offset')(x)
x = Conv2D(128, (3, 3), strides = (2, 2), padding = 'same', name = 'conv_3', trainable = trainable)(x_offset)
x = BatchNormalization(axis = bn_axis)(x)
x = Activation('relu')(x)

#Conv_4 layer
x_offset = ConvOffset2D(128, name = 'conv_4_offset')(x)
x = Conv2D(128, (3, 3), padding = 'same', name = 'conv_4', trainable = trainable)(x_offset)
x = BatchNormalization(axis = bn_axis)(x)
x = Activation('relu')(x)

#Pooling layer
x = GlobalAveragePooling2D()(x)

#fc layer
outputs = Dense(classes, activation = 'softmax', name = 'fc', trainable = trainable)(x)

return inputs, outputs

#determine the input for tensorflow: channels_last
input_tensor = Input(shape = (img_rows, img_cols, 1))
inputs, outputs = deform_cnn(input_tensor = input_tensor, trainable = True)

#train model
model = Model(inputs = inputs, outputs = outputs)
model.summary()

When I run this, I am getting the following ValueError:

ValueError Traceback (most recent call last)
in <cell line: 3>()
1 #determine the input for tensorflow: channels_last
2 input_tensor = Input(shape = (img_rows, img_cols, 1))
----> 3 inputs, outputs = deform_cnn(input_tensor = input_tensor, trainable = True)
4
5

2 frames
/content/drive/MyDrive/ColabNotebooks/xingshulicc-deformable-cnn/layers.py in tf__call(self, x)
14 retval_ = ag__.UndefinedReturnValue()
15 x_shape = ag__.converted_call(ag__.ld(x).get_shape, (), None, fscope)
---> 16 offsets = ag__.converted_call(ag__.converted_call(ag__.ld(super), (ag__.ld(ConvOffset2D), ag__.ld(self)), None, fscope).call, (ag__.ld(x),), None, fscope)
17 offsets = ag__.converted_call(ag__.ld(self).to_bc_h_w_2, (ag_.ld(offsets), ag__.ld(x_shape)), None, fscope)
18 x = ag__.converted_call(ag__.ld(self).to_bc_h_w, (ag_.ld(x), ag__.ld(x_shape)), None, fscope)

ValueError: Exception encountered when calling layer "conv_2_offset" (type ConvOffset2D).

in user code:

File "/content/drive/MyDrive/ColabNotebooks/xingshulicc-deformable-cnn/layers.py", line 41, in call  *
    offsets = super(ConvOffset2D, self).call(x)
File "/usr/local/lib/python3.10/dist-packages/keras/layers/convolutional/base_conv.py", line 318, in call  **
    outputs.set_shape(out_shape)

ValueError: Dimension 3 in both shapes must be equal, but are 64 and 32. Shapes are [?,28,28,64] and [?,28,28,32].

Call arguments received by layer "conv_2_offset" (type ConvOffset2D):
• x=tf.Tensor(shape=(None, 28, 28, 32), dtype=float32)

[saradhimpardha]

As much as I know about deformable cnns, the number of channels in the deformable offset layer should be twice the number of channels in the input feature map, but this code is not accepting it that way. When the 'filters' is reduced to half i.e., from 32 to 16, likewise in every offset layer, then the model was built successfully.