深度学习算法优化系列六|使用TensorFlowLite对LeNet进行训练中量化

在深度学习算法优化系列三 | Google CVPR2018 int8量化算法 这篇推文中已经详细介绍了Google提出的Min-Max量化方式&＃xff0c;关于原理这一小节就不再赘述了&＃xff0c;感兴趣的去看一下那篇推文即可。昨天已经使用tflite测试了训练后量化&＃xff0c;所以今天主要来看一下训练时量化时怎么做的。注意训练中的量化实际上是伪量化&＃xff0c;伪量化是完全量化的第一步&＃xff0c;它只是模拟了量化的过程&＃xff0c;并没有实现量化&＃xff0c;只是在训练过程中添加了伪量化节点&＃xff0c;计算过程还是用float32计算。
然后训练得出.pb文件&＃xff0c;放到指令TFLiteConverter里去实现第二步完整的量化&＃xff0c;最后生成tflite模型&＃xff0c;实现int8计算。

• tensorflow 1.31.1
• python3.5

导入一些需要用到的头文件。

#coding&＃61;utf-8
import re
import time
import numpy as np
import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim import get_variables_to_restore
from tensorflow.python.framework import graph_util
import tensorflow.examples.tutorials.mnist.input_data as input_data


设置一些超参数&＃xff0c;分别为dropout层的丢弃比率&＃xff0c;学习率&＃xff0c;批量大小&＃xff0c;模型需要保存的路径以及训练的迭代次数。

# 参数设置
KEEP_PROB &＃61; 0.5
LEARNING_RATE &＃61; 1e-5
BATCH_SIZE &＃61; 30
PARAMETER_FILE &＃61; "./checkpoint/variable.ckpt-100000"
MAX_ITER &＃61; 100000


构建我们的训练网络&＃xff0c;这里使用LeNet&＃xff0c;想使用其他网络或者自己的网络相应修改即可。注意一下这里使用了tensorflow中的变量重用函数&＃xff0c;方便的控制在测试阶段不使用Dropout。关于Lenet可以详细的看一下我之前的推文&＃xff0c;地址如下&＃xff1a;卷积神经网络学习路线&＃xff08;六&＃xff09;| 经典网络回顾之LeNet 同时在LeNet类中已经定义好损失函数和优化器。在训练中引入量化的操作要复杂一些&＃xff0c;首先训练时在损失计算后面&＃xff0c;优化器定义前面需要引入tf.contrib.quantize.create_training_graph()。代码如下&＃xff1a;其中tf.contrib.quantize.create_training_graph()函数中的80000代表程序迭代到8w次之后才开始做量化&＃xff0c;原理在深度学习算法优化系列三中讲了。因为这里要训练10w次&＃xff0c;所以后2w次会做训练时的伪两家

# Build LeNet
class Lenet:def __init__(self, is_train&＃61;True):self.raw_input_image &＃61; tf.placeholder(tf.float32, [None, 784], "inputs")self.input_images &＃61; tf.reshape(self.raw_input_image, [-1, 28, 28, 1])self.raw_input_label &＃61; tf.placeholder("float", [None, 10], "labels")self.input_labels &＃61; tf.cast(self.raw_input_label, tf.int32)self.dropout &＃61; KEEP_PROBself.is_train &＃61; is_trainwith tf.variable_scope("Lenet") as scope:self.train_digits &＃61; self.build(True)scope.reuse_variables()self.pred_digits &＃61; self.build(False)self.loss &＃61; slim.losses.softmax_cross_entropy(self.train_digits, self.input_labels)# 获取当前的计算图&＃xff0c;用于后续的量化self.g &＃61; tf.get_default_graph()if self.is_train:# 在损失函数之后&＃xff0c;优化器定义之前&＃xff0c;在这里会自动选择计算图中的一些operation和activation做伪量化tf.contrib.quantize.create_training_graph(self.g, 80000)self.lr &＃61; LEARNING_RATEself.train_op &＃61; tf.train.AdamOptimizer(self.lr).minimize(self.loss)else:# 用于预测时&＃xff0c;将之前训练时构造的伪量化的operation和activation实际量化&＃xff0c;用于后续的推断tf.contrib.quantize.create_eval_graph(self.g)self.predictions &＃61; tf.arg_max(self.pred_digits, 1, name&＃61;"predictions")self.correct_prediction &＃61; tf.equal(tf.argmax(self.pred_digits, 1), tf.argmax(self.input_labels, 1))self.train_accuracy &＃61; tf.reduce_mean(tf.cast(self.correct_prediction, "float"))def build(self, is_trained&＃61;True):with slim.arg_scope([slim.conv2d], padding&＃61;&＃39;VALID&＃39;,weights_initializer&＃61;tf.truncated_normal_initializer(stddev&＃61;0.01),weights_regularizer&＃61;slim.l2_regularizer(0.0005)):net &＃61; slim.conv2d(self.input_images, 6, [5, 5], 1, padding&＃61;&＃39;SAME&＃39;, scope&＃61;&＃39;conv1&＃39;)net &＃61; slim.max_pool2d(net, [2, 2], scope&＃61;&＃39;pool2&＃39;)net &＃61; slim.conv2d(net, 16, [5, 5], 1, scope&＃61;&＃39;conv3&＃39;)net &＃61; slim.max_pool2d(net, [2, 2], scope&＃61;&＃39;pool4&＃39;)net &＃61; slim.conv2d(net, 120, [5, 5], 1, scope&＃61;&＃39;conv5&＃39;)net &＃61; slim.flatten(net, scope&＃61;&＃39;flat6&＃39;)net &＃61; slim.fully_connected(net, 84, scope&＃61;&＃39;fc7&＃39;)net &＃61; slim.dropout(net, self.dropout, is_training&＃61;is_trained, scope&＃61;&＃39;dropout8&＃39;)digits &＃61; slim.fully_connected(net, 10, scope&＃61;&＃39;fc9&＃39;)return digits


开始训练LeNet的模型&＃xff0c;代码如下&＃xff1a;

def train():mnist &＃61; input_data.read_data_sets("MNIST_data/", one_hot&＃61;True)test_images &＃61; mnist.test.imagestest_labels &＃61; mnist.test.labelssess &＃61; tf.Session()batch_size &＃61; BATCH_SIZEparamter_path &＃61; PARAMETER_FILEmax_iter &＃61; MAX_ITERlenet &＃61; Lenet()variables &＃61; get_variables_to_restore()save_vars &＃61; [variable for variable in variables if not re.search("Adam", variable.name)]saver &＃61; tf.train.Saver(save_vars)sess.run(tf.initialize_all_variables())# 用来显示标量信息tf.summary.scalar("loss", lenet.loss)# merge_all 可以将所有summary全部保存到磁盘&＃xff0c;以便tensorboard显示。如果没有特殊要求&＃xff0c;# 一般用这一句就可一显示训练时的各种信息了。summary_op &＃61; tf.summary.merge_all()# 指定一个文件用来保存图train_summary_writer &＃61; tf.summary.FileWriter("logs", sess.graph)for i in range(max_iter):batch &＃61; mnist.train.next_batch(batch_size)if i % 100 &＃61;&＃61; 0:train_accuracy, summary &＃61; sess.run([lenet.train_accuracy, summary_op], feed_dict&＃61;{lenet.raw_input_image: batch[0],lenet.raw_input_label: batch[1]})train_summary_writer.add_summary(summary)print("step %d, training accuracy %g" % (i, train_accuracy))if i % 500 &＃61;&＃61; 0:test_accuracy &＃61; sess.run(lenet.train_accuracy, feed_dict&＃61;{lenet.raw_input_image: test_images,lenet.raw_input_label: test_labels})print("\n")print("step %d, test accuracy %g" % (i, test_accuracy))print("\n")sess.run(lenet.train_op, feed_dict&＃61;{lenet.raw_input_image: batch[0],lenet.raw_input_label: batch[1]})saver.save(sess, paramter_path)print("saved model")


训练完成之后我们会得到checkpoint模型文件&＃xff1a;

接下来我们需要将checkpoint模型文件转换成freeze pb文件&＃xff0c;当然保存为saved_model模型也是可以的&＃xff0c;因为昨天介绍过了保存为saved_model然后转换为tflite模型&＃xff0c;所以今天介绍这种新的也是更加常用的转换方式。将checkpoint模型保存为pb模型的代码如下&＃xff1a;

# 将存储了伪量化信息的freeze pb文件转换成完全量化的tflite文件&＃xff0c;可以看见量化完之后文件内存基本减小到1/4
def convert_to_tflite():converter &＃61; tf.lite.TFLiteConverter.from_frozen_graph("pb_model/freeze_eval_graph.pb", ["inputs"], ["predictions"])converter.inference_type &＃61; tf.lite.constants.QUANTIZED_UINT8converter.quantized_input_stats &＃61; {"inputs":(0., 1.)} # mean, std_dev&＃xff0c;需要自己从训练集&＃xff08;增强后&＃xff0c;输入网络之前的&＃xff09;统计出来converter.allow_custom_ops &＃61; Trueconverter.default_ranges_stats &＃61; (0, 255)converter.post_training_quantize &＃61; Truetflite_model &＃61; converter.convert()open("tflite_model/eval_graph.tflite", "wb").write(tflite_model)


这样就会在tflite_model文件夹下生成eval_graph.tflite模型了。

注意下面几点&＃xff1a;

• ["inputs"], ["predictions"]是freeze pb中的输入节点和输出节点。
• quantized_input_states是定义输入的均值和方差&＃xff0c;tensorflow lite的文档中说这个mean和var的计算方式是&＃xff1a;mean是0到255之间的整数值&＃xff0c;映射到浮点数 0.0f。std_dev &＃61; 255 /&＃xff08;float_max - float_min。按照文档的默认数值效果也还行&＃xff0c;这里只是为了展示整个量化过程就不再仔细的去推敲mean和stddev的取值了&＃xff0c;这里用默认值来做。
• default_ranges_states是指量化后的值的范围&＃xff0c;其中255就是2^8 - 1。
• 官方文档如下&＃xff1a;
  
  同时跟进这个函数可以发现&＃xff1a;

可以大概算出来mean和stddev取[0,1]是满足这个关系的。

接下来我们写3个测试函数来测试一下原始的checkpoint模型&＃xff0c;pb模型和tflite模型的表现。代码如下&＃xff1a;

# 使用原始的checkpoint进行预测
def origin_predict():mnist &＃61; input_data.read_data_sets("MNIST_data/", one_hot&＃61;True)sess &＃61; tf.Session()saver &＃61; tf.train.import_meta_graph("./checkpoint/variable.ckpt-100000.meta")saver.restore(sess, "./checkpoint/variable.ckpt-100000")input_node &＃61; sess.graph.get_tensor_by_name(&＃39;inputs:0&＃39;)pred &＃61; sess.graph.get_tensor_by_name(&＃39;predictions:0&＃39;)labels &＃61; [label.index(1) for label in mnist.test.labels.tolist()]predictions &＃61; []start_time &＃61; time.time()for i in range(10):for image in mnist.test.images:prediction &＃61; sess.run(pred, feed_dict&＃61;{input_node: [image]}).tolist()[0]predictions.append(prediction)end_time &＃61; time.time()correct &＃61; 0for prediction, label in zip(predictions, labels):if prediction &＃61;&＃61; label:correct &＃43;&＃61; 1print(correct / len(labels))print((end_time - start_time))sess.close()# 使用freeze pb文件进行预测
def freeze_pb_predict():mnist &＃61; input_data.read_data_sets(&＃39;MNIST_data/&＃39;, one_hot&＃61;True)with tf.Session() as sess:with tf.gfile.FastGFile("pb_model/freeze_eval_graph.pb", &＃39;rb&＃39;) as f:# 使用tf.GraphDef()定义一个空Graphgraph_def &＃61; tf.GraphDef()graph_def.ParseFromString(f.read())# 返回一个上下文管理器,使得这个Graph对象成为当前默认的graph.当你想在一个进程里面创建# 多个图的时候,就应该使用这个函数.为了方便起见,一个全局的图对象被默认提供,要是你没有# 显式创建一个新的图的话,所有的操作(ops)都会被添加到这个默认的图里面来.#sess.graph.as_default()# 导入Graphtf.import_graph_def(graph_def, name&＃61;&＃39;&＃39;)# tf.global_variables_initializer()添加节点用于初始化所有的变量(GraphKeys.VARIABLES)。# 返回一个初始化所有全局变量的操作&＃xff08;Op&＃xff09;。在你构建完整个模型并在会话中加载模型后&＃xff0c;运行这个节点。sess.run(tf.global_variables_initializer())# 获取输出Tensor和输出Tensorinput_node &＃61; sess.graph.get_tensor_by_name(&＃39;inputs:0&＃39;)pred &＃61; sess.graph.get_tensor_by_name(&＃39;predictions:0&＃39;)labels &＃61; [label.index(1) for label in mnist.test.labels.tolist()]predictions &＃61; []start_time &＃61; time.time()for image in mnist.test.images:prediction &＃61; sess.run(pred, feed_dict&＃61;{input_node: [image]}).tolist()[0]predictions.append(prediction)end_time &＃61; time.time()correct &＃61; 0for prediction, label in zip(predictions, labels):if prediction &＃61;&＃61; label:correct &＃43;&＃61; 1print(correct / len(labels))print((end_time - start_time))# 使用tflite进行预测
def tflite_predict():mnist &＃61; input_data.read_data_sets("MNIST_data/", one_hot&＃61;True)labels &＃61; [label.index(1) for label in mnist.test.labels.tolist()]images &＃61; mnist.test.imagesimages &＃61; np.array(images, dtype&＃61;"uint8")# 根据tflite文件生成解析器interpreter &＃61; tf.contrib.lite.Interpreter(model_path&＃61;"./tflite_model/eval_graph.tflite")# 用allocate_tensors()分配内存interpreter.allocate_tensors()# 获取输入输出tensorinput_details &＃61; interpreter.get_input_details()output_details &＃61; interpreter.get_output_details()predictions &＃61; []start_time &＃61; time.time()for i in range(10):for image in images:# 填充输入tensorinterpreter.set_tensor(input_details[0][&＃39;index&＃39;], [image])# 前向推理interpreter.invoke()# 获取输出tensorscore &＃61; interpreter.get_tensor(output_details[0][&＃39;index&＃39;])[0][0]# # 结果去掉无用的维度# result &＃61; np.squeeze(score)# #print(&＃39;result:{}&＃39;.format(result))# # 输出结果是长度为10&＃xff08;对应0-9&＃xff09;的一维数据&＃xff0c;最大值的下标就是预测的数字predictions.append(score)end_time &＃61; time.time()correct &＃61; 0for prediction, label in zip(predictions, labels):if prediction &＃61;&＃61; label:correct &＃43;&＃61; 1print((end_time - start_time))print(correct / len(labels))


这里值得注意的问题是&＃xff0c;测试其中一个模型的时候我们需要注释掉另外两个&＃xff0c;因为我发现在tensorflow-1.13.1中一个使用同一个py文件启动了2个Session的话&＃xff0c;第二个Session执行的内容会出问题&＃xff0c;当然我使用第一个Session之后是确认关闭了的&＃xff0c;所以我认为这可能是一个BUG。我的源码也放github了&＃xff0c;有兴趣的话可以帮忙看看是不是我代码的问题。

我在测试的时候模型掉了大约四个点&＃xff0c;我们可以大概想一下为什么这里精度掉得还是比较多&＃xff0c;可能对于Mnist数据集来说LeNet已经算是一个大模型了&＃xff0c;所以训练后量化比训练时量化表现更好。并且这个模型再训练10w轮&＃xff0c;量化后模型的精度可以降低到2个百分点。

这两天测试了一下TFite的int8量化&＃xff0c;可以看出来量化确实是有效果的&＃xff0c;并且tflite也确实可以实现加速&＃xff0c;况且这是在pc端的测试&＃xff0c;放到设备端加速应该会更明显吧。

• Tensorflow-Lite官方文档&＃xff1a;https://tensorflow.google.cn/lite
• Tensorflow后量化官方实例&＃xff1a;https://github.com/tensorflow/tensorflow/blob/d035a83459330c87bbc527e3d480b65f32841997/tensorflow/contrib/lite/tutorials/post_training_quant.ipynb
• 我的github地址&＃xff1a;https://github.com/BBuf/model_quantization

欢迎关注我的微信公众号GiantPandaCV&＃xff0c;期待和你一起交流机器学习&＃xff0c;深度学习&＃xff0c;图像算法&＃xff0c;优化技术&＃xff0c;比赛及日常生活等。