如何在 Tensorflow 2.0 中创建 serving_input_fn 用于图像预处理？答案

【问题标题】：How to create a serving_input_fn in Tensorflow 2.0 for image preprocessing?如何在 Tensorflow 2.0 中创建 serving_input_fn 用于图像预处理？
【发布时间】：2019-04-25 20:57:25
【问题描述】：

我正在使用 Tensorflow 2.0，并且能够训练 CNN 来对 3 通道图像进行图像分类。我在数据输入管道（如下所示）中执行图像预处理，并希望在服务模型本身中包含预处理功能。我的模型使用 TF Serving Docker 容器和 Predict API。

用于训练的数据输入管道基于https://www.tensorflow.org/alpha/tutorials/load_data/images 的文档。

我的管道图像预处理函数是load_and_preprocess_from_path_label：

def load_and_preprocess_path(image_path):

    # Load image
    image = tf.io.read_file(image_path)
    image = tf.image.decode_png(image)

    # Normalize to [0,1] range
    image /= 255

    # Convert to HSV and Resize
    image = tf.image.rgb_to_hsv(image)
    image = tf.image.resize(image, [HEIGHT, WIDTH])

    return image

def load_and_preprocess_from_path_label(image_path, label):

    return load_and_preprocess_path(image_path), label

通过图像路径列表，管道使用 load_and_preprocess_from_path_label 中的 tf 函数预取并执行图像预处理：

all_image_paths, all_image_labels = parse_labeled_image_paths()
x_train, x_test, y_train, y_test = sklearn.model_selection.train_test_split(all_image_paths, all_image_labels, test_size=0.2)

# Create a TensorFlow Dataset of training images and labels
ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
image_label_ds = ds.map(load_and_preprocess_from_path_label)

BATCH_SIZE = 32
IMAGE_COUNT = len(all_image_paths)

ds = image_label_ds.apply(tf.data.experimental.shuffle_and_repeat(buffer_size=IMAGE_COUNT))
ds = ds.batch(BATCH_SIZE)
ds = ds.prefetch(buffer_size=AUTOTUNE)

# Create image pipeline for model
image_batch, label_batch = next(iter(ds))
feature_map_batch = model(image_batch)

# Train model
model.fit(ds, epochs=5)

我发现以前的 Tensorflow 示例使用 serving_input_fn()，并使用了在 Tensorflow 2.0 中似乎不再存在的 tf.placeholder。

Tensorflow 2.0 中 serving_input_fn 的示例显示在 https://www.tensorflow.org/alpha/guide/saved_model 上。由于我使用的是 Predict API，看起来我需要类似于：

serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(...)

# Save the model with the serving preprocessing function
model.export_saved_model(MODEL_PATH, serving_input_fn)

理想情况下，服务模型将接受任何大小的 3 通道图像样本的 4D 张量，并在分类之前对它们执行初始图像预处理（解码图像、归一化、转换为 HSV 和调整大小）。

如何在 Tensorflow 2.0 中创建具有类似于我的 load_and_preprocess_path 函数的预处理函数的 serving_input_fn？

【问题讨论】：

标签： tensorflow2.0

【解决方案1】：

我在升级时遇到了类似的问题。看来，在 Tensorflow 2 中实现这一点的方法是提供一个函数，保存的模型可以用来进行预测，例如：

def serve_load_and_preprocess_path(image_paths: tf.Tensor[tf.string]):
    # loaded images may need converting to the tensor shape needed for the model        
    loaded_images = tf.map_fn(load_and_preprocess_path, image_paths, dtype=tf.float32)
    predictions = model(loaded_images)
    return predictions


serve_load_and_preprocess_path = tf.function(serve_load_and_preprocess_path)
serve_load_and_preprocess_path = serve_load_and_preprocess_path.get_concrete_function(
    image_paths=tf.TensorSpec([None,], dtype=tf.string))

tf.saved_model.save(
    model,
    MODEL_PATH,
    signatures=serve_load_and_preprocess_path
)

# check the models give the same output
loaded = tf.saved_model.load(MODEL_PATH)
loaded_model_predictions = loaded.serve_load_and_preprocess_path(...)
np.testing.assert_allclose(trained_model_predictions, loaded_model_predictions, atol=1e-6)

【讨论】：

感谢您的回复！我收到ValueError: Trying to read from list with wrong element dtype. List has type float but expected type string for 'map/TensorArrayV2Stack/TensorListStack' (op: 'TensorListStack') with input shapes: [], [3]. 有什么想法吗？谢谢！
哦好的，你还需要专门说loaded_images = tf.map_fn(load_and_preprocess_path, image_paths, dtype=tf.float32)。见github.com/tensorflow/tensorflow/issues/31055
谢谢，我会更新我的答案以包含 dtype。这最终对你有用吗？
是的，工作。我在下面发布了我的答案。它稍微简化了您的解决方案。

【解决方案2】：

扩展和简化@harry-salmon 的答案。对我来说，以下工作：

def save_model_with_serving_signature(model, model_path):
    @tf.function(input_signature=[tf.TensorSpec(shape=[None, ], dtype=tf.string)])
    def serve_load_and_preprocess_path(image_paths):
        return model(tf.map_fn(load_and_preprocess_path, image_paths, dtype=tf.float32))

    tf.saved_model.save(
        model,
        model_path,
        signatures=serve_load_and_preprocess_path
    )

注意：map 函数中的 dtype=tf.float32 很重要，没有它就无法工作。我找到了解决方案here。我还通过简单地添加一个装饰器来简化具体的功能工作（详见this）。

【讨论】：