Depthwise cross-correlation with tensorflow 2

Question

I want to implement the depthwise cross-correlation layer described in SiamRPN++ with tensorflow 2 and keras. It should be a subclass of keras layer to allow a flexible usage. My implementation compiles correctly, but in training tensorflow throws the error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Specified a list with shape [8,24,32] from a tensor with shape [24,8,32]

Here's my code. What am I doing wrong?

class CrossCorr(Layer):
"""
Implements the cross correlation laer of siam_rpn_plus_plus
"""

def __init__(self, **kwargs):
    super().__init__(**kwargs)

def build(self, inputs):
    super(CrossCorr, self).build(inputs)  # Be sure to call this at the end

def call(self, inputs, **kwargs):
    def _corr(search_img, filter):
        x = tf.expand_dims(search_img, 0)
        f = tf.expand_dims(filter, -1)
        # use the feature map as kernel for the depthwise conv2d of tensorflow
        return tf.nn.depthwise_conv2d(input=x, filter=f, strides=[1, 1, 1, 1], padding='SAME')
    
    # Iteration over each batch
    out = tf.map_fn(
        lambda filter_simg: _corr(filter_simg[0], filter_simg[1]),
        elems=inputs,
        dtype=inputs[0].dtype
    )
    return tf.squeeze(out, [1])

def compute_output_shape(self, input_shape):
    return input_shape

To call it I use:

def _conv_block(inputs, filters, kernel, strides, kernel_regularizer=None):
    x = Conv2D(filters, kernel, padding='same', strides=strides, 
    kernel_regularizer=kernel_regularizer)(inputs)
    x = BatchNormalization()(x)
    return Activation(relu)(x)


def cross_correlation_layer(search_img, template_img, n_filters=None):
    n_filters = int(search_img.shape[-1]) if n_filters is None else n_filters

    tmpl = _conv_block(template_img, n_filters, 3, 1, kernel_regularizer=L1L2(1e-5, 1e-4))
    search = _conv_block(search_img, n_filters, 3, 1, kernel_regularizer=L1L2(1e-5, 1e-4))

    # calculate cross correlation by striding the generated "filter" over the image in depthwise manner
    cc = CrossCorr()([search, tmpl])
    # 1D conv to make it a seperable convolution
    cc = Conv2D(filters=n_filters, kernel_size=1, strides=1)(cc)
    # apply one more filter over it
    fusion = _conv_block(cc, n_filters, 3, 1)
    return fusion

Answer 1

After attempting to run your code, I realised that the Layer (which you call via _conv_block) is expecting a set of batched images, the following is modified from the above in your cross_correlation_layer function

def cross_correlation_layer(search_img, template_img, n_filters=None):
    n_filters = int(search_img.shape[-1]) if n_filters is None else n_filters

    template_img = tf.expand_dims(template_img, -1)
    search_img = tf.expand_dims(search_img, -1)

    tmpl = _conv_block(template_img, n_filters, 3, 1, kernel_regularizer=L1L2(1e-5, 1e-4))
    search = _conv_block(search_img, n_filters, 3, 1, kernel_regularizer=L1L2(1e-5, 1e-4))

The rest of your function is left unchanged, hope this helps, (Not related to your question, but the word 'filter' is a reserved keyword in Python, we try to avoid using those as parameter names)