{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c8Cx-rUMVX25"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "I9sUhVL_VZNO"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Y8E0lw5eYWm"
      },
      "source": [
        "# 训练后 float16 量化"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CGuqeuPSVNo-"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <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>\n",
        "  <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>\n",
        "  <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>\n",
        "  <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>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BTC1rDAuei_1"
      },
      "source": [
        "## 概述\n",
        "\n",
        "现在，[TensorFlow Lite](https://tensorflow.google.cn/lite/) 支持在模型从 TensorFlow 转换到 TensorFlow Lite FlatBuffer 格式期间将权重转换为 16 位浮点值。这样可以将模型的大小缩减至原来的二分之一。某些硬件（如 GPU）可以在这种精度降低的算术中以原生方式计算，从而实现比传统浮点执行更快的速度。可以将 Tensorflow Lite GPU 委托配置为以这种方式运行。但是，转换为 float16 权重的模型仍可在 CPU 上运行而无需其他修改：float16 权重会在首次推断前上采样为 float32。这样可以在对延迟和准确率造成最小影响的情况下显著缩减模型大小。\n",
        "\n",
        "在本教程中，您将从头开始训练一个 MNIST 模型，并在 TensorFlow 中检查其准确率，然后使用 float16 量化将此模型转换为 Tensorflow Lite FlatBuffer 格式。最后，检查转换后模型的准确率，并将其与原始 float32 模型进行比较。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2XsEP17Zelz9"
      },
      "source": [
        "## 构建 MNIST 模型"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dDqqUIZjZjac"
      },
      "source": [
        "### 设置"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gyqAw1M9lyab"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "logging.getLogger(\"tensorflow\").setLevel(logging.DEBUG)\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "import numpy as np\n",
        "import pathlib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eQ6Q0qqKZogR"
      },
      "source": [
        "### 训练并导出模型"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hWSAjQWagIHl"
      },
      "outputs": [],
      "source": [
        "# Load MNIST dataset\n",
        "mnist = keras.datasets.mnist\n",
        "(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
        "\n",
        "# Normalize the input image so that each pixel value is between 0 to 1.\n",
        "train_images = train_images / 255.0\n",
        "test_images = test_images / 255.0\n",
        "\n",
        "# Define the model architecture\n",
        "model = keras.Sequential([\n",
        "  keras.layers.InputLayer(input_shape=(28, 28)),\n",
        "  keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
        "  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation=tf.nn.relu),\n",
        "  keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
        "  keras.layers.Flatten(),\n",
        "  keras.layers.Dense(10)\n",
        "])\n",
        "\n",
        "# Train the digit classification model\n",
        "model.compile(optimizer='adam',\n",
        "              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "              metrics=['accuracy'])\n",
        "model.fit(\n",
        "  train_images,\n",
        "  train_labels,\n",
        "  epochs=1,\n",
        "  validation_data=(test_images, test_labels)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5NMaNZQCkW9X"
      },
      "source": [
        "在此示例中，您只对模型进行了一个周期的训练，因此只训练到约 96% 的准确率。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xl8_fzVAZwOh"
      },
      "source": [
        "### 转换为 TensorFlow Lite 模型\n",
        "\n",
        "现在，您可以使用 TensorFlow Lite [Converter](https://tensorflow.google.cn/lite/models/convert) 将训练后的模型转换为 TensorFlow Lite 模型。\n",
        "\n",
        "现在使用 `TFLiteConverter` 加载模型："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_i8B2nDZmAgQ"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "tflite_model = converter.convert()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F2o2ZfF0aiCx"
      },
      "source": [
        "将其写入 `.tflite` 文件："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vptWZq2xnclo"
      },
      "outputs": [],
      "source": [
        "tflite_models_dir = pathlib.Path(\"/tmp/mnist_tflite_models/\")\n",
        "tflite_models_dir.mkdir(exist_ok=True, parents=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ie9pQaQrn5ue"
      },
      "outputs": [],
      "source": [
        "tflite_model_file = tflite_models_dir/\"mnist_model.tflite\"\n",
        "tflite_model_file.write_bytes(tflite_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7BONhYtYocQY"
      },
      "source": [
        "要改为在导出时将模型量化为 float16，首先将 `optimizations` 标记设置为使用默认优化。然后将 float16 指定为目标平台支持的类型："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HEZ6ET1AHAS3"
      },
      "outputs": [],
      "source": [
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.target_spec.supported_types = [tf.float16]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xW84iMYjHd9t"
      },
      "source": [
        "最后，像往常一样转换模型。请注意，为了方便调用，转换后的模型默认仍将使用浮点输入和输出。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yuNfl3CoHNK3"
      },
      "outputs": [],
      "source": [
        "tflite_fp16_model = converter.convert()\n",
        "tflite_model_fp16_file = tflite_models_dir/\"mnist_model_quant_f16.tflite\"\n",
        "tflite_model_fp16_file.write_bytes(tflite_fp16_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PhMmUTl4sbkz"
      },
      "source": [
        "请注意，生成文件的大小约为 `1/2`。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JExfcfLDscu4"
      },
      "outputs": [],
      "source": [
        "!ls -lh {tflite_models_dir}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L8lQHMp_asCq"
      },
      "source": [
        "## 运行 TensorFlow Lite 模型"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-5l6-ciItvX6"
      },
      "source": [
        "使用 Python TensorFlow Lite 解释器运行 TensorFlow Lite 模型。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ap_jE7QRvhPf"
      },
      "source": [
        "### 将模型加载到解释器中"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jn16Rc23zTss"
      },
      "outputs": [],
      "source": [
        "interpreter = tf.lite.Interpreter(model_path=str(tflite_model_file))\n",
        "interpreter.allocate_tensors()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J8Pztk1mvNVL"
      },
      "outputs": [],
      "source": [
        "interpreter_fp16 = tf.lite.Interpreter(model_path=str(tflite_model_fp16_file))\n",
        "interpreter_fp16.allocate_tensors()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2opUt_JTdyEu"
      },
      "source": [
        "### 在单个图像上测试模型"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AKslvo2kwWac"
      },
      "outputs": [],
      "source": [
        "test_image = np.expand_dims(test_images[0], axis=0).astype(np.float32)\n",
        "\n",
        "input_index = interpreter.get_input_details()[0][\"index\"]\n",
        "output_index = interpreter.get_output_details()[0][\"index\"]\n",
        "\n",
        "interpreter.set_tensor(input_index, test_image)\n",
        "interpreter.invoke()\n",
        "predictions = interpreter.get_tensor(output_index)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XZClM2vo3_bm"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pylab as plt\n",
        "\n",
        "plt.imshow(test_images[0])\n",
        "template = \"True:{true}, predicted:{predict}\"\n",
        "_ = plt.title(template.format(true= str(test_labels[0]),\n",
        "                              predict=str(np.argmax(predictions[0]))))\n",
        "plt.grid(False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3gwhv4lKbYZ4"
      },
      "outputs": [],
      "source": [
        "test_image = np.expand_dims(test_images[0], axis=0).astype(np.float32)\n",
        "\n",
        "input_index = interpreter_fp16.get_input_details()[0][\"index\"]\n",
        "output_index = interpreter_fp16.get_output_details()[0][\"index\"]\n",
        "\n",
        "interpreter_fp16.set_tensor(input_index, test_image)\n",
        "interpreter_fp16.invoke()\n",
        "predictions = interpreter_fp16.get_tensor(output_index)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CIH7G_MwbY2x"
      },
      "outputs": [],
      "source": [
        "plt.imshow(test_images[0])\n",
        "template = \"True:{true}, predicted:{predict}\"\n",
        "_ = plt.title(template.format(true= str(test_labels[0]),\n",
        "                              predict=str(np.argmax(predictions[0]))))\n",
        "plt.grid(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LwN7uIdCd8Gw"
      },
      "source": [
        "### 评估模型"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "05aeAuWjvjPx"
      },
      "outputs": [],
      "source": [
        "# A helper function to evaluate the TF Lite model using \"test\" dataset.\n",
        "def evaluate_model(interpreter):\n",
        "  input_index = interpreter.get_input_details()[0][\"index\"]\n",
        "  output_index = interpreter.get_output_details()[0][\"index\"]\n",
        "\n",
        "  # Run predictions on every image in the \"test\" dataset.\n",
        "  prediction_digits = []\n",
        "  for test_image in test_images:\n",
        "    # Pre-processing: add batch dimension and convert to float32 to match with\n",
        "    # the model's input data format.\n",
        "    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)\n",
        "    interpreter.set_tensor(input_index, test_image)\n",
        "\n",
        "    # Run inference.\n",
        "    interpreter.invoke()\n",
        "\n",
        "    # Post-processing: remove batch dimension and find the digit with highest\n",
        "    # probability.\n",
        "    output = interpreter.tensor(output_index)\n",
        "    digit = np.argmax(output()[0])\n",
        "    prediction_digits.append(digit)\n",
        "\n",
        "  # Compare prediction results with ground truth labels to calculate accuracy.\n",
        "  accurate_count = 0\n",
        "  for index in range(len(prediction_digits)):\n",
        "    if prediction_digits[index] == test_labels[index]:\n",
        "      accurate_count += 1\n",
        "  accuracy = accurate_count * 1.0 / len(prediction_digits)\n",
        "\n",
        "  return accuracy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T5mWkSbMcU5z"
      },
      "outputs": [],
      "source": [
        "print(evaluate_model(interpreter))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Km3cY9ry8ZlG"
      },
      "source": [
        "在 float16 量化模型上重复评估，以获得如下结果："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-9cnwiPp6EGm"
      },
      "outputs": [],
      "source": [
        "# NOTE: Colab runs on server CPUs. At the time of writing this, TensorFlow Lite\n",
        "# doesn't have super optimized server CPU kernels. For this reason this may be\n",
        "# slower than the above float interpreter. But for mobile CPUs, considerable\n",
        "# speedup can be observed.\n",
        "print(evaluate_model(interpreter_fp16))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L7lfxkor8pgv"
      },
      "source": [
        "在此示例中，您已将模型量化为 float16，但准确率没有任何差别。\n",
        "\n",
        "您还可以在 GPU 上评估 fp16 量化模型。要使用降低的精度值执行所有算术，请确保在您的应用中创建 `TfLiteGPUDelegateOptions` 结构，并将 `precision_loss_allowed` 设置为 `1`，如下所示：\n",
        "\n",
        "```\n",
        "//Prepare GPU delegate.\n",
        "const TfLiteGpuDelegateOptions options = {\n",
        "  .metadata = NULL,\n",
        "  .compile_options = {\n",
        "    .precision_loss_allowed = 1,  // FP16\n",
        "    .preferred_gl_object_type = TFLITE_GL_OBJECT_TYPE_FASTEST,\n",
        "    .dynamic_batch_enabled = 0,   // Not fully functional yet\n",
        "  },\n",
        "};\n",
        "```\n",
        "\n",
        "有关 TFLite GPU 委托以及如何在您的应用中进行使用的详细文档，请参阅[此处](https://tensorflow.google.cn/lite/performance/gpu_advanced?source=post_page---------------------------)。"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "post_training_float16_quant.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}