How to reproduce the Bottleneck Blocks in Mobilenet V3 with Keras API?

Using Keras API, I am trying to write the MobilenetV3 as explained in this article: https://arxiv.org/pdf/1905.02244.pdf with the architecture as described in this picture:

For that, I need to implement the bottloneck_blocks from the previous article https://arxiv.org/pdf/1801.04381.pdf. See image for architecture:

I managed to glue together the Initial and final Conv layers:

from tensorflow.keras.layers import Input, Conv2D, Add, AvgPool2D, UpSampling2D


first_input = Input(shape=(256, 256, 3))
firt_conv   = Conv2D(16,3, strides=2, name="FirstConv2d", padding="same")(first_input)

bneck1  = add_bottleneck_block(firt_conv, 16, 16)
bneck2  = add_bottleneck_block(bneck1, 64, 24, strides=2)

#... Skiping all the other BottleNeck Blocks for simplicity 

lastBneck      = add_bottleneck_block(second2LastBneck, 960, 160, bneck_depth=5)
middleConv     = Conv2D(160, 1 , strides=1, name="MiddleConv", )(bneck3)
pool7          = AvgPool2D(7, strides=1, padding='same', name="7x7Pool")(middleConv)
SecondLastConv = Conv2D(1280, 1, strides=1, name="SecondLastConv")(pool7)
lastConv       = Conv2D(3,1, strides=1, name="lastConv1x1")(SecondLastConv)
upScale        = UpSampling2D(2)(lastConv) # This layer is application specific for my training.


v3 = tf.keras.models.Model(inputs=[first_input], outputs=upScale)

v3.compile(optimizer='adam', loss=tf.keras.losses.BinaryCrossentropy(),)
v3.summary()

Where the bottleneck_block is given in the next snippet of code (modified from https://towardsdatascience.com/mobilenetv2-inverted-residuals-and-linear-bottlenecks-8a4362f4ffd5)

def bottleneck_block(x, expand=64, squeeze=16, strides=1, bneck_depth=3):
  """
  Bottleneck block with Activation and batch normalization commented since
  I don't believe this is the issue in my problem
  """
  m = tf.keras.layers.Conv2D(expand, (1,1), strides=1)(x)
  #m = tf.keras.layers.BatchNormalization()(m)
  #m = tf.keras.layers.Activation('relu6')(m)
  m = tf.keras.layers.DepthwiseConv2D(bneck_depth, padding='same', strides=strides)(m)
  #m = tf.keras.layers.BatchNormalization()(m)
  #m = Activation('relu6')(m)
  m = tf.keras.layers.Conv2D(squeeze, (1,1), strides=1)(m)
  #m = tf.keras.layers.BatchNormalization()(m)
  return tf.keras.layers.Add()([m, x])

However, in bneck2 I get the following error:

ValueError: Operands could not be broadcast together with shapes (16, 16, 24) (128, 128, 16)

I know the error means the dimension of the inputs and outputs are off, but I don’t know how to fix it to structure the network as the MobileNetV3.

What am I missing here?

For reference, here is source code in the tensorflow repo for the same network: https://github.com/tensorflow/models/blob/a174bf5b1db0e2c1e04697ff5aae5182bd1c60e7/research/slim/nets/mobilenet/mobilenet_v3.py#L130

Answer

The Solution is to modify the bottleneck_block as described in the V3 author’s repo:

import tensorflow as tf
def bottleneck_block(x, expand=64, squeeze=16, strides=1, bneck_depth=3, se=False):
  """
  se stands for squeeze_excite
  """

  m = tf.keras.layers.Conv2D(expand, (1,1), strides=1)(x)
  m = tf.keras.layers.BatchNormalization()(m)
  #m = tf.keras.layers.Activation('relu6')(m)
  m = tf.keras.layers.DepthwiseConv2D(bneck_depth, padding='same', strides=strides)(m)
  m = tf.keras.layers.BatchNormalization()(m)
  #m = Activation('relu6')(m)
  if se:
    m = squeeze_excite_block(m, ratio=4)
  m = tf.keras.layers.Conv2D(squeeze, (1,1), strides=1, padding='same')(m)
  m = tf.keras.layers.BatchNormalization()(m)

  if (
    # stride check enforces that we don't add residuals when spatial
    # dimensions are None
    strides == 1 and
    # Depth matches
    m.get_shape().as_list()[3] == x.get_shape().as_list()[3]
  ):
    m = tf.keras.layers.Add()([m, x])

  return m

The check in dimension and stride prevents the error I initially got when adding two nets that do not match the dimension