Creating deformed convolution using attention mask in Keras

Question

I wanted to create deformable convolution network in Keras and compare its performance with standard convolution in Keras.

I tried on MNIST fashion data set.

Code for Standard convolution in its simplest form works well giving 85% accuracy in one epoch.

img_rows, img_cols = 28, 28
(x_train, y_train), (x_test, y_test) = fashion_mnist.load_data()

x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)/255.
x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)/255.

y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)

rows = 28
cols = 28
channels = 1
filters1 = 32
output_class = 10

x = Input(shape= (rows,cols,channels))
first_cnn = Conv2D(filters1, kernel_size=(3, 3), strides=(1, 1),activation='relu')
feature_map =  first_cnn(x)
flat_unit = Flatten()
flat_layer = flat_unit(feature_map)
output = Dense(output_class, activation='softmax')(flat_layer)
modelcnn = Model(input = [x], output = [output])


modelcnn.compile(loss='categorical_crossentropy',optimizer=RMSprop(), \
                      metrics=['accuracy'])
modelcnn.fit(x_train , y_train , batch_size =128,epochs = 1)

However, when I create my custom convolution layer for dynamic shaped filters based on attention mask, I get following error apparently in the final 'Dense' layer

InvalidArgumentError: 2 root error(s) found. (0) Invalid argument: Matrix size-incompatible: In[0]: [128,18432], In[1]: [784,10] [[{{node dense_2/MatMul}}]] [[Mean_1/_129]] (1) Invalid argument: Matrix size-incompatible: In[0]: [128,18432], In[1]: [784,10] [[{{node dense_2/MatMul}}]] 0 successful operations. 0 derived errors ignored.

My code is here:

img_rows, img_cols = 28, 28
(x_train, y_train), (x_test, y_test) = fashion_mnist.load_data()

x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)/255.
x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)/255.

y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)


class My_Convolution_Layer(Layer):

    def __init__(self, **kwargs):
        # number of filters is 32 meaning 32 patterns in the mask like convolution
        # the relative postions is 25
        self.output_dim = 32
        self.input_dim = 25
        super(My_Convolution_Layer, self).__init__(**kwargs)

    def build(self, input_shape):
        # Create a trainable weight variable for this layer. which are the same except that
        # they will fire selectively 
        # they could also be dynamically generated 
        self.kernel = self.add_weight(name='kernel', 
                                      shape=(self.input_dim, self.output_dim),
                                      initializer='uniform',
                                      trainable=True)
        super(My_Convolution_Layer, self).build(input_shape)# Be sure to call this at the end

    def call(self, x):
        image, mask = x
        #feature_map
        featuremap= [] # python datastrcture
        # the shape contaisn batch size so be careful
        for rows in range(0,mask.shape[1]):
            for cols in range(0,mask.shape[2]):

                tmp_slice_of_image =image[:,rows:rows+5,cols:cols+5,:]
                flatten_tmp_slice_of_image =  Reshape((25,))(tmp_slice_of_image)
                featuremap.append( flatten_tmp_slice_of_image)

        featuremap_t = Reshape((24,24,25))(K.stack(featuremap, axis =1))
        hardmard_product = \
                multiply([featuremap_t, mask] )
        convolution = K.dot(hardmard_product , self.kernel)
        convolution_relu = keras.activations.relu( convolution )
        return convolution_relu

rows = 28
cols = 28
channels = 1
output_class = 10

x = Input(shape= (rows,cols,channels))
attention_mask_cnn = Conv2D(25, kernel_size=(5, 5), strides=(1, 1),activation='tanh')
attention_mask = attention_mask_cnn(x)
my_convolution_layer = My_Convolution_Layer()
feature_map = my_convolution_layer([x, attention_mask] )
flat_unit = Flatten()
flat_layer = flat_unit(feature_map)
output = Dense(output_class,activation='softmax')(flat_layer)
model_dcnn = Model(input = [x], output = [output])

model_dcnn.compile(loss='categorical_crossentropy',optimizer=RMSprop(), \
                      metrics=['accuracy'])
model_dcnn.fit(x_train , y_train ,batch_size =128, epochs =1)