Как воспроизвести блоки узких мест в Mobil enet V3 с Keras API?

Question

Используя Keras API, я пытаюсь написать MobilenetV3, как описано в этой статье: https://arxiv.org/pdf/1905.02244.pdf с архитектурой, описанной на этом рисунке:

Для этого мне нужно реализовать bottloneck_blocks из предыдущей статьи https://arxiv.org/pdf/1801.04381.pdf. См. Изображение для архитектуры:

Мне удалось склеить начальный и конечный слои Conv:

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()

Где bottleneck_block указано в следующем фрагменте кода (изменено с 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])

Однако в bneck2 появляется следующая ошибка:

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

Я знаю, что ошибка означает, что измерение входов и выходов отключено, но я не знаю, как это исправить, чтобы структурировать сеть как MobileNetV3.

Чего мне здесь не хватает?

Для справки, вот исходный код в тензорном репо для той же сети: https://github.com/tensorflow/models/blob/a174bf5b1db0e2c1e04697ff5aae5182bd1c60e7/research/slim/nets/mobilenet/mobilenet_v3.py#L130

Answer 1

Решение состоит в том, чтобы изменить bottleneck_block, как описано в репозитории V3 автора :

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

Проверка размера и шага предотвращает ошибку, которую я изначально получил при добавлении двух сети, которые не соответствуют размеру

Answer 2

В ваших нижних слоях есть Add ().

Теперь, Add ожидает два тензора одинаковой формы. Но так как вы пропустили так много слоев, когда эта строка запущена, tf.keras.layers.Add()([m, x]) - m и x имеют разные размеры.

Итак, либо спроектируйте меньшую сеть с меньшим количеством слоев, либо просто внедрите все промежуточные слои .