zeh
zeh

Reputation: 1317

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:

enter image description here

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

enter image description here

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

Upvotes: 0

Views: 1366

Answers (2)

Zabir Al Nazi Nabil
Zabir Al Nazi Nabil

Reputation: 11218

In your bottolneck layers, there are Add() ops.

Now, Add expects two tensors with the same shape. But, as you have skipped so many layers when this line is run, tf.keras.layers.Add()([m, x]) - m and x have different dimensions.

So, either design a smaller network with fewer layers or just implement all of the intermediate layers.

Upvotes: 1

zeh
zeh

Reputation: 1317

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

Upvotes: 1

Related Questions