Model optimization techniques
By: Vishnu Nkumar
Tags: Machine Learning, Keras, TensorFlow, Model Optimization, Deep Learning

Tensorflow 2.3 has been released do please have a look in their website

When I used to look at my model sizes once training gets over I always gets frustrated because of the number it shows up. But no more worries as the machine learning world has developed immensely over the last decade. In this article, I will brief you through the methods of model optimization using TensorFlow.

Tensorflow is an open-sourced deep learning framework created by Google and it should be in the list of every machine learning researchers and aspirants. Arguably every new model is based out of either Tensorflow or Pytorch in this decade. Initially, there were other players in the market like MXNet, Caffe etc but now their share is very meagre when compared to Tensorflow and Pytorch. With weight clustering added to the already present pruning and quantization, there are three major ways by which one can reduce the size of their model without much reduction in metrics.

PRUNING:

Pruning is essentially nullifying some of the weights of a neural network. So what happens here is during training, some of the neurons are not used up for training and only the important neurons are trained. The same applies while inferencing the nullified neurons no longer take part in. Pruning helps in significant compression of the model i.e it can reduce up to 6x size and with a 0.1 percentage loss in metrics. Below is a way of how you can implement pruning in a Keras model after you compile and train the model in a conventional way.

import tensorflow_model_optimization as tfmot prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude # Compute end step to finish pruning after 2 epochs.

batch_size = 128

epochs = 2

validation_split = 0.1 # 10% of training set will be used for validation set. num_images = train_images.shape[0] * (1 - validation_split)

end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs # Define model for pruning.

pruning_params = {

'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.50,

final_sparsity=0.80,

begin_step=0,

end_step=end_step)

} model_for_pruning = prune_low_magnitude(model, **pruning_params) # `prune_low_magnitude` requires a recompile.

model_for_pruning.compile(optimizer='adam',

loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),

metrics=['accuracy']) model_for_pruning.summary()

QUANTIZATION:

In this method, we won’t be nullifying any neurons but we will be converting the model weights from floating-point numbers to integers. There are mainly two methods of quantization in practice. One is quantized aware training and the other one is post quantization. Quantized aware training involves the conversion of weights to integer during training while post-quantization is the conversion of weights to integer while saving the model. Below are the ways by which you can implement both the methods in your Keras model

Quantized aware training

import tensorflow_model_optimization as tfmot



quantize_model = tfmot.quantization.keras.quantize_model



# q_aware stands for for quantization aware.

q_aware_model = quantize_model(model)



# `quantize_model` requires a recompile.

q_aware_model.compile(optimizer='adam',

loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),

metrics=['accuracy'])



q_aware_model.summary()

Post-quantization

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)

converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]

tflite_quant_model = converter.convert()

WEIGHT CLUSTERING

This is the method by which I am more fascinated about because of the way its implemented. Weights in a single layer are clustered into N clusters and their respective centroid values are found. These weights are then replaced by the centroid values which will significantly help in model compression. The metrics obtained are better in comparison to the previous methods. The model metrics and size depends on the N value, if it is too low we will get a highly compressed model but the metrics will be very poor and vice versa. Below is the way how it is implemented in a Keras model

``import tensorflow_model_optimization as tfmot



cluster_weights = tfmot.clustering.keras.cluster_weights

CentroidInitialization = tfmot.clustering.keras.CentroidInitialization



clustering_params = {

'number_of_clusters': 16,

'cluster_centroids_init': CentroidInitialization.LINEAR

}



# Cluster a whole model

clustered_model = cluster_weights(model, **clustering_params)



# Use smaller learning rate for fine-tuning clustered model

opt = tf.keras.optimizers.Adam(learning_rate=1e-5)



clustered_model.compile(

loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),

optimizer=opt,

metrics=['accuracy'])



clustered_model.summary()

So this is my first shot at explaining something from the field I have been into. I hope you find this useful. Thanks for having a read at it. Your comments are more welcome than ever and I appreciate your time and effort

Please continue looking into my account, I will post new articles frequently