SavedModel TF Serve Fix

[glenn-jocher]

Fix for #7205 proposed by @tylertroy

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Enhanced TensorFlow SavedModel export functionality with NMS (Non-Max Suppression) support.

📊 Key Changes

• 📝 Rearranged the order of some lines of code for better readability.
• ✨ Modified the TensorFlow export code to support exporting models with NMS directly incorporated for TensorFlow SavedModels.

🎯 Purpose & Impact

• The main purpose of these changes is to improve the TensorFlow model export process by including NMS within the SavedModel itself.
• Users who rely on TensorFlow for deploying models will benefit from this, as it simplifies the deployment pipeline by embedding crucial post-processing steps (like NMS) into the model, reducing the need for additional code after loading the model.
• The impact of this change streamlines model deployment, potentially improving performance and stability since the model export is more aligned with TensorFlow's execution framework. 🚀

[glenn-jocher]

@tylertroy your proposed fix breaks existing workflows, i.e. SavedModel inference with detect.py:

!git clone https://github.com/ultralytics/yolov5 -b fix/tf_serve
%cd yolov5

!python export.py --include saved_model
!python detect.py --weights yolov5s_saved_model

EDIT: Since all TF exports rely on SavedModel as a prerequisite they are all probably also impacted. So this small change has knockon impacts down the export and inference workflows that would need to be resolved for before update could be considered for merging with master.

[tylertroy]

@glenn-jocher As I understand it, any time the --nms flag is applied, then the saved_model no longer works with detect.py. I understood my change to only effect output when --nms is applied. Not so?

Edit: Now I understand the issue you describe. Perhaps an if statement can clear this up.

diff --git a/export.py b/export.py
index 7517dc4..0c44cc0 100644
--- a/export.py
+++ b/export.py
@@ -276,7 +276,7 @@ def export_saved_model(model, im, file, dynamic,
            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], [spec]) if tf_nms else tf.function(lambda x: frozen_func(x)[0], [spec])
            tfm.__call__(im)
            tf.saved_model.save(
                tfm,

[glenn-jocher]

@tylertroy yes it's true that NMS models are no longer detect.py compatible. It seems like the proposed update (or at least my implementation here) affects SavedModels regardless of the --nms flag though.

Maybe the solution is an if nms: ... else: ... statement there to only apply the change to --nms exports?

[tylertroy]

@glenn-jocher
Great idea, I also made this suggestion in my above comment as an edit. Probably should have just made a new message! See above.

[glenn-jocher]

@tylertroy PR is merged. Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐