tensorflow QAT

tensorflow qat

https://www.wpgdadatong.com/tw/blog/detail/70672

在边缘运算的重点技术之中，除了简化复杂的模块构架，来简化参数量以提高运算速度的这项模块轻量化网络构架技术之外。另一项技术就是各家神经网络框架（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 模型運作概念之示意圖
何谓量化？在此文章是泛指数值程度上的量化，亦指有限范围的数值表示方式。其作用是为了降低数值数据量与模块大小，来提升传输与执行（推理）速度！！而所谓的训练后之量化（Post-training quantization）就是利用训练完成的模块，再次进行量化的一种优化方式。主要特色就是仅须要储存后的模块（SaveModel / .h5 /ckpt），且不需要训练时的数据库即可量化。
举例来说，如下图所示，是须将原本数值分布为-3e38到+3e38的浮点数型态float，量化为数值分布-231到231的整数型态int，并以原本数据的最大值与最小值来找出有效的数值范围，将有一定概率大幅度减少数据量。
在这里插入图片描述
量化优势？劣势？对于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的准确度。但其实该量化方式存在比较多不稳定因素，像是各机器学习框架的转换、版本不匹配、缺乏正确训练数据集或是不能微调等等因素，故官方比较推荐使用『感知量化训练量化』来获得更稳定的量化体验。

使用方式 :

大致上 Quantization Aware Training 的應用核心可以分為四個步驟，分別為建立原生模組、建立感知量化模組、進行感知量化訓練、進行轉換等步驟。

必要套件 :

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

Step 1：建立原生模块

port 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 )

感知量化模型是需要搭配原生模型与原生模型所训练的数据集。这里利用MNIST手写识别的示例来进行演示，如下代码所示；包含模块构架建立，以及利用MNIST DataSets进行训练。
其原生模块构架，如下图所示。
在这里插入图片描述
Step 2：建立感知量化模块（普通用法）
最简单的感知量化训练方式，就是直接利用Tensorflow Model Optimization的量化套件进行应用。
如同下代码与结果所示，将原生模块代入至quantize_model量化套件中，即构成感知量化的模块构架；若仔细观察的话，则会发现构架层的名称皆冠上quant的字眼，也就是程序会去模拟低精度的运算，亦可称Fake Quantization。

import tensorflow_model_optimization as tfmot
quantize_model = tfmot.quantization.keras.quantize_model
# q_aware stands for for quantization aware.
q_aware_model = quantize_model(model)
q_aware_model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])

q_aware_model.summary()

在这里插入图片描述
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代码。
实验测试数据结果：
其测试准确度数据结果如下：
在这里插入图片描述
四.结语
依目前实验结果而论，感知量化训练（Quantization Aware Training）能够尽可能去逼近原始模块的准确度，甚至还可能有些许的小幅度提升。而就推理速度来看，感知量化训练的普通用法其实相当于训练后之量化的全整数量化之结果，也表示模块内的参数已转换为低精度的表示方式！！换句话说，若实现过感知量化训练的优化则不必多此一举再做一次训练后之量化的优化。但碍于硬件加速器或处理器的不同，可能会出现各式各样的问题，就如同上述表格中的X，这就代表感知量化训练所生成的模块，仍有一定机率不能顺利运行于NPU上。故重新结合训练后之量化方式来达到更加的应用体验！！此外，读者不仿思考一下，为何感知量化训练的普通用法与进阶用法，所呈现的推理速度表现会有落差？其实此结果是完全符合，上述所向各位描述的一致，也就是进阶用法仅量化了完全连接层（Dense）一个构架层而已，自然表现就会比普通用法来得慢！！最后，探讨一下是否推荐使用『感知量化训练』来进行优化？以作者角度而言，若是开发者刚好使用Keras框架来开发模块的话，是个很棒的选择！！而事实上，每个神经网络框架都有拥护者，不可能所有人都用此框架开发。故对于活用度而言，感知量化训练略显于不足，故仍推荐先活用训练后之量化的方式进行优化！