Блог

У меня есть две tf.keras модели:

• Модель обертывает подклассный tf.keras слой (содержащий tf.keras.layers.Conv2D ) tf.keras.layers.TimeDistributed .
• ModelB непосредственно оборачивается tf.keras.layers.Conv2D tf.keras.layers.TimeDistributed .

Насколько я понимаю, эти модели должны быть идентичными. Взглянув на краткое описание модели, мы можем подтвердить, что обе модели имеют одинаковое количество параметров (как и должно быть).

Однако при подсчете количества провалов мы видим, что в приведенном ниже примере у ModelB в 3 раза больше провалов, чем у ModelA. На самом деле, после некоторых проб и ошибок выясняется, что у T/2 x ModelB больше провалов, чем у ModelA, где T -количество временных шагов в распределенном времени.

Я что-то неправильно понимаю здесь? Может быть, ошибка в том, как я вычисляю провалы, или что-то не так с распределением времени?

Код для воспроизведения приведен ниже.

 import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import  convert_variables_to_constants_v2_as_graph

def get_flops(model, inputs):
  # Adapted from https://github.com/tensorflow/tensorflow/issues/32809#issuecomment-768977280
  input_signature = tf.TensorSpec([1, *inputs.shape[1:]])
  concrete = tf.function(lambda inputs: model(inputs))
  concrete_func = concrete.get_concrete_function(input_signature)
  frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(concrete_func)
  with tf.Graph().as_default() as graph:
      tf.graph_util.import_graph_def(graph_def, name='')
      run_meta = tf.compat.v1.RunMetadata()
      opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
      flops = tf.compat.v1.profiler.profile(graph=graph, run_meta=run_meta, cmd="op", options=opts)
      return flops.total_float_ops

class Layer2D(tf.keras.layers.Layer):
  def __init__(self, filters):
    super(Layer2D, self).__init__()
    self.conv2d = tf.keras.layers.Conv2D(
      filters=filters, kernel_size=(3, 3), padding='same')
  def call(self, inputs, training=False):
    """Input shape: [B, H, W, C]"""
    return self.conv2d(inputs)

class ModelA(tf.keras.Model):
  """Model which wraps subclassed Keras layer in TimeDistributed"""
  def __init__(self, filters):
    super(ModelA, self).__init__()
    self.time_dist_layer_2d = tf.keras.layers.TimeDistributed(Layer2D(filters))
  def call(self, inputs, training=False):
    """Input shape: [B, T, H, W, C]"""
    return self.time_dist_layer_2d(inputs)

class ModelB(tf.keras.Model):
  """Model which directly wraps tf.keras.layers.Conv2D in TimeDistributed"""
  def __init__(self, filters):
    super(ModelB, self).__init__()
    self.conv2d = tf.keras.layers.TimeDistributed(
      tf.keras.layers.Conv2D(
        filters=filters, kernel_size=(3, 3), padding='same'))
  def call(self, inputs, training=False):
    """Input shape: [B, T, H, W, C]"""
    return self.conv2d(inputs)

model_a = ModelA(8)
model_b = ModelB(8)

inputs = tf.random.uniform([4, 6, 224, 224, 3], dtype=tf.float32)

# Build models by calling them
model_a(inputs)
model_b(inputs)

model_a.summary()
model_b.summary()

print("model_a flops: {}".format(get_flops(model_a, inputs)))
print("model_b flops: {}".format(get_flops(model_b, inputs)))