TFlite NMS

[oranmichel]

@zldrobit
Hello,
what is the current reason that we are not able to add the NMS inside the TFlite graph?
i saw that there is a (TFL::NonMaxSuppressionV5Op)
thanks, Oran

[zldrobit]

@oranmichel
I have experimented TFLite models with two kinds of NMS:

• Class-aware NMS, namely tf.image.combined_non_max_suppression, cannot be transformed to TFLite without Flex ops.
• Class-agnostic NMS, namely tf.image.non_max_suppression, does not support int8 quantization.
• Class-agnostic NMS with padded output, namely tf.image.non_max_suppression_padded, has shuffled output order.

Therefore, adding NMS won't satisfy all TFLite use cases (without Flex ops and/or int8 quantization support). The support of NMS in TFLIte/TF.js is summarized in the following table.

op	TFLite support	TFLite GPU delegate support	TF.js support	int8 quantization support
NMS	indirectly support1	not supported due to dynamic-sized output	OK	No
CNMS	flex required	not tested	not supported	not tested
NMSP	shuffled output order2	not supported due to dynamic-sized tensors3	wrong result	not tested

• 1: tf.image.non_max_suppression does not return a dynamic valid detection number after converting to TFLite models tensorflow/tensorflow#51629
• 2: Convert to tflite will change output order for models with multiple outputs tensorflow/tensorflow#33303
• 3: tf.image.non_max_suppression_padded returns dynamic-sized tensors: fails on TFlite GPU Delegate tensorflow/tensorflow#52388
• NMS: tf.image.non_max_suppression
• CNMS: tf.image.combined_non_max_suppression
• NMSP: tf.image.non_max_suppression_padded

EDIT: The experimented TensorFlow version is 2.6.0.

[zldrobit]

You could try the experimental branch https://github.com/zldrobit/yolov5/commits/tf-android-flex-ops and refer to https://www.tensorflow.org/lite/guide/ops_select as a manual.
Yet, I still cannot export int8 quantized TFLite model with Flex ops.

[oranmichel]

@zldrobit
Hi, did you manage to run the model with the nms inside the graph on the android gpu? or just on the cpu.
thanks, Oran

[zldrobit]

@oranmichel Lately, I am able to integrate tf.image.combined_non_max_suppression in YOLOv5 TFLite model.
The fp16 model works with CPU, but does not output any bounding boxes with GPU delegate.
And I failed to convert YOLOv5 to an int8 quantized model:

TFLite export failure: Model output is not dequantized.
Traceback (most recent call last):
  File "models/tf.py", line 479, in <module>
    tflite_model = converter.convert()
  File "/opt/conda/lib/python3.8/site-packages/tensorflow/lite/python/lite.py", line 873, in convert
    return super(TFLiteKerasModelConverterV2,
  File "/opt/conda/lib/python3.8/site-packages/tensorflow/lite/python/lite.py", line 637, in convert
    result = _modify_model_io_type(result, **flags_modify_model_io_type)
  File "/opt/conda/lib/python3.8/site-packages/tensorflow/lite/python/util.py", line 843, in modify_model_io_type
    _modify_model_output_type(model_object, inference_output_type)
  File "/opt/conda/lib/python3.8/site-packages/tensorflow/lite/python/util.py", line 765, in _modify_model_output_type
    raise ValueError("Model output is not dequantized.")
ValueError: Model output is not dequantized.

You could try it with the tf-android-flex-ops branch.

I also experimented with the class-agnostic tf.image.non_max_suppression, but it requires input box with 2 dimension, which won't run without changing the input and output of YOLOv5:

TensorFlow saved_model export failure: boxes must be 2-D[1,6300,1,4] [Op:NonMaxSuppressionV3]
Traceback (most recent call last):
  File "models/tf.py", line 410, in <module>
    y = tf_model.predict(img)
  File "models/tf.py", line 330, in predict
    ind = tf.image.non_max_suppression(
  File "/opt/conda/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py", line 201, in wrapper
    return target(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/tensorflow/python/ops/image_ops_impl.py", line 3553, in non_max_suppression
    return gen_image_ops.non_max_suppression_v3(boxes, scores, max_output_size,
  File "/opt/conda/lib/python3.8/site-packages/tensorflow/python/ops/gen_image_ops.py", line 2733, in non_max_suppression_v3
    _ops.raise_from_not_ok_status(e, name)
  File "/opt/conda/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 6862, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: boxes must be 2-D[1,6300,1,4] [Op:NonMaxSuppressionV3]

Hence, tf.image.non_max_suppression does not support batch inference.
The code is in tf-android-flex-ops-class-agnostic-nms.

[glemarivero]

Do you have a working example with tflite with nms and fp16?
Thanks!

[zldrobit]

@glemarivero You could export an fp16 TFLite model with nms using

python3 models/tf.py --weights weights/yolov5s.pt --cfg models/yolov5s.yaml --img 320 --tf-nms --agnostic-nms

However, it cannot be used with GPU delegate since the nms operator tf.non_max_suppression returns dynamic size tensors.
In order to get static-sized tensors, I tried tf.non_max_suppression_padded and couldn't get a correct result after conversion.
So the best I can get is a fp16 TFLite model with nms running in CPU. You could see #1127 for more details about TFLite conversion.

The Android project in https://github.com/zldrobit/yolov5/tree/tf-android does not support TFLite models with nms, because users would prefer fp16 w/o nms to fp16 w/ nms for GPU acceleration.