#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2019 The TensorFlow Authors.
End of explanation
"""


import logging
logging.getLogger("tensorflow").setLevel(logging.DEBUG)

import tensorflow as tf
from tensorflow import keras
import numpy as np
import pathlib

tf.float16

"""
Explanation: 训练后 float16 量化
<table class="tfo-notebook-buttons" align="left">
  <td><a target="_blank" href="https://tensorflow.google.cn/lite/performance/post_training_float16_quant"><img src="https://tensorflow.google.cn/images/tf_logo_32px.png">在  TensorFlow.org 上查看</a></td>
  <td><a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/lite/performance/post_training_float16_quant.ipynb"><img src="https://tensorflow.google.cn/images/colab_logo_32px.png">在 Google Colab 中运行 </a></td>
  <td><a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/lite/performance/post_training_float16_quant.ipynb">     <img src="https://tensorflow.google.cn/images/GitHub-Mark-32px.png">     在 GitHub 上查看源代码</a></td>
  <td><a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/lite/performance/post_training_float16_quant.ipynb"><img src="https://tensorflow.google.cn/images/download_logo_32px.png">下载笔记本</a></td>
</table>

概述
现在，TensorFlow Lite 支持在模型从 TensorFlow 转换到 TensorFlow Lite FlatBuffer 格式期间将权重转换为 16 位浮点值。这样可以将模型的大小缩减至原来的二分之一。某些硬件（如 GPU）可以在这种精度降低的算术中以原生方式计算，从而实现比传统浮点执行更快的速度。可以将 Tensorflow Lite GPU 委托配置为以这种方式运行。但是，转换为 float16 权重的模型仍可在 CPU 上运行而无需其他修改：float16 权重会在首次推断前上采样为 float32。这样可以在对延迟和准确率造成最小影响的情况下显著缩减模型大小。
在本教程中，您将从头开始训练一个 MNIST 模型，并在 TensorFlow 中检查其准确率，然后使用 float16 量化将此模型转换为 Tensorflow Lite FlatBuffer 格式。最后，检查转换后模型的准确率，并将其与原始 float32 模型进行比较。
构建 MNIST 模型
设置
End of explanation
"""


# Load MNIST dataset
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Normalize the input image so that each pixel value is between 0 to 1.
train_images = train_images / 255.0
test_images = test_images / 255.0

# Define the model architecture
model = keras.Sequential([
  keras.layers.InputLayer(input_shape=(28, 28)),
  keras.layers.Reshape(target_shape=(28, 28, 1)),
  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation=tf.nn.relu),
  keras.layers.MaxPooling2D(pool_size=(2, 2)),
  keras.layers.Flatten(),
  keras.layers.Dense(10)
])

# Train the digit classification model
model.compile(optimizer='adam',
              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
model.fit(
  train_images,
  train_labels,
  epochs=1,
  validation_data=(test_images, test_labels)
)

"""
Explanation: 训练并导出模型
End of explanation
"""


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

"""
Explanation: 在此示例中，您只对模型进行了一个周期的训练，因此只训练到约 96% 的准确率。
转换为 TensorFlow Lite 模型
现在，您可以使用 Python TFLiteConverter 将训练的模型转换为 TensorFlow Lite 模型。
现在使用 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: 将其写入 .tflite 文件：
End of explanation
"""


converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]

"""
Explanation: 要改为在导出时将模型量化为 float16，首先将 optimizations 标记设置为使用默认优化。然后将 float16 指定为目标平台支持的类型：
End of explanation
"""


tflite_fp16_model = converter.convert()
tflite_model_fp16_file = tflite_models_dir/"mnist_model_quant_f16.tflite"
tflite_model_fp16_file.write_bytes(tflite_fp16_model)

"""
Explanation: 最后，像往常一样转换模型。请注意，为了方便调用，转换后的模型默认仍将使用浮点输入和输出。
End of explanation
"""


!ls -lh {tflite_models_dir}

"""
Explanation: 请注意，生成文件的大小约为 1/2。
End of explanation
"""


interpreter = tf.lite.Interpreter(model_path=str(tflite_model_file))
interpreter.allocate_tensors()

interpreter_fp16 = tf.lite.Interpreter(model_path=str(tflite_model_fp16_file))
interpreter_fp16.allocate_tensors()

"""
Explanation: 运行 TensorFlow Lite 模型
使用 Python TensorFlow Lite 解释器运行 TensorFlow Lite 模型。
将模型加载到解释器中
End of explanation
"""


test_image = np.expand_dims(test_images[0], axis=0).astype(np.float32)

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

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

import matplotlib.pylab as plt

plt.imshow(test_images[0])
template = "True:{true}, predicted:{predict}"
_ = plt.title(template.format(true= str(test_labels[0]),
                              predict=str(np.argmax(predictions[0]))))
plt.grid(False)

test_image = np.expand_dims(test_images[0], axis=0).astype(np.float32)

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

interpreter_fp16.set_tensor(input_index, test_image)
interpreter_fp16.invoke()
predictions = interpreter_fp16.get_tensor(output_index)

plt.imshow(test_images[0])
template = "True:{true}, predicted:{predict}"
_ = plt.title(template.format(true= str(test_labels[0]),
                              predict=str(np.argmax(predictions[0]))))
plt.grid(False)

"""
Explanation: 在单个图像上测试模型
End of explanation
"""


# A helper function to evaluate the TF Lite model using "test" dataset.
def evaluate_model(interpreter):
  input_index = interpreter.get_input_details()[0]["index"]
  output_index = interpreter.get_output_details()[0]["index"]

  # Run predictions on every image in the "test" dataset.
  prediction_digits = []
  for test_image in test_images:
    # Pre-processing: add batch dimension and convert to float32 to match with
    # the model's input data format.
    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)
    interpreter.set_tensor(input_index, test_image)

    # Run inference.
    interpreter.invoke()

    # Post-processing: remove batch dimension and find the digit with highest
    # probability.
    output = interpreter.tensor(output_index)
    digit = np.argmax(output()[0])
    prediction_digits.append(digit)

  # Compare prediction results with ground truth labels to calculate accuracy.
  accurate_count = 0
  for index in range(len(prediction_digits)):
    if prediction_digits[index] == test_labels[index]:
      accurate_count += 1
  accuracy = accurate_count * 1.0 / len(prediction_digits)

  return accuracy

print(evaluate_model(interpreter))

"""
Explanation: 评估模型
End of explanation
"""


# NOTE: Colab runs on server CPUs. At the time of writing this, TensorFlow Lite
# doesn't have super optimized server CPU kernels. For this reason this may be
# slower than the above float interpreter. But for mobile CPUs, considerable
# speedup can be observed.
print(evaluate_model(interpreter_fp16))

"""
Explanation: 在 float16 量化模型上重复评估，以获得如下结果：
End of explanation
"""


