使用 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)

使用 tf.map_fn 导出模型

Exporting a model with tf.map_fn

python

tensorflow

tensorflow2.x