zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

SavedModel TF Serve Fix (#7228) 

* SavedModel TF Serve Fix

Fix for https://github.com/ultralytics/yolov5/issues/7205 proposed by @tylertroy

* Update export.py
This commit is contained in:
Glenn Jocher 2022-04-01 00:07:23 +02:00 committed by GitHub
parent 4d157f578a
commit 734ab033fd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
export.py
View File

@ -285,12 +285,12 @@ def export_saved_model(model,
if keras: if keras:
keras_model.save(f, save_format='tf') keras_model.save(f, save_format='tf')
else: else:
m = tf.function(lambda x: keras_model(x)) # full model
spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype) 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) m = m.get_concrete_function(spec)
frozen_func = convert_variables_to_constants_v2(m) frozen_func = convert_variables_to_constants_v2(m)
tfm = tf.Module() 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) tfm.__call__(im)
tf.saved_model.save(tfm, tf.saved_model.save(tfm,
f, f,