Экспорт модели с tf.map_fn

Question

У меня есть модель тензорного потока, которую мне нужно экспортировать в сохраненную модель. Ниже приведен упрощенный код модели, которую я пытаюсь экспортировать.

import tensorflow as tf


def foo(x):
  return tf.reduce_sum(x)


inputs = tf.keras.layers.Input(shape=(128,128,3))
y = tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='SAME')(inputs)
y = tf.keras.layers.ReLU()(y)

outputs = tf.map_fn(foo, y, dtype=(tf.float32))

model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
model.save('./export', save_format='tf')

, но при экспорте модели я получаю следующую ошибку.

/Users/bruce/.venv/bin/python /Users/bruce/test_project/mymodel/test.py
Traceback (most recent call last):
  File "/Users/bruce/test_project/mymodel/test.py", line 12, in <module>
    outputs = tf.map_fn(foo, y, dtype=(tf.float32))
  File "/Users/bruce/.venv/lib/python3.6/site-packages/tensorflow_core/python/ops/map_fn.py", line 228, in map_fn
    for elem in elems_flat]
  File "/Users/bruce/.venv/lib/python3.6/site-packages/tensorflow_core/python/ops/map_fn.py", line 228, in <listcomp>
    for elem in elems_flat]
  File "/Users/bruce/.venv/lib/python3.6/site-packages/tensorflow_core/python/ops/tensor_array_ops.py", line 1078, in __init__
    name=name)
  File "/Users/bruce/.venv/lib/python3.6/site-packages/tensorflow_core/python/ops/tensor_array_ops.py", line 716, in __init__
    self._tensor_array = [None for _ in range(size)]
TypeError: 'Tensor' object cannot be interpreted as an integer

Я не могу удалить часть tf.map_fn, которая выполняет важную обработку, которая мне нужна в сохраненной модели при ее развертывании.

Answer 1

Вам необходимо использовать собственный слой:

class MyMapLayer(tf.keras.layers.Layer):
    def __init__(*args, **kwargs)
        super().__init__(*args, **kwargs)

    def foo(self, x):
        return tf.reduce_sum(x)

    def call(self, inputs, **kwargs):
        return tf.map_fn(self.foo, inputs, dtype=(tf.float32))

Затем в вашей модели:

inputs = tf.keras.layers.Input(shape=(128,128,3))
y = tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='SAME')(inputs)
y = tf.keras.layers.ReLU()(y)

outputs = MyMapLayer()(y)

model = tf.keras.models.Model(inputs=inputs, outputs=outputs)