<filename>Learning_Tensorflow/tf_Iite/post_training_quantization.py
# -*- coding: utf-8 -*-
"""Post_Training_Quantization.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/13p3RtIv5kG0yu8cjDwbS29WJNdAVf1vJ

# Post training Quatization

- We have a trained model that we want to optimize.
- Optimizing this model by using tflite does weight pruning, adjusts the outputs and weights to int8.
- Also tflite optimizes the size of the model in this process.
- This is useful for low compute devices like mobiles and also useful when we are using edge devices.

# Load the libraries
"""

import tensorflow as tf
print(tf.__version__)

import matplotlib.pyplot as plt
import numpy as np

"""# Load the old Model which was trained Unquantized"""

unqunatized_model = tf.keras.models.load_model('/content/fashion_unquantized_model.hdf5')

unqunatized_model.summary()

"""# Use Tflite to convert the Model"""

converter = tf.lite.TFLiteConverter.from_keras_model(unqunatized_model)
tflite_model = converter.convert()

model_size = len(tflite_model) / 1024
print("Model size = %d KBs"%(model_size))

"""# We can set several optimizations also"""

converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_LATENCY]
# DEFAULT IS FLOAT 16 format for quantization.
# We need a dataset_gen once to check for data.
# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# converter.inference_input_type = tf.uint8
# converter.inference_output_type = tf.uint8
# converter.representative_dataset = dataset_gen
tflite_quantized_model = converter.convert()

model_size = len(tflite_quantized_model) / 1024
print("Model size = %d KBs"%(model_size))

converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# converter.inference_input_type = tf.uint8
# converter.inference_output_type = tf.uint8
# converter.representative_dataset = dataset_gen
tflite_tiny_quantized_model = converter.convert()

model_size = len(tflite_tiny_quantized_model) / 1024
print("Model size = %d KBs"%(model_size))

"""# Check on the dataset once"""

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

print(x_test.shape)

print(x_train.shape)

x_test = np.expand_dims(x_test, axis=-1)

"""# Invoke the interpreter"""

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

interpreter.get_tensor_details()

all_layers = interpreter.get_tensor_details()
for layer in all_layers:
    print(interpreter.get_tensor(layer["index"]))

input_tensor_index = interpreter.get_input_details()[0]["index"]
output = interpreter.tensor(interpreter.get_output_details()[0]["index"])

prediction_output = []
for test_image in x_test:
    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)
    interpreter.set_tensor(input_tensor_index, test_image)
    interpreter.invoke()
    out = np.argmax(output()[0])

    prediction_output.append(out)

accurate_count = 0
for index in range(len(prediction_output)):
    if prediction_output[index] == y_test[index]:
        accurate_count += 1
    
    accuracy = accurate_count / len(prediction_output)

print(accuracy)

"""# Time It"""

import time

start_time = time.time()

for test_image in x_test:
    test_image = np.expand_dims(test_image, axis=0)
    model.predict(test_image)

end_time = time.time()
print("Total time for prediction with normal model = %d seconds" %(end_time - start_time))

start_time = time.time()

for test_image in x_test:
    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)

    interpreter.set_tensor(input_tensor_index, test_image)
    interpreter.invoke()
    out = np.argmax(output()[0])

end_time = time.time()
print("Total time for prediction with normal model = %d seconds" %(end_time - start_time))

"""# Gist of what you should do"""

# Make tflite model
converter = tf.lite.TFLiteConverter.from_keras_model(unqunatized_model)
tflite_model = converter.convert()

# Set the interpreter
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

# Set i/o for the interpreter
input_tensor_index = interpreter.get_input_details()[0]["index"]
output = interpreter.tensor(interpreter.get_output_details()[0]["index"])

# Get the output
for test_image in x_test:
    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)

    interpreter.set_tensor(input_tensor_index, test_image)
    interpreter.invoke()
    out = np.argmax(output()[0])