如何在Tensorflow 2.0中创建serving_input_fn进行图像预处理?

时间:2019-04-25 20:57:26

标签: tensorflow2.0

我正在使用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(),并使用了 tf.placeholder ,它似乎不再存在于Tensorflow 2.0中。

https://www.tensorflow.org/alpha/guide/saved_model上显示了Tensorflow 2.0中 serving_input_fn 的示例。由于我使用的是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?

2 个答案:

答案 0 :(得分:3)

升级时,我遇到了类似的问题。似乎在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)

答案 1 :(得分:1)

扩展和简化@ 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
    )

注意:dtype=tf.float32函数中的map很重要,没有它就无法工作。我找到了解决方法here。我还通过简单地添加装饰器简化了具体的功能工作(有关详细信息,请参见this)。