! pip uninstall -y tensorflow
! pip install -U tf-nightly

import tensorflow as tf
tf.enable_eager_execution()

! git clone --depth 1 https://github.com/tensorflow/models

import sys
import os

if sys.version_info.major >= 3:
    import pathlib
else:
    import pathlib2 as pathlib

# Add `models` to the python path.
models_path = os.path.join(os.getcwd(), "models")
sys.path.append(models_path)

"""
Explanation: Post Training Quantization
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/tutorials/post_training_quant.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/tutorials/post_training_quant.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
</table>

Overview
TensorFlow Lite now supports
converting weights to 8 bit precision as part of model conversion from
tensorflow graphdefs to TFLite's flat buffer format. Weight quantization
achieves a 4x reduction in the model size. In addition, TFLite supports on the
fly quantization and dequantization of activations to allow for:


Using quantized kernels for faster implementation when available.


Mixing of floating-point kernels with quantized kernels for different parts
    of the graph.


Note that the activations are always stored in floating point. For ops that
support quantized kernels, the activations are quantized to 8 bits of precision
dynamically prior to processing and are de-quantized to float precision after
processing. Depending on the model being converted, this can give a speedup over
pure floating point computation.
In contrast to
quantization aware training
, the weights are quantized post training and the activations are quantized dynamically 
at inference in this method.
Therefore, the model weights are not retrained to compensate for quantization
induced errors. It is important to check the accuracy of the quantized model to
ensure that the degradation is acceptable.
In this tutorial, we train an MNIST model from scratch, check its accuracy in
tensorflow and then convert the saved model into a Tensorflow Lite flatbuffer
with weight quantization. We finally check the
accuracy of the converted model and compare it to the original saved model. We
run the training script mnist.py from
Tensorflow official mnist tutorial.
Building an MNIST model
Setup
End of explanation
"""


saved_models_root = "/tmp/mnist_saved_model"

# The above path addition is not visible to subprocesses, add the path for the subprocess as well.
# Note: channels_last is required here or the conversion may fail. 
!PYTHONPATH={models_path} python models/official/mnist/mnist.py --train_epochs=1 --export_dir {saved_models_root} --data_format=channels_last

"""
Explanation: Train and export the model
End of explanation
"""


saved_model_dir = str(sorted(pathlib.Path(saved_models_root).glob("*"))[-1])
saved_model_dir

"""
Explanation: For the example, we only trained the model for a single epoch, so it only trains to ~96% accuracy.
Convert to a TFLite model
The savedmodel directory is named with a timestamp. Select the most recent one:
End of explanation
"""


import tensorflow as tf
tf.enable_eager_execution()
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

"""
Explanation: Using the python TFLiteConverter, the saved model can be converted into a TFLite model.
First load the model using the TFLiteConverter:
End of explanation
"""


tflite_models_dir = pathlib.Path("/tmp/mnist_tflite_models/")
tflite_models_dir.mkdir(exist_ok=True, parents=True)

tflite_model_file = tflite_models_dir/"mnist_model.tflite"
tflite_model_file.write_bytes(tflite_model)

"""
Explanation: Write it out to a tflite file:
End of explanation
"""


# Note: If you don't have a recent tf-nightly installed, the
# "post_training_quantize" line will have no effect.
tf.logging.set_verbosity(tf.logging.INFO)
converter.post_training_quantize = True
tflite_quant_model = converter.convert()
tflite_model_quant_file = tflite_models_dir/"mnist_model_quant.tflite"
tflite_model_quant_file.write_bytes(tflite_quant_model)

"""
Explanation: To quantize the model on export, set the post_training_quantize flag:
End of explanation
"""


!ls -lh {tflite_models_dir}

"""
Explanation: Note how the resulting file, with post_training_quantize set, is approximately 1/4 the size.
End of explanation
"""


import numpy as np
mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()
images, labels = tf.to_float(mnist_test[0])/255.0, mnist_test[1]

# Note: If you change the batch size, then use 
# `tf.lite.Interpreter.resize_tensor_input` to also change it for
# the interpreter.
mnist_ds = tf.data.Dataset.from_tensor_slices((images, labels)).batch(1)

"""
Explanation: Run the TFLite models
We can run the TensorFlow Lite model using the python TensorFlow Lite
Interpreter. 
load the test data
First let's load the mnist test data to feed to it:
End of explanation
"""


interpreter = tf.lite.Interpreter(model_path=str(tflite_model_file))
interpreter.allocate_tensors()
input_index = interpreter.get_input_details()[0]["index"]
output_index = interpreter.get_output_details()[0]["index"]

tf.logging.set_verbosity(tf.logging.DEBUG)
interpreter_quant = tf.lite.Interpreter(model_path=str(tflite_model_quant_file))

interpreter_quant.allocate_tensors()
input_index = interpreter_quant.get_input_details()[0]["index"]
output_index = interpreter_quant.get_output_details()[0]["index"]


"""
Explanation: Load the model into an interpreter
End of explanation
"""


for img, label in mnist_ds.take(1):
  break

interpreter.set_tensor(input_index, img)
interpreter.invoke()
predictions = interpreter.get_tensor(output_index)

import matplotlib.pylab as plt

plt.imshow(img[0])
template = "True:{true}, predicted:{predict}"
_ = plt.title(template.format(true= str(label[0].numpy()),
                              predict=str(predictions[0,0])))
plt.grid(False)

"""
Explanation: Test the model on one image
End of explanation
"""


def eval_model(interpreter, mnist_ds):
  total_seen = 0
  num_correct = 0

  for img, label in mnist_ds:
    total_seen += 1
    interpreter.set_tensor(input_index, img)
    interpreter.invoke()
    predictions = interpreter.get_tensor(output_index)
    if predictions == label.numpy():
      num_correct += 1

    if total_seen % 500 == 0:
        print("Accuracy after %i images: %f" %
              (total_seen, float(num_correct) / float(total_seen)))

  return float(num_correct) / float(total_seen)

print(eval_model(interpreter, mnist_ds))

"""
Explanation: Evaluate the models
End of explanation
"""


print(eval_model(interpreter_quant, mnist_ds))


"""
Explanation: We can repeat the evaluation on the weight quantized model to obtain:
End of explanation
"""


archive_path = tf.keras.utils.get_file("resnet_v2_101.tgz", "https://storage.googleapis.com/download.tensorflow.org/models/tflite_11_05_08/resnet_v2_101.tgz", extract=True)
archive_path = pathlib.Path(archive_path)
archive_dir = str(archive_path.parent)

"""
Explanation: In this example, we have compressed model with no difference in the accuracy.
Optimizing an existing model
We now consider another example. Resnets with pre-activation layers (Resnet-v2) are widely used for vision applications.
  Pre-trained frozen graph for resnet-v2-101 is available at the
  Tensorflow Lite model repository.
We can convert the frozen graph to a TFLite flatbuffer with quantization by:
End of explanation
"""


! cat {archive_dir}/resnet_v2_101_299_info.txt

graph_def_file = pathlib.Path(archive_path).parent/"resnet_v2_101_299_frozen.pb"
input_arrays = ["input"] 
output_arrays = ["output"]
converter = tf.lite.TFLiteConverter.from_frozen_graph(
  str(graph_def_file), input_arrays, output_arrays, input_shapes={"input":[1,299,299,3]})
converter.post_training_quantize = True
resnet_tflite_file = graph_def_file.parent/"resnet_v2_101_quantized.tflite"
resnet_tflite_file.write_bytes(converter.convert())



!ls -lh {archive_dir}/*.tflite

"""
Explanation: The info.txt file lists the input and output names. You can also find them using TensorBoard to visually inspect the graph.
End of explanation
"""


