From 7aed8001c83f65cfbc959750455ebb0ab00d695d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 1 Apr 2022 00:07:23 +0200 Subject: [PATCH] SavedModel TF Serve Fix (#7228) * SavedModel TF Serve Fix Fix for https://github.com/ultralytics/yolov5/issues/7205 proposed by @tylertroy * Update export.py --- export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/export.py b/export.py index e146dad42980..cc7a74db9af2 100644 --- a/export.py +++ b/export.py @@ -285,12 +285,12 @@ def export_saved_model(model, if keras: keras_model.save(f, save_format='tf') else: - m = tf.function(lambda x: keras_model(x)) # full model spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype) + m = tf.function(lambda x: keras_model(x)) # full model m = m.get_concrete_function(spec) frozen_func = convert_variables_to_constants_v2(m) tfm = tf.Module() - tfm.__call__ = tf.function(lambda x: frozen_func(x)[0], [spec]) + tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x)[0], [spec]) tfm.__call__(im) tf.saved_model.save(tfm, f,