深度学习 10 tf2模型转tflite，并进行权重量化_tf模型转tflite模型

在tf2官方给出的量化方式有两种，训练后量化和感知量化。训练后量化是指模型训练好之后在进行 权重量化；感知量化是指在训练过程中自动根据训练数据的反馈进行量化。感知量化对模型的精度影响最小。

1 训练后量化

# -*- coding: utf-8 -*-
import tensorflow as tf
import numpy as np
import glob,os,cv2
 
# 生成代表性数据集
def get_representative_dataset():
    img_dir = r"数据"
    img_ext = ".bmp"
    img_paths = glob.glob(os.path.join(img_dir, '*' + img_ext))
    for path in img_paths:
        img=cv2.imdecode(np.fromfile(path,dtype=np.uint8),0)
        img= cv2.resize(img, (512,512), interpolation = cv2.INTER_NEAREST)/255
        img=img.reshape((1,512,512,1)).astype(np.float32)
        yield  img,img
        
#如果使用官方文档提供的下面的converter获取方法，会发现，最后转成出来的lite模型的输入依然为float类型。
converter = tf.lite.TFLiteConverter.from_saved_model('saved_model')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()
open("float_model.tflite", "wb").write(tflite_quant_model)
 
#----------------加载并获取基本信息---------------------
interpreter = tf.lite.Interpreter(model_path='float_model.tflite')
#interpreter = tf.lite.Interpreter(model_content=concrete_func_tflite)
interpreter.allocate_tensors()#给所有的tensor分配内存
 
#获取 input 和 output tensor
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
 
#输出字典，我们需要利用里面的输出来设置输入输出的相关参数
print(input_details)
print(output_details)
 
#----------------进行预测---------------------
input_shape = input_details[0]['shape']#设置输入维度
data_gen=get_representative_dataset()
x,x=data_gen.__next__()
input_data = tf.constant(x)
interpreter.set_tensor(input_details[0]['index'],input_data)
 
#执行预测
interpreter.invoke()
output_results = interpreter.get_tensor(output_details[0]['index'])
print(output_results.shape)#十个标签概率值的概率分布
import matplotlib.pyplot as plt
plt.imshow(output_results[0,:,:,0])
plt.show()
plt.imshow(x[0,:,:,0])
plt.show()

2 训练前量化

tf中有一些layer不支持量化，因此需要进行规避（不对该层量化）

可以根据类型进行规避，也可以根据layer的name进行规避

from keras import layers
from keras.models import load_model
import tensorflow as tf
model=load_model('model_name.h5')
import tensorflow_model_optimization as tfmot
from tensorflow.keras import layers
import tensorflow as tf
# Helper function uses `quantize_annotate_layer` to annotate that only the 
# Dense layers should be quantized.
def apply_quantization_to_dense(layer):
  if isinstance(layer, tf.keras.layers.Conv2DTranspose):#根据类型进行规避
      return layer
  elif isinstance(layer, tf.keras.layers.BatchNormalization):
      return layer
  elif isinstance(layer, tf.keras.layers.Concatenate):
      return layer
  elif isinstance(layer, tf.keras.layers.Activation):
      return layer
  elif layer.name=='img_out':#根据name进行规避
      return layer
  return tfmot.quantization.keras.quantize_annotate_layer(layer)
annotated_model = tf.keras.models.clone_model(
    model,
    clone_function=apply_quantization_to_dense,
)
#已经被量化的模型
model = tfmot.quantization.keras.quantize_apply(annotated_model)
model.summary()
#后续即可正常训练使用

更多权重量化方式请参考  
https://tensorflow.google.cn/lite/performance/post_training_quantization