Custom convolutions and none type object in keras custom layer for gating operation

Question

I am trying to implement gated pooling layer in keras and this will need to first find the max and average pooling of the layer. After that it computes a region wise parameter using the same trainable mask across all the depth dimensions.

For example if input is size (batch, 28, 28, 6), the max and average pooling will return (batch, 14, 14, 6) with strides (2,2). The gating operation should return (batch, 14, 14, 6) after a dot product of the mask ( of shape (2,2) ) with the input region and then find the weighted sum of the max and average pooling. But I am unable to do this operation because I cannot declare the output dimension with 'None' type for batch.

I have multiple threads and various operations. The reshape operation is not working in the custom layer and adding it as a separate layer in the model also throws up errors. I have tried many other things as well but I am unable to handle this problem. Some help would be greatly appreciated.

class Gated_pooling(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(Gated_pooling, self).__init__(**kwargs)

  def build(self, input_shape):
    self.batch, self.row, self.col, self.channel = input_shape.as_list()
    self.output_size = self.row//2
    self.mask = self.add_weight(name='mask', shape=(2,2,1,1),
                               initializer='truncated_normal',
                               trainable=True
                               )
    self.maxpool = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2), padding='VALID')
    self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(2, 2), strides=(2,2), padding='VALID')
    super(Gated_pooling, self).build(input_shape)

  def call(self, x):
    x1 = self.maxpool(x)
    x2 = self.maxpool(x)
    xs = tf.zeros(x1.shape)
    i=0
    for c in tf.split(x, self.channel, 3):
      xs[:, :, i] = tf.nn.conv2d(input=c, filters=self.mask, strides=[2,2,1,1], padding='VALID')
      i+=1

    #z = tf.math.sigmoid(tf.reshape(xs, [self.batch, self.output_size, self.output_size, self.channel]))
    z = tf.math.sigmoid(xs)
    output = tf.add(tf.multiply(z, x1), tf.multiply((1-z), x2))
    #return tf.reshape(output, [self.batch, self.output_size, self.output_size, self.channel])
    return output

Answer 1

this seems to work...

class Gated_pooling(tf.keras.layers.Layer):

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

        self.mask = self.add_weight(name='mask', shape=(2,2,1,1),
                                   initializer='truncated_normal',
                                   trainable=True
                                   )
        self.maxpool = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2), padding='VALID')
        self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(2, 2), strides=(2,2), padding='VALID')

    def call(self, x):

        self.batch, self.row, self.col, self.channel = x.shape
        self.output_size = self.row//2

        x1 = self.maxpool(x)
        x2 = self.maxpool(x)
        xs = []

        for c in tf.split(x, self.channel, 3):
            xs.append(tf.nn.conv2d(c, filters=self.mask, strides=[2,2], padding='VALID'))

        xs = tf.concat(xs, axis=-1)

        z = tf.math.sigmoid(xs)
        output = tf.add(tf.multiply(z, x1), tf.multiply((1-z), x2))

        return output

usage:

X = np.random.uniform(0,1, (32,10,10,3)).astype('float32')

gp = Gated_pooling()
gp(X)