Model Compression with TensorFlow Lite: A Look into Reducing Model Size | by Cawin Chan | Towards Data Science

A significant problem in the arms race to produce more accurate models is complexity, which leads to the problem of size. These models are usually huge and resource-intensive, which leads to greater space and time consumption. (Takes up more space in memory and slower in prediction as compared to smaller models)
A large model size is a common byproduct when attempting to push the limits of model accuracy in predicting unseen data in deep learning applications. For example, with more nodes, we can detect subtler features in the dataset. However, for project requirements such as using AI in embedded systems that depend on fast predictions, we are limited by the available computational resources. Furthermore, prevailing edge devices do not have networking capabilities, as such, we are not able to utilize cloud computing. This results in the inability to use massive models which would take too long to get meaningful predictions.
As such, we will need to optimize our performance to size, when designing our model.
To overly simplify for the gist of understanding machine learning models, a neural network is a set of nodes with weights(W) that connect between nodes. You can think of this as a set of instructions that we optimize to increase our likelihood of generating our desired class. The more specific this set of instructions are, the greater our model size, which is dependent on the size of our parameters (our configuration variables such as weight).
We are to classify between 3 different classes from a given dataset in a Raspberry Pi(RPi). RPi is traditionally not an embedded device, however, in our case RPi was a step towards embedded devices.
A bit hard to read? Here is a simpler view.
Here we have the average accuracy of prediction for all classes.
It seems like on average the winner is the Transformer model as it has higher average accuracy, however, if we compare the performance to the model size. We see that the transformer has a higher total parameter count, then the Encoder-Decoder LSTM.
For the Transformer, a marginal increase in accuracy of 0.1% required an increase of 1.39Mb of parameters, thus it had a lower accuracy over size ratio.
As such the Encoder-Decoder model is the clear winner, however, 3.16MB is still too large of a model to allow for the effective usage of an embedded device.
A re-visit from my previous article, installing TensorFlow 2.3.0 in Raspberry Pi3+/4,
With our autoencoder, on average it takes 1760ms to predict in our embedded device, however as our sensors are sensing every second, this won’t do.
Great, now what?
These are the model compression methods, each with increasing degrees of difficulty in implementation. An excellent and comprehensive survey on each of these techniques has been done here.
Quantization
Pruning
Low Ranked Approximation
Knowledge Distillation
TensorFlow Lite deals with the first two methods and does a great job in abstracting the hard parts of model compression.
TensorFlow Lite covers:
Post-Training Quantization— Reduce Float16 — Hybrid Quantization — Integer Quantization
Post-Training Quantization— Reduce Float16 — Hybrid Quantization — Integer Quantization
2. During-Training Quantization
3. Post-Training Pruning
4. Post-Training Clustering
The most common and easiest to implement method would be post-training quantization. The usage of quantization is the limiting of the bits of precision of our model parameters as such this reduces the amount of data that is needed to be stored.
I will only go through post-training Hybrid/Dynamic range quantization because it is the easiest to implement, has a great amount of impact in size reduction with minimal loss.
To reference our earlier neural network diagram, our model parameters(weights) which refer to the lines connecting each note can be seen to represent its literal weight(significance) or the importance of the node to predict our desired outcome.
Originally, we gave 32-bits to each weight, known as the tf.float32(32-bit single-precision floating-point), to reduce the size of our model, we would essentially shave off from 32-bits to 16-bits(tf.float16) or 8-bits( tf.int8) depending on the type of quantization used.
We can intuitively see that this poses significant exponential size reductions as with a bigger and more complex the model, the greater the number of nodes and subsequently the greater number of weights which leads to a more significant size reduction especially for fully-connected neural networks, which has each layer of nodes connected to each of the nodes in the next layer.
Smaller model sizes — Models can actually be stored into embedded devices (ESP32 has ~4Mb of Flash Memory)
Faster Prediction Rates — This speeds up actionability, which provides viability for real-time decisions.
Lowered power consumption — An often overlooked feature as environments that train models often come with a constant supply of power, embedded devices usually run on a battery thus this is the most important feature.
You might be quick to think that reducing the amount of information we store for each weight, would always be detrimental to our model, however, quantization promotes generalization which was a huge plus in preventing overfitting — a common problem with complex models.
By Jerome Friedman, the father of gradient boost, empirical evidence shows that lots of small steps in the right direction result in better predictions with test data. By quantization, it is possible to get an improved accuracy due to the decreased sensitivity of the weights.
Imagine if, in our dataset, we get lucky, every time we try to detect a cookie our dataset shows us a chocolate chip cookie, our cookie detection would get a high training accuracy, however, if in real-life we only have raisin cookies, it would have a low test accuracy. Generalization is like blurring our chocolate chip so that our model realizes as long as there is this blob, it is a cookie.
The same can be said for other compression methods such as pruning. It is in this vein whereby dropout also can improve unseen accuracy as randomly dropping nodes during training promotes generalization as well.
I will be using Tensorflow 2.3.0.
To convert your model post-training to the quantized model:
First, save your tf.keras model in the recommended SavedModel format (tf.lite.TFLiteConverter.from_keras_model() works too however I was met with many errors, be warned):
import tensorflow as tftf.keras.models.save_model(model, str(saved_model_dir))
Second, ensure that the operations(ops) are supported by TensorFlow Lite, check out which operations are supported here.
As of 06/09/20, the activation function of SELU is not supported by TensorFlow(TF), I found that out the hard way, in my experience, RELU is a good substitute with minimal loss.
Since the TensorFlow Lite builtin operator library only supports a limited number of TensorFlow operators, not every model is convertible.
If your model requires TensorFlow operators that are not supported, not all is lost, converter.allow_custom_ops = True & converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]#, tf.lite.OpsSet.SELECT_TF_OPS , should be attempted to allow for a custom implementations.ALLOW_CUSTOM_OPS — Allow for the custom implementation of unsupported operators.TFLITE_BUILTINS — Transforms the model using TensorFlow Lite built-in operators.SELECT_TF_OPS — Converts the model using the TensorFlow operator.
I had an autoencoder model with 2 LSTMs, using allow_custom_ops = True & tf.lite.OpsSet.TFLITE_BUILTINS without my own custom implementations worked for me. This should skip any operators that are not supported.
Third, we will convert the file.
import pathlibconverter=tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)tflite_models_dir = pathlib.Path.cwd() / "tflite_model_dir"tflite_models_dir.mkdir(exist_ok=True, parents=True)converter.allow_custom_ops = Trueconverter.optimizations = [tf.lite.Optimize.DEFAULT]converter.target_spec.supported_ops=[tf.lite.OpsSet.TFLITE_BUILTINS]#, tf.lite.OpsSet.SELECT_TF_OPS tflite_model_quant = converter.convert()
Note: pathlib is an optional library used to simplify windows pathing.
Lastly, instead of predicting using our quantized model, we will run an inference.
The term inference refers to the process of executing a TensorFlow Lite model on-device in order to make predictions based on input data.
To perform an inference with a TensorFlow Lite model, we must run it through an interpreter. The TensorFlow Lite interpreter is designed to be lean and fast. The interpreter uses a static graph ordering and a custom (less-dynamic) memory allocator to ensure minimal load, initialization, and execution latency.
# Load the TFLite model and allocate tensors.    interpreter = tf.lite.Interpreter(model_path=path_to_tflite_model)    interpreter.allocate_tensors()# Get input and output tensors.    input_details = interpreter.get_input_details()    output_details = interpreter.get_output_details()#Predict model with processed data             input_shape = input_details[0]['shape']            input_data = np.array(x_3d_pca_processed, dtype=np.float32)#x_3d_pca_processed is the input, change to your inputinterpreter.set_tensor(input_details[0]['index'], input_data)            interpreter.invoke()            prediction = interpreter.get_tensor(output_details[0]['index'])  outcome = np.argmax(prediction.mean(axis=1),axis=1)#outcome is my own implementation to get the class with the highest probability, remove if necessary
That’s it!
For a more in-depth explanation of these TFlite tools, click here. For more on TFlite tools, click here.
With our original model as the baseline, we compared with:- Original Model- Post-Quantized model- Post-Pruned model- Post-Quantized and Post-Pruned model on our computer before implementing our model in our “embedded” device.
It works surprisingly well! We have a clear winner with post-quantization having a negligible loss in accuracy and its 81% reduction in size and 91% reduction in latency.
We will now proceed to integrate our TFlite model into our RPi. Our RPi will have Tensorflow 2.3.0 installed, it has to be running on Debian buster, more details in my previous article. Once you dragged the TFLite file to our RPi, run the inference.
In our RPi, I achieved a prediction speed of around 90ms for our 0.60Mb TFlite model vs our 1760ms with 3.16Mb original! It took me by surprise how great of a performance improvement TFlite was able to churn despite my “custom implementations” and how well an RPI could handle a TF model.
As seen from above, TensorFlow’s estimations match our findings, similar to Resnet V2, we experienced a 0% accuracy loss and even gained a 5.2x (3.16mb to 0.60mb) smaller model which exceeded expectations.
Do take into consideration that in Deep Learning which extends to model compression, there is no hard and fast solution to any problem. the fact of the matter is that it is hard to tell how much savings we would get, the best we can do is try it out ourselves and analyze whether there is an improvement in model size with little loss in accuracy.
With that being said, I believe that attempting post-quantization is an excellent first step towards model compression, due to its ease in implementation, significant reduction, and negligible loss. Only when nothing else works, should you attempt the more complex methods of model compression.
Pros:- Easiest and only tool (06/09/20) to implement model compression- Minimal effect on accuracy (Depending on model)- Major speed up in prediction
Cons: - Requires the latest Tensorflow version 2 - Relatively new (06/09/20), Many operations (ops) are not supported yet such as SELU- Requires converting model which can fail- Possible complications when running inference compared to our good friend.predict() as it is more convoluted.
Despite its cons, TensorFlow Lite serves as a powerful tool with great potential that surpassed my expectations. I foresee in the near future, model compression being more widely used as the demand for AI in embedded devices inevitably grows, which gives TFLite a reason to provide greater operation coverage. With its shortcomings that can be mitigated by custom implementations, TensorFlow Lite for model compression is worth a shot.
With that being said, model compression should not be seen as a one-trick pony, instead, it should be used after we have attempted to optimize the performance to the model size and are unable to reduce the model size, without significant accuracy loss. Exploring different possible models and locating a better model architecture is often a better solution, I explored over 5 different model architectures before choosing our Autoencoder. Get a storage container instead of trying to cram a fridge into a suitcase.
To support my work: