Question

我正在用keras（带有tensorflow后端）编写两个模型，这些模型需要多次转换图像。现在，经过一些搜索，我想出了一种有效的方法。基本上归结为使用tf.map_fn将tf.contrib.image.translate映射到一批图像上，如以下代码所示：

import tensorflow as tf
import tensorflow.keras as keras

def sample(inputs):
    """ Creates a random direction for every image in the batch."""
    batch_size = tf.shape(inputs)[0]
    eps = tf.random.normal((batch_size, 2))
    return eps


def translate_single_set(args):
    """
    Translates a single fixed set of images. Used as callable in tf.map_fn.
    args should be [images, translations]
    """
    assert isinstance(args, list)
    assert len(args) == 2
    return tf.contrib.image.translate(args[0], args[1])


translator_inputs = keras.layers.Input(shape=(28, 28))
shift = keras.layers.Lambda(sample)(translator_inputs)

translated = keras.layers.Lambda(
    lambda args: tf.map_fn(translate_single_set, args, dtype=tf.float32)
)([translator_inputs, shift])

translator = keras.Model(inputs=translator_inputs, outputs=translated)


# use this on mnist
from tensorflow.keras.datasets import mnist
(original_train, _), (original_test, _) = mnist.load_data()

original_train = original_train.astype('float32')/255
original_test = original_test.astype('float32')/255

binary_train = np.round(original_train)
binary_test = np.round(original_test)
# Shuffle the data
np.random.shuffle(binary_train)
np.random.shuffle(binary_test)

translator.predict(binary_train, batch_size=100)

问题是它非常慢。仅当使用cpu时，最后一行执行需要几秒钟，而在gpu上则需要数十秒。由于我的模型具有很多可训练的权重，所以我确实很想使用GPU。但是，我的第一个模型在翻译图像时比不翻译时慢了十到二十倍（在gpu上，在cpu上只慢了五到七倍）。第二个模型尚未完成，但需要比第一个模型多25倍的翻译。

是否有一种方法可以更有效地翻译（批量）图像？我是否认为主要瓶颈可能在于使用tf.map_fn是正确的吗？（我不知道如何检查）

如何使Tensorflow / Keras中的图像转换效率更高

0 个答案: