tensorflow量化感知训练_TensorFlow量化训练

前段时间研究了tflite和量化相关的操作, 经测试量化尤其在具有专门DSP加速的硬件上(比如MTK8183)有着很好的加速效果,大约3X的提升; tensorflow提供了tflite转化工具toco,使用命令大致如下:

bazel-bin/tensorflow/contrib/lite/toco/toco --input_file=mobilenet_v1_1.0_128_frozen.pb \

--input_format=TENSORFLOW_GRAPHDEF \

--output_format=TFLITE \

--output_file=/tmp/mobilenet_v1_1.0_128.tflite \

--inference_type=FLOAT \

--input_data_types=FLOAT \

--input_arrays=input \

--output_arrays=MobilenetV1/Predictions/Reshape_1 \

--input_shapes=1,128,128,3 \

--logtostderr

这里面遇到的一个坑是,发现用pip直接安装的tensorflow, 运行这个命令会包error:, 因为option的名字不一样,那是用python重新封装的,而从源码编译toco(用官方的bazel工具)是C++的; 上述命令只是转化为tflite的格式,如果需要量化,只需要改成这样:

bazel-bin/tensorflow/contrib/lite/toco/toco --input_file=mobilenet_v1_1.0_128_frozen.pb \

--input_format=TENSORFLOW_GRAPHDEF \

--output_format=TFLITE \

--output_file=/tmp/mobilenet_v1_1.0_128.tflite \

--inference_type=QUANTIZED_UINT8 \

--input_data_types=QUANTIZED_UINT8 \

--input_arrays=input \

--output_arrays=MobilenetV1/Predictions/Reshape_1 \

--input_shapes=1,128,128,3 \

--logtostderr

对于, 一般的模型, 它会报错, 原因是没有提供MinMax的范围(量化需要这两个值进行缩放), 对于不同的类型,这两个值的获取方式不一样,对于里面的weights比较好办,只要统计一下最大最小值就好了,但是对于graph里面的其他tensor,因为输入的不同,所以每层输出的值也会不同,这里就需要在训练的时候去统计每个node的[Min, Max]; 当然你可以用这个语句去手动制定范围,但是精度可能就会收到很大影响.

tf.quantization.fake_quant_with_min_max_args

而且使用toco的时候需要加上:

--reorder_across_fake_quant=true

或者你也可以在使用toco的时候加上:

--default_ranges_min=0 \

--default_ranges_max=1

但是这样就更加糟糕了,它会给所有没有[Min,Max]的node,使用这两个值作为默认值   所以,由此我们知道: 1) 我们需要在训练的时候统计[Min, Max] 2) 为了量化能有更好的精度,我们希望[Min, Max]最好在一个比较小的范围内, 所以一些best practise包括: a)使用relu6作为激励函数; b) 使用batch normalization   Google提供了 quantization-aware的训练方法,可以参考

https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/quantize 简单总结下: 1) 在你定义好网络结构之后, 加上

tf.contrib.quantize.create_training_graph(input_graph=g,

quant_delay=2000000)

它会把网络重新改造,具体可以看tensorboar; 这里面,要注意quant_delay这个参数, 这个参数的意思的在此之前,网络会使用float进行正常的训练,收敛到一个比较好的水平;之后,会执行一个fake-quant的操作. 改造之前的一个卷积层:

改造之后的一个卷积层:

关键的操作就在于这个act_quant节点,统计了最大最小的值,然后执行量化操作,最上面的delay_quant里面有两个switch op, 就用根据设置的quant_delay来决定是否使用量化的结果进行forward

2) 按照正常的方式保存你的checkpoint; 3) 导出用于eval的网络,这个最好是在另外一个文件里面(否则可能会有这个问题

https://github.com/tensorflow/tensorflow/issues/19936),同样的方式创建网络结构后,调用:

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

然后load checkpoint,最后保存graph和新的checkpoint.

tf.train.write_graph(sess.graph_def, output_dir, 'graph.pb', as_text=True)

saver = tf.train.Saver(max_to_keep=100)

saver.save(sess, './workspace/'+args.model+'/chk', global_step=1)

4) freeze graph

python3 -m tensorflow.python.tools.freeze_graph \

--input_graph=./graph.pb \

--output_graph=exported_freezed_inference_graph.pb \

--input_checkpoint=./chk-1 \

--output_node_names="your_output_name"

5) 现在就可以使用toco命令导出量化的模型了;

也许等你执行完以上操作,发现自己的模型里面仍然会报某些node还是没有[Min, Max], 根据我的观察,在某些情况下(具体不明),发现只有添加了BatchNorm的层在调用cerate_train_graph的时候才会添加fake_quant; 但是有的时候不需要BatchNorm也可以,很奇怪~   最后,你还可以用这个命令查看你到处tflite量化模型的op和[Min,Max]的范围

bazel-bin/tensorflow/contrib/lite/toco/toco \

--input_file=exported_freezed_inference_graph.pb \

--input_format=TENSORFLOW_GRAPHDEF \

--output_format=GRAPHVIZ_DOT \

--output_file=exported_freezed_inference_graph_opt.dot \

--inference_type=QUANTIZED_UINT8 \

--input_data_types=QUANTIZED_UINT8 \

--input_arrays=image \

--output_arrays=Openpose/MConv_Stage1_L_5_pointwise/Conv2D \

--input_shapes=1,128,128,3 \

--logtostderr \

--default_ranges_min=0 \

--default_ranges_max=1

dot -Tpdf exported_freezed_inference_graph_opt.dot -o exported_freezed_inference_graph_opt.pdf