Convert Mask R-CNN model to TFLite with Tensorflow 2.3

Mask R-CNN is one of the important models in the object detection world. It was published in 2018 and it has multiple implementations based on Pytorch (detectron2) and Tensorflow (object detection).

In this quick tutorial, we will explore how we can export Mask R-CNN to tflite so that it can be used on mobile devices such as Android smartphones. We are going to use leekunhee/Mask_RCNN version of Mask R-CNN which is a fork of the greatest implementation done by matterport/Mask_RCNN. The latter implementation is based on Keras and Tensorflow 1.13. While the fork done by leekunhee is a porting to Tensorflow 2.0.

We will start by describing the minor changes we have done to keep everything updated and compatible with Tensorfow 2.3 and after that, we will explain the steps to export the model to tflite and make it ready for mobile devices deployment.

Migrate functions to Tensorflow 2.3.0

Moving to Tensorflow 2.3.0 implies using the new set of functions provided by the new API. Most of the work was done by leekunhee, but we noticed that few of the legacy code remains and that we can update it.

The main file that defines the model is located at mrcnn/model.py. Here is the list of the changed things:

At the beginning of the file, we can see this line tf.compat.v1.disable_eager_execution() which is, actually, not needed. Eager execution is one of the greatest features that TF enabled by default starting from version 2.0.

In that same file, we can see the use of some functions coming from tf.compat.v1 such as tf.compat.v1.where so we replaced it with tf.where.

These are the small changes that we did on the initial model. Let’s go to the TFLite conversion part.

Convert Keras Model to TFLite

Tensorflow supports converting Keras models very easily. According to the documentation, all you have to do is to use tf.lite.TFLiteConverter.from_keras_model and given the model, it will convert it to TFLite.

We used Netron which is an amazing tool that helped us in debugging the model and knowing where to look and what to change.

So since we have a Keras Model, we tried to follow the documentation of Tensorflow. We recommend you to open samples/demo.ipynb notebook and to run it first as it is.

# After running a first Keras inference as defined in demo.ipynb

# Create the TFLite converter with the keras model
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)

# Let’s extract the keras model
keras_model = model.keras_model

tflite_model = converter.convert()

Hmm … This seems to run forever and never ends. At this point, we didn’t know what to do and after multiple reading of the documentation, we decided to allow custom ops on tflite model. Therefore, we added the following line:

# After running a first Keras inference as defined in demo.ipynb

# Let’s extract the keras model
keras_model = model.keras_model

# Create the TFLite converter with the keras model
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
converter.allow_custom_ops=True

tflite_model = converter.convert()

This code works and gives us a tflite model. When we try to run an inference using the tflite model and tf.lite.Interpreter we get the following error:

RuntimeError: Encountered unresolved custom op: NonMaxSuppressionV3.Node number 323 (NonMaxSuppressionV3) failed to prepare .

This error is coming from Node number 323 in the graph. We used Netron to visualize the tflite model and find that node:

Netron visualization of NonMaxSuppressionV3 in the graph

After looking at the documentation of the supported operations on TFLite, we noticed that NonMaxSuppression is a supported operation in TF_OPS, which is good news. This time, we added a few other parameters to TFLiteConverter so the conversion code became:

# After running a first Keras inference as defined in demo.ipynb

# Let’s extract the keras model
keras_model = model.keras_model

# Create the TFLite converter with the keras model
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)

converter.allow_custom_ops=True
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
]
converter.optimizations = [tf.lite.Optimize.DEFAULT]

tflite_model = converter.convert()

The conversion went well but when we tried to run an inference, we had a new exception which is saying that the operation tf.sets.intersect is not supported. So we had to reimplement that operation which is basically a simple operation that gets two tensors of different sizes and returns the intersection between the two sets. We used the solution suggested by vijay m on Stackoverflow, which is based on using operations such as sub, abs, and sum and then using array broadcasting to get the right result. This solution uses simple operations that are supported by TFLite.

And finally, after multiple attempts and many hours of tedious work, we succeeded in converting Mask-RCNN to TFLite. You can find in this colab notebook the reproduction of all of this.

What is next?

• Load the model in an Android device.
• Do some benchmarking.

Special thanks

• Special thanks to Kartoue Mady Demdah <kdemdah@olameter.com> and Taieb Lamine Ben Cheikh <tbencheikh@olameter.com> for their precious contribution to this work.
• Special thanks to Olameter for giving us the opportunity to work on this and for giving us permission to publish the results.