Applying a mask to Conv2D kernel in Keras

user_4257302 0
Viewed 1542

I'm looking to apply a mask to the kernel of a Conv2D layer in Keras. I'm having a bit of difficulty understanding kernel shape.

For kernel_size = 3, and filters = 1, the shape of the kernel is (3, 3, 4, 1) => (kernel_size, kernel_size, ???, filters)

What does the 3rd dimension in the kernel represent?

How can I take an NxN mask and multiply it to each of the kernel filters?

This is the code I have so far. I am not sure if it will work as I expect it to because I don't fully understand the kernel shape.

class MaskedConv2D(tf.keras.layers.Layer):
    def __init__(self, *args, **kwargs):
        super(MaskedConv2D, self).__init__()
        self.conv2d = Conv2D(*args, **kwargs)
        
    def build(self, input_shape):
        self.conv2d.build(input_shape[0])
        self._convolution_op = self.conv2d._convolution_op
        
    def masked_convolution_op(self, filters, kernel, mask):
        m = K.expand_dims(K.expand_dims(mask[0, ...], axis=2), axis=3) # (3, 3) => (3, 3, 1, 1)
        m = K.tile(m, (1, 1, kernel.shape[2], kernel.shape[3])) # (3, 3, 1, 1) => (3, 3, 4, 1)
        return self._convolution_op(filters, tf.math.multiply(kernel, m))        
        
    def call(self, inputs):
        x, mask = inputs
        self.conv2d._convolution_op = functools.partial(self.masked_convolution_op, mask=mask)
        return self.conv2d.call(x)
keras tensorflow convolution machine-learning python
1 Answers

First: The Kernel size

The kernel size for 2D convolution is as follows

[ height, width, input_filters, output_filters ]

The third dimension is of the same size as the input filters. This is critical.

Let's consider how convolution is done manually. Here are the steps:

  • take a patch from the batch of images (BatchSize,height,width,filters)
  • reshape it to [BatchSize, height * width * filters]
  • matrix multiply it by the reshaped kernel [height * width * filters, output_filters]
  • at this point our data is in the shape [BatchSize, output_filters]
  • add in the bias of shape [output_filters] by broadcasting across the whole batch

The output is the filters for each patch.

How to apply a NxN mask to the filter

Given that we know the weights in the convolution are shaped like [ height, width, input_filters, output_filters ] and we want to properly apply a mask of [ height, width ], can can just broadcast that mask like so

masked_weight = weight * mask.reshape([height,width,1,1])

Our Tensorflow keras layer could be written like so

class MaskedConv2D(tf.keras.layers.Layer):
    def __init__(self, *args, **kwargs):
        super(MaskedConv2D, self).__init__()
        self.conv2d = tf.keras.layers.Conv2D(*args, **kwargs)
        
    def build(self, input_shape):
        self.conv2d.build(input_shape[0])
        self._convolution_op = self.conv2d._convolution_op
        
    def masked_convolution_op(self, filters, kernel, mask):
        return self._convolution_op(filters, tf.math.multiply(kernel, tf.reshape(mask, mask.shape + [1,1] )))
        
    def call(self, inputs):
        x, mask = inputs
        self.conv2d._convolution_op = functools.partial(self.masked_convolution_op, mask=mask)
        return self.conv2d.call(x)

and we can test it with the following script

mcon = MaskedConv2D(filters=2,kernel_size=[3,3])

# hack: initialize it by running some data through it
mcon((np.ones([1,4,4,3], dtype=np.float32), tf.constant([[1,1,0],[1,1,1],[0,1,1]], dtype=tf.float32)))

# set all the weights to 1 for testing
mcon.set_weights([ np.ones([3,3,3,2]) , np.zeros([2]) ])

# pass in a matrix of 1s and mask out 2 elements for each input filter
mcon((np.ones([1,4,4,3], dtype=np.float32), tf.constant([[1,1,0],[1,1,1],[0,1,1]], dtype=tf.float32)))

which has the predictable output

<tf.Tensor: shape=(1, 2, 2, 2), dtype=float32, numpy=
    array([[[[21., 21.],
             [21., 21.]],
            [[21., 21.],
             [21., 21.]]]], dtype=float32)>
user_40239510
Related