How to create a serving_input_fn in Tensorflow 2.0 for image preprocessing? - tensorflow2.0

Question

I am using Tensorflow 2.0 and am able to train a CNN for image classification of 3-channel images. I perform image preprocessing within the data input pipeline (shown below) and would like to include the preprocessing functionality in the served model itself. My model is served with a TF Serving Docker container and the Predict API.
The data input pipeline for training is based on the documentation at https://www.tensorflow.org/alpha/tutorials/load_data/images.
My pipeline image preprocessing function is 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
With lists of image paths, the pipeline prefetches and performs image preprocessing using tf functions within load_and_preprocess_from_path_label:
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)
Previous Tensorflow examples I've found use serving_input_fn(), and utilized tf.placeholder which seems to no longer exist in Tensorflow 2.0.
An example for serving_input_fn in Tensorflow 2.0 is shown on https://www.tensorflow.org/alpha/guide/saved_model. Since I am using the Predict API, it looks like I would need something similar to:
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)
Ideally, the served model would accept a 4D Tensor of 3-channel image samples of any size and would perform the initial image preprocessing on them (decode image, normalize, convert to HSV, and resize) before classifying.
How can I create a serving_input_fn in Tensorflow 2.0 with a preprocessing function similar to my load_and_preprocess_path function?

Answer 1

I faced a similar issue when upgrading. It appears that the way to achieve this in Tensorflow 2 is to provide a function which the saved model can use to make the predictions, something like:
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)

Answer 2

Expanding and simplifying #harry-salmon answer. For me the following worked:
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
)
Note: dtype=tf.float32 in map function was important and didn't work without it. I found solution here. Also I simplified the concrete function work by simply adding a decorator (see this for details).