【ATU Book-i.MX8系列 - TFLite 进阶】 模组量化(下)

一.   概述

在边缘运算的重点技术之中，除了简化复杂的模组架构，来简化参数量以提高运算速度的这项模组轻量化网路架构技术之外。另一项技术就是各家神经网路框架 (TensorFlow、Pytorch etc…) 的 模组优化 能力，主要探讨 TensorFlow Lite 的 训练后之量化方式(Post-training quantization) 与 感知量化训练(Quantization-aware Training) ，依序分为上与下两篇幅，本篇将介绍后者资讯为主。所谓的量化就是将以最小精度的方式，来进行模组推理，使模组应用至各种 Edge Device 之中，并达到足够成本效益，如下图所示。顺带一提， 恩智浦 NXP i.MX8M Plus 的 NPU(Neural Processing Unit) 神经处理单元，属于纯整数的 AI 加速器，就仅适用于 8位元的整数运算才能获得最佳效益 !! 此系列的后续章节，也会利用 NPU 来实现算法加速之目的。

TensorFlow 模型应用概念之示意图

利用 TensorFlow Lite 量化方式 所构成的 模组 ，就是将训练完成的轻量化模组，透过量化与优化的方式来提升推理速度 !! 如下模型运作概念图所示，储存模型完成后，即可依序执行冻结模型、优化模型到最后轻量化模型 (.tflite)，让模型运行在移动式装置时可达到最佳化的效果。

 ※ MobileNet 模组是一种轻量化模组的架构，而此篇重点是如何透过模组量化转换为轻量化模组(tflite)

TensorFlow 模型运作概念之示意图

如下图所示，本系列是隶属于 机器学习开发环境 eIQ 之 推理引擎层 (Inference Engines Layer) 中的 TensorFlow Lite 进阶系列，故后续将向读者介绍 “模组量化(下)”
若新读者欲理解更多人工智能、机器学习以及深度学习的资讯，可点选查阅下方博文

 大大通精彩博文   【ATU Book-i.MX8系列】博文索引

TensorFlow Lite 进阶系列博文-文章架构示意图

二.  量化理论

何谓量化 ? 在此文章是泛指数值程度上的量化，亦指有限范围的数值表示方式。其作用是为了降低数值资料量与模组大小，来提升传输与执行(推理)速度!! 而所谓的训练后之量化(Post-training quantization) 就是利用训练完成的模组，再次进行量化的一种优化方式。主要特色就是仅须要储存后的模组( SaveModel / .h5 /ckpt)，且不需要训练时的资料库即可量化。

举例来说，如下图所示，是须将原本数值分布为 -3e³⁸ 到 +3e³⁸ 的浮点数型态 float，量化为数值分布 -2³¹  到 2³¹ 的整数型态 int ，并以原本数据的最大值与最小值来找出有效的数值范围，将有一定概率大幅度减少资料量。

数据量化示意图

如何量化? 下列以最实际公式进行演示

假设原始数值的范围为 [-2 : 6.0] 的浮点表示，将其量化至目标范围 [-128 : 127] 的 8bit 整数范围

第一步，找出量化后能够表示的最小刻度

第二步，找出相对应的量化定点值

第三步，找出相对应的量化定点值

即可找出浮点数为 3.0 时，所对应的量化数值为 32。若将上述量化方式，将浮点数数值范围量化为整数范围，即如同下方表格所示。

量化优势? 劣势 ? 对于 TensorFlow Lite轻量化的应用而言

优势：

- 减少模组尺寸：最多能缩减 75% 的大小
- 加快推理速度 : 使用整数计算大幅度提升速度
- 支援硬体较佳 : 能使处理八位元的处理器进行推理
- 传输速度提升 : 因模组尺寸缩小，能更获得更好的传输品质

缺点 :

    -  精度损失 : 因为数值的表示范围缩减，故模组的准确度将会大幅度的降低

三.  感知量化训练(Quantization-aware Training)

感知量化训练(Quantization-aware Training, QAT) 亦是一种量化手段，其原理与上一小节所介绍的量化方式雷同，目的也是以降低精度的方式，缩小模组所需计算的资料量，来提升模组运算速度，且保持一定准确度的一种优化手段。相较于上一小节所介绍的 训练后之量化方式(Post-training quantization) 最大的不同，就是需要利用 原生模组 与 训练集(DataSets) 来作重新训练，而 感知量化训练 会于训练时，去模拟低精度的运算，来保持最佳的模组准确度。

故理论上， ‘感知量化训练量化’ 的准确度会来得比 ‘训练后之量化方式’ 来的准确，如下图所示 ; 感知量化训练的模组(QAT Model) 准确度能够逼近于原生模组(Baseline Model)，反之，训练后之量化方式(Post-Training full quantized Model) 则降低了约莫 0.05 的准确度。但其实该量化方式存在比较多不稳定因素，像是各机器学习框架的转换、版本不匹配、缺乏正确训练资料集或是不能微调等等因素，故官方比较推荐使用 ‘感知量化训练量化’ 来获得更稳定的量化体验。

 ※ 补充说明 : 目前仅应用于 Keras ( TensorFLow 2.X ) 的所建立的模组或 API

TensorFlow Lite 各种量化方式的准确度示意图
资料出处 – 官方文件

 使用方式 : 

大致上 Quantization Aware Training 的应用核心 可以分为四个步骤，分别为建立原生模组、建立感知量化模组、进行感知量化训练、进行转换等步骤。

 必要套件 :

$ pip install -q tensorflow
$ pip install -q tensorflow-model-optimization

Step 1 :  建立原生模组

感知量化模型是需要搭配 原生模型 与 原生模型所训练的资料集。这里利用 MNIST 手写识别的范例来进行演示，如下代码所示 ; 包含模组架构建立，以及利用 MNIST DataSets 进行训练。

import tempfile
import os
import tensorflow as tf
from tensorflow import keras

# Load MNIST dataset
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
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='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=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),metrics=['accuracy'])
model.fit( train_images, train_labels,epochs=1, validation_split=0.1, batch_size=32 )

其原生模组架构，如下图所示。

 Step 2 : 建立感知量化模组 (普通用法)

最简单的感知量化训练方式，就是直接利用 Tensorflow Model Optimization 的量化套件进行应用。
如同下代码与结果所示，将 原生模组 代入至 quantize_model 量化套件中，即构成感知量化的模组架构 ; 若仔细观察的话，则会发现架构层的名称皆冠上 quant 的字眼，也就是程式会去模拟低精度的运算，亦可称 Fake Quantization 。

import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model

其感知量化模组架构，如下图所示。

Step 2 : 建立感知量化模组 (进阶用法)

进阶的感知量化训练方式，就是直接利用Tensorflow Model Optimization 的量化套件进行微调，选择适当的架构层进行量化。
如下代码所示，将利用 clone_model 复制原生模组(baseline model) 并透过 apply_quantization_to_dense 来选择须量化的架构层，而此范例仅量化 Dense 架构层。

def apply_quantization_to_dense(layer):
    if isinstance(layer, tf.keras.layers.Dense):
        return tfmot.quantization.keras.quantize_annotate_layer(layer)
    return layer

annotated_model = tf.keras.models.clone_model(
    base_model,
    clone_function=apply_quantization_to_dense,
)

qat_model = tfmot.quantization.keras.quantize_apply(annotated_model)
qat_model.compile(optimizer='adam', loss=tf.keras.losses.SpareCategoricalCrossentropy(from_logits=True),metrics=['accuracy'])

其感知量化模组架构，如下图所示，请读者仔细观察与普通用法的不同，就能够发现前面几架构层，以无 quant 的字眼，表示这是原生的架构层。

若欲理解更多进阶用法，可以参考官方范例以及 查看可量化的架构层(请搜寻Default8BitQuantizeRegistry 的字眼，查阅所描述的架构层)，同理，若欲尝试量化没有支援的架构层，可以透过 QuantizeConfig 方式进行量化，可以参考 Medium 网志的解说 !!

Step 3 : 进行感知量化训练

接着，需要对 感知量化模组(q_ware_model) 再次进行训练，得以模拟低精度运算，如下代码所示。

train_images_subset = train_images[0:1000] 
train_labels_subset = train_labels[0:1000]
q_aware_model.fit(train_images_subset, train_labels_subset,batch_size=500, epochs=1, validation_split=0.1)
q_aware_annotate_model.fit(train_images_subset, train_labels_subset,batch_size=500, epochs=1, validation_split=0.1)

训练完成后，即可验证模组的准确度的表现状况，如下代码与结果所示 ; Baseline test accuary 为原始模组所呈现的准确度，反之 Quant test accuracy 为感知量化的准确度 !!  而感知量化的方式略高于 0.02 的准确度，故更好的准确度表现 !!

_, baseline_model_accuracy = model.evaluate(
    test_images, test_labels, verbose=0)

_, q_aware_model_accuracy = q_aware_model.evaluate(
   test_images, test_labels, verbose=0)

_, q_aware_annotate_model_accuracy = q_aware_annotate_model.evaluate(
   test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Quant test accuracy:', q_aware_model_accuracy)
print('Quant test accuracy(annotate):', q_aware_annotate_model_accuracy)

 Step 4 : 进行转换

基本上，前面步骤已经完成大致 感知量化训练 的操作，但因为最终目标是须应用于移动装置之中。故需要于感知量化训练且生成新的模组之后，利用上一小节的 训练后之量化方式(Post-training quantization) 的方式进行量化转换，如下代码所示 !! 然而，细心的读者应该可以发现一些小细节，就是这个代码又做一次优化，这里先卖个关子，原因将置于结论再向读者探讨 !!

# quantized_aware_tranning_model(Dynamic)
run_model = tf.function(lambda x: q_aware_model(x))
concrete_func = run_model.get_concrete_function(tf.TensorSpec([1,28,28], model.inputs[0].dtype))
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
converter.optimizations = [tf.lite.Optimize.DEFAULT]
quantized_aware_tranning_model_dynamic = converter.convert()
with open("quantized_aware_tranning_model_dynamic.tflite",'wb') as f:
    f.write(quantized_aware_tranning_model_dynamic)
print("quantized_aware_tranning_model_dynamic done!!")

完成后，会生成对应的 .tflite 档案，即可直接应用!!

  量化使用分析 : 

这里以 MNIST 手写数字识别模组为基准，将测试原生模组 与 感知量化训练(普通用法与进阶用法) 所生成的模组，来搭配训练后之量化等转换方式来验证准确度为何 !! 同时，也测试经过训练后之量化的方式，是否对于模组准确度或是应用有何影响? 其测试代码就如同上一章所介绍的方式，或可以直接查看以及执行运行 Colab 代码。

 实验测试数据结果 :

其测试准确度数据结果如下 :

TensorFlow Lite 各种量化转换之准确度分析

同时，将模组移植至 NXP i.MX8M Plus 并使用 AI 芯片 NPU 来测试推理速度数据结果如下 :
 ※ 补充说明 :  测试版本为BSP L5.15.5 版本

 TensorFlow Lite 各种量化转换之推理速度分析 (NPU)

亦将模组移植至 NXP i.MX8M Plus 并使用 ARM 核心加速(XNNPACK) 来测试推理速度，如下 :
 ※ 补充说明 : 其中 X 表示模组无法正常运行，暂不列入考量

TensorFlow Lite 各种量化转换之推理速度分析 (ARM)

四.  结语

依目前实验结果而论，感知量化训练(Quantization Aware Training) 能够尽可能去逼近原始模组的准确度，甚至还可能有些许的小幅度提升。而就推理速度来看，感知量化训练的普通用法其实相当于训练后之量化的全整数量化之结果，也表示模组内的参数已转换为低精度的表示方式 !! 换句话说，若实现过 感知量化训练 的优化则不必多此一举再做一次训练后之量化的优化。但碍于硬体加速器或处理器的不同，可能会出现各式各样的问题，就如同上述表格中的 X ，这就代表感知量化训练所生成的模组，仍有一定机率不能顺利运行于 NPU 上。故重新结合训练后之量化方式来达到更加的应用体验 !! 此外，读者不仿思考一下，为何感知量化训练的普通用法与进阶用法，所呈现的推理速度表现会有落差 ? 其实此结果是完全符合，上述所向各位描述的一致，也就是进阶用法仅量化了完全连接层(Dense) 一个架构层而已，自然表现就会比普通用法来得慢 !! 最后，探讨一下是否推荐使用‘感知量化训练’ 来进行优化? 以作者角度而言，若是开发者刚好使用 Keras 框架来开发模组的话，是个很棒的选择 !! 而事实上，每个神经网路框架都有拥护者，不可能所有人都用此框架开发。故对于活用度而言，感知量化训练略显于不足，故仍推荐先活用 训练后之量化的方式 进行优化 !! 倘若，有闲之时再将框架、模组、算法移植至 Keras 来实现感知量化训练，才能发会最大的应用效益 !!  因此后续章节，将会向各位介绍不同的神经网路或深度学习框架之间的转换，敬请期待 !!

五.  参考文件

[1] 官方文件 - i.MX Machine Learning User's Guide pdf
[2] 官方文件 - TensorFlow Lite 转换工具
[3] 官方文件 - Quantization Aware Training
[4] 官方文件 - Colab
[5] 第三方文件 - Tensorflow模型量化(Quantization)原理及其实现方法

如有任何相关 TensorFlow Lite 进阶技术问题，欢迎至博文底下留言提问 !!
接下来还会分享更多 TensorFlow Lite 进阶文章 !!敬请期待 【ATU Book-i.MX8系列 – TFLite 进阶】 !!