tensorflow的量化教程(2)

Google提供的 quantization-aware-trainning的量化训练方法,具体可以参考https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/quantize

对应的论文

https://arxiv.org/abs/1712.05877

还有这个白皮书

https://arxiv.org/abs/1806.08342

简单总结下:

1) 如何训练

 在你定义好网络结构之后, 加上下面这句话，即可量化训练

tf.contrib.quantize.create_training_graph(input_graph=g,
                                          quant_delay=2000000)

 其中，g为你创建好的模型，quant_delay表示网络在进行量化训练前 采用浮点训练的次数，在源码中，quant_delay默认为0。

这个默认值比较适合finetune一个已经训练好的模型。如果重头开始训，quant_delay应该设置为浮点模型收敛时的steps，模型训到此steps之后，量化训练便激活。（从头训练一个模型，quant_delay的值不一定非要取浮点模型收敛时的step值，有时取0，即直接量化训效果会更好）

 可以用tensorboard可视化模型，来查看量化训练是否成功，应该是像下面这样，嵌入伪量化操作

                                                                                               图1

          图2和图3分别为weights_quant和act_quant的展开

                                                                                                  图2 

                                                                                            图3         

      在量化训练时，需要统计出模型层融合后每层输出的[min, max]，如图3所示，至于这个[min, max]的作用，后面会讲到, 注意，还有一个[min,max]就是图2权重参数的，这两个意义不一样。权重参数的min max 可以直接算出，而每层输出的min max是在量化训练过程中通过指数移动均值统计出来。论文中提到，为了使量化训练有更好的精度，推荐使用relu6，让输出限制在较小的范围内，亲测有效。

2) 如何评估及转模型

 网络量化训练好以后，需要评估量化模型的性能，需要添加下面这句话

tf.contrib.quantize.create_eval_graph(input_graph=g)

注意， tf.contrib.quantize.create_eval_graph()和tf.contrib.quantize.create_training_graph()不能同时出现在同一程序中，不然会出问题。

在桌面PC或者服务器上用tensorflow训练好的模型不能直接用在TFLite上运行，需要先转换成.tflite格式文件才行。当然，现在的tensorflow已经支持在训练过程中直接生成.tflite文件，精简了部署流程，不过仅仅是对浮点模型有效，量化模型还不支持

下面介绍tflite文件的生成最常用的3步

1、在tensorflow训练脚本中保存图模型（.pb）和变量(.ckpt）)文件

2、利用freeze_graph脚本将图模型和权重文件冻结为一个文件(.pb)

3、利用toco工具，生成最终的tflite文件  

                                                           

 

下面详细说明每步具体流程

Tensorflow在训练期间通常不将权重存储在模型文件中，而是分开保存在一个叫checkpoints的检查点文件(.ckpt格式)。

模型和权重文件分开保存，

 

第一步，导出图模型文件和变量文件，在tensorflow训练脚本中加入相关代码即可，例如：

    #保存图模型，不含有权重信息

    g = tf.get_default_graph()

    graph_def = g.as_graph_def()

    tf.train.write_graph(graph_def, "./model", 'graph.pb', as_text=False)

    #保存训练期间的权重信息

    saver = tf.train.Saver()

    ……

    saver.save(sess, os.path.join( ‘./model’, 'model.ckpt'))

 

第二步，生成frozen_graph.pb文件（模型图结构文件和权重信息文件冻结为一个文件），从最新的检查点文件中提取所有变量的值，并将变量转换为常量，通过指定的输出节点，将与前向推理无关的外部节点去除，将常量固化到图里面，保存到指定的输出文件中。

有两种方法，一种是直接编译tensorflow-master/tensorflow/python/tools目录下的freeze_graph.py脚本，另一种是用bazel构建，这里，示例第一种：

python freeze_graph.py \

--input_graph=/path/to/graph.pb \

--input_checkpoint=/path/to/model.ckpt    \                     

--output_graph=/path/to/freeze.pb \

--output_node_names=fc6/act_quant/FakeQuantWithMinMaxVars   \

--input_binary=True

 

第三步，生成tflite，利用toco工具，可以直接在命令行使用，不需要bazel构建，命令如下：

1. 生成量化的tflite，例如：

toco \

 --output_file=quant_test.tflite \

--graph_def_file=freeze.pb \

--input_arrays=x-input  \

--output_arrays=fc6/act_quant/FakeQuantWithMinMaxVars \

--output_format=TFLITE \

--inference_type=QUANTIZED_UINT8\

--mean_values=0 \

--std_dev_values=1 

（如果是直接将浮点的pb转为量化的tflite，将 –inference_type=QUANTIZED_UINT8 改为—inference_type=FLOAT并添加

--default_ranges_min -default_ranges_max

  default_ranges_min和default_ranges_max是用于指定输出层最小最大值范围，由于没有训练，该值只能由用户自己指定，这个值得设定会影响最终的量化精度。虽然再怎么设定，精度也高不到哪去，但却可以模拟训练量化后模型在移动设备的性能

如果是量化训练好的pb文件，在转tflite时则不需要指定-default_ranges_min ，-default_ranges_max)

官方提供的量化训练好的模型可以在这里下载

https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet_v1.md

下面是在3288（4核A17，1.8GHZ）上的测试结果

跑mobilenet_v1(1.0, 224x224)，单线程浮点模型的结果如下，50次循环平均耗时421.637ms

 

跑mobilenet_v1，单线程量化模型的结果如下：50次循环平均耗时182.976ms

     人脸特征点定位模型误差性能指标

模型	浮点模型	直接量化模型	量化训练模型
MSE	0.000106	0.00289	0.000107

人脸特征点定位模型耗时，单位ms，蓝色单线程，红色双线程

tflite框架浮点双线程优化的不是很好（摊手.jpg）

 

 

---------------------------------------------------------------------------------------------------------------------------------------------------

过程中遇到的问题核解决方案

利用toco工具直接将浮点训练的.pb文件转换为量化的.tflite文件时，曾经遇到过下面这个问题

tflite_convert: error: --default_ranges_min and --default_ranges_max must be used together

 

我的tensorflow版本是1.9，据说将版本更新到1.10可以解决这个问题，不过没试过，怕版本更新后，会引起其他tensorflow工程的不兼容。后来采用了下面这种方法，成功解决这个问题。

pip install --user tf-nightly

 

另外安利一款看模型的神器Netron，用它打开固化话的pb模型，显示如下