tensorflow2.0 - 如何在 Tensorflow 2.0 中创建serving_input_fn 进行图像预处理？

Question

我正在使用 Tensorflow 2.0 并且能够训练 CNN 进行 3 channel 图像的图像分类。我在数据输入管道(如下所示)中执行图像预处理，并希望在所服务的模型本身中包含预处理功能。我的模型使用 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 示例使用了 services_input_fn()，并使用了 tf.placeholder，它似乎在 Tensorflow 2.0 中不再存在。

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 channel 图像样本的 4D 张量，并在分类之前对它们执行初始图像预处理(解码图像、标准化、转换为 HSV 和调整大小)。

如何在 Tensorflow 2.0 中使用类似于我的 load_and_preprocess_path 函数的预处理函数创建一个serving_input_fn？

Answer 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)