使用TFLite对超分辨率模型进行INT8量化

神经网络量化

神经网络量化是将网络权重由高精度（FLOAT32）转换为低精度（INT8）的过程，量化后的网络体积变小，计算量下降，因此更适用于移动设备的部署，并且精度不会损失太多。

TFLite模型的INT8量化

假设有一个训练好的TensorFlow超分辨率模型model，现在我们希望使用TFLite对其进行量化，从而部署到移动设备中。在量化之前，我们需要对输入数据进行建模，TFLite通过统计分析，得到最优的量化参数。具体实现上，我们需要写一个生成器来返回输入图像数据。举例如下：

def representation_dataset_gen():
	size = 100
	for i in range(size):
		lr_path = "data/LR/%04d.pt"%(i+1)
		print("representative data: %03d/%03d"%(i+1, size))
		with open(lr_path, "rb") as f:
			lr = pickle.load(f)
		lr = lr.astype(np.float32)
		lr = np.expand_dims(lr, 0)
		yield [lr]

在此，我们设计了函数representation_dataset_gen()，生成100张具有代表性的输入图片。

TFLite的量化过程如下：首先我们需要设定好图像的输入尺度，然后读取训练好的TensorFlow模型，通过TFLite提供的量化器对模型进行量化，具体实现如下：

def quantize(model_path, quantized_model_path):
	tensor_shape = [1, None, None, 3]	# 输入图像大小，None代表长宽不固定
	rep = representation_dataset_gen	# 输入图像生成器
	model = tf.saved_model.load(model_path)		# 读入训练模型
	
	# 获取TensorFlow模型的推理图，并且固定输入图像的尺寸
	concrete_func = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
	concrete_func.inputs[0].set_shape(tensor_shape)

	# 将TensorFlow模型转换为TFLite，并进行量化
	converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
	converter.experimental_new_converter = True
	converter.experimental_new_quantizer = True
	converter.optimizations = [tf.lite.Optimize.DEFAULT]
	converter.representative_dataset = rep
	converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
	converter.inference_input_type = tf.uint8
	converter.inference_output_type = tf.uint8
	
	# 得到量化模型
	quantized_tflite_model = converter.convert()
	# 保存模型
	open(quantized_model_path, "wb").write(quantized_tflite_model)

由此，我们便得到了INT8量化后的TFLite模型，并可以将其部署到Android工程中（见我上篇博文）