0

我有一些张量流模型，我需要在保存的模型中导出。下面是我试图导出的模型的简化代码。

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在保存的模型中进行一些我需要的基本处理的部分。

4

1 回答 1

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)
于 2020-05-06T07:23:25.347 回答