Python TF map_fn太慢

我想在tensorflow中创建一个自定义层，它应该对传入的tensor应用函数

f

。因此，如果批次包含张量

T=[T1，T2，…，Tn]

它应该返回张量

[f（T1），f（T2），…，f（Tn）]

这样做的预期方法似乎是使用

tf.map\u fn

函数。然而，我注意到这个函数非常慢。下面是一个MWE，它在我的笔记本电脑上产生以下性能：

• ~61us/步，无标识层
• ~62us/带有标识层的步骤

  输出=输入

• ~120us/步，带标识层

  output=tf.map\u fn（…）

有没有办法加快批量大小的迭代

import tensorflow as tf
mnist = tf.keras.datasets.mnist
(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

class Identity(tf.keras.layers.Layer):
    def __init__(self,  **kwargs):
        super(Identity, self).__init__(**kwargs)

    def call(self, inputs):
        output = tf.map_fn(lambda x: x, inputs)
#        output = inputs
        return output  

    def compute_output_shape(self, input_shape):
        return input_shape

model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(512, activation=tf.nn.relu),
        Identity(),
        tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5, batch_size=100)

---

不，不是那样，

tf.map\u fn

很慢，正如您在绝对最小的示例中所看到的，

tf.while\u loop

更慢。您可以更改

parallel_iterations

参数，但它仍然必须单独操作每个元素。你必须尽可能使用矢量化的方式，使用作用于整个阵列的TensorFlow操作。（公平地说，120us实际上非常快，我的意思是

tf。map\u fn

比矢量化的方法慢，只要存在一个）在你的案例中，输入的形状是什么？@MikhailBerlinkov我正在修改标准的退出层。输入形状是前一层的结果。好吧，我认为加速计算的唯一方法是通过矢量化（并使用现有的有效操作）或实现您自己的低级并行化代码。