Converting Sandberg’s FaceNet pre-trained model to TensorFlow Lite (using an “unorthodox way”)

I converted some pre-trained FaceNet models to TensorFlow Lite to see how it works on my smartphone. I thought it was going to be an easy task, but I ran into several difficulties. Since this task took me a while, I thought it would be useful to explain how I did it, since this information could save a lot of time for those trying to do it.

The FaceNet

FaceNet: A Unified Embedding for Face Recognition and Clustering. FaceNet is a face recognition system developed in 2015 by researchers at Google that achieved the state-of-the-art results on a range of face recognition benchmark datasets (99.63% on the LFW). This work introduces the novel concept of triplet loss.

Hiroki Taniai’s Keras Model

In this great article, Jason Brownlee describes how to develop a Face Recognition System Using FaceNet in Keras. He uses this pre-trained Keras FaceNet model provided by Hiroki Taniai, trained on the MS-Celeb-1M dataset.

Hiroki’s Keras model size is about 92 MB, that’s too much for a smartphone, but what perhaps if we convert the model to TensorFlow Lite…

Let’s go to Python and convert it with the TocoConverter using post-training quantizantion technique to shrink the model size:

That was easy… Let’s see the size of the generated file:

Figure 1. TensorFlow Lite model light weight after post-training quantization

The model file is about 23 MB, that’s much better.

Let’s test it with Python using some random data:

EMBEEDINGS shape: (1, 128) norm 4.3764944
Process finished with exit code 0

So, the model gives us as output a 128 D vector but is not normalized.

• So, is that a problem?

Well, actually we could use it to compare faces, but because it is not normalized it is not clear what threshold to use to decide if two embeedings correspond to the same person. According to FaceNet work this value should be around 1.24. See the following paper extract:

Extract from FaceNet recommended threshold for face classification

Let’s see which other options are there available…

Converting David Sandberg’s Implementation to TFLite

Today the most important implementation of FaceNet is without a doubt the David Sandberg’s implementation, today it counts with about 30 contributors in github. This implementation is very well done and has become a facto standard for the FaceNet. It provides several pre-traned models over the Inception ResNet v1. This implementation is on TensorFlow.

• So, easy, why don’t we just convert this implementation from TensorFlow to TensorFlow Lite?

Not so fast, this is not an easy task. Even though converting the FaceNet model from Keras to TensorFlow Lite is barely one line of code, converting from TensorFlow to TensorFlow Lite requires five steps: — First we have to strip the training branch, then we have to create a new inference graph by setting the parameter is_training=False in the inception_resnet_v1 inference layers creation, and save the graph to disk, and then we need to freeze the graph to combine the structure of the net and its weights, and then we can finally run the “tflite_convert” command with its 11 parameters to get our brand new TensorFlow Lite model. — Piece of cake!

This post explains in detail how to do all this stuff. I followed its instructions, and got my 23 MB shrunk .tflite file. But! the data-type of the new model’s input was uint8, and the resulting embeedings were wrong. — Did I do something wrong on my way? The last part of the post shows some code using dummy float data, but would be great to add a sanity check to see if the embeedings are fine after conversion.

TF -> Keras -> TF lite approach

So, let’s try something different: Why don’t we convert the model from TensorFlow to Keras, and then from Keras to TensorFlow Lite? Figure 2 depicts the idea:

Figure 2. “Unorthodox way” to convert a TensorFlow model to TensorFlow Lite

• Ok, but how easy is to convert a TensorFlow trained model to a Keras trained model?

Actually, this task is even more difficult. To do such conversion it is necessary to have exactly the same model definition in both platforms, (TensorFlow and Keras). If we have that, we can create an instance of the model in Keras initialized with random weights and copy the weights from one model to the other.

Fortunately Hiroki Taniai kindly did this open-heart surgery job for us here. He basically mapped the layers from Sandberg’s model to his model using regular expressions. Then he exported to numpy arrays the weights stored in the checkpoint file, and then he wrote the weights into the Keras model using the set_weights() method.

There are some minor issues:

1. He did the job for an old version of Sandberg’s pretrained model. Sandberg’s current model has been trained with a final Dense layer of 512 units (instead of 128), and thus the conversion fails.
2. The Keras model does not have normalize its embeedings

But this two issues can be easily solved:

• The first issue is simply solved by creating the Keras model with classes=512, instead of the default 128.
• The second issue gets fixed simply by adding a normalization layer at the end of the classification block.

Figure 3. Normalization Layer Adding a final Lambda layer will give us l2-normalized output embeedings.

Here I provide a Google Colab, that does the job.

What it does is as follows:

1. Download Hiroki Taniai’s Keras FaceNet implementation
2. Override the inception_resnet_v1.py file with my patched version (which does adds an extra layer to the model to have normalized embeedings as output)
3. Download Sandberg’s pre-trained model (20180402–114759) from here, and unzips it
4. Extract the tensors from the checkpoint file and writes the weights to numpy arrays on disk, mapping the name of each corresponding layer.
5. Create a new Keras model with random weights (Important: using 512 classes).
6. Write the weights for each corresponding layer reading from the numpy arrays.
7. Store the model with the Keras format .h5
8. Convert Keras to TensorFlow Lite using the command “tflite_convert”.

After running all the Colab steps, we should find the converted model file under the path /contents/keras-facenet/model. And simply download it from the contextual menu.

Sanity check

After the conversion was done, I checked that the model didn’t break.

First, I loaded this 3 different images provided in the repository (2 from Bill Gates, and 1 from Larry Page). The faces are cropped using the mtcnn python package.

Now check if embeedings are good, so let’s compute it’s embeedings:

distance bill vs bill 0.7266881 
distance bill vs larry 1.2134411

So by using the euclidean distance, the embeedings seems to be good. Those who want can also check it with the cosine distance.