天天看点

tensorboard ckpt pb 模型的输出节点_网络模型压缩量化(mnist的CNN举例)

一、背景

Neural Network模型一般都会占用很大的磁盘空间,比如AlexNet的模型文件就超过了200 MB.模型包含了数百万的参数,绝大部分的空间都用来存储这些模型的参数了。这些参数是浮点数类型的,普通的压缩算法很难压缩它们的空间。

  一般模型的内部的计算都采用了浮点数计算,浮点数的计算会消耗比较大的计算资源(空间和cpu/gpu时间),如果在不影响模型准确率的情况下,模型内部可以采用其他简单数值类型进行计算的话,计算速度会提高很多,消耗的计算资源会大大减小,尤其是对于移动设备来说,这点尤其重要。由此引入量化技术。

量化即通过减少表示每个权重所需的比特数来压缩原始网络。

二、量化理论

  对量化的实现是通过把常见操作转换为等价的八位版本达到的。涉及的操作包括卷积,矩阵乘法,激活函数,池化操作,以及拼接。对模型可以压缩到1/4。

我们发现模型属于同一层的参数值会分布在一个较小的区间内,比如-10, 30之间,我们可以记下这个最小值和最大值,在采用8位数量化(可以有其他选择)的情况下,可以把同一层的所有参数都线性映射(也可以采用非线性映射进一步压缩空间)到区间[-10, 30]之间的255个8位整数中的最接近的一个数,

这样模型文件的大小就基本压缩到了原来的25%,而且在加载的时候可以恢复到原来的数值,当然恢复的数值跟压缩之前会有差异,但事实证明,模型对这种压缩造成的噪音表现的很健壮。

满足第一个需求就这么简单,模型不需要改变,需要改变的是模型的存储和加载部分;但是要解决第二个问题就不这么简单了,事实证明,模型内部的计算如果都采用8位整形数来计算的话,模型的表现会相差很大,原因是模型的训练过程中,需要计算梯度,然后运用一定的学习算法,不断的更新参数,每次的更新量可能是很小的。所以模型的训练过程需要高精度的浮点型数值的。那么模型的预测过程是否可以使用整形数值代替呢?事实证明是可行的。

TensorFlow中模型的量化过程是这样实现的:

将模型中实现了对应的量化操作的操作符替换成量化操作符,经过转换后,输入输出依旧是float,将input从浮点数转换成8 bit,只不过中间的计算是用过8 bit来计算的。

tensorboard ckpt pb 模型的输出节点_网络模型压缩量化(mnist的CNN举例)

原始的计算图

Quantize:

quantize取input中的min和max,分别对应被量化的input中的最小值(0)和最大值(255),把[min, max]这个区间均匀分成255个小区间,把input中的值对应到对应的区间中。

Dequantize:

反量化操作则是把上述操作反向执行。

原因:

1. 权重、活化张量的数值通常分布在一个相对较小的范围中(weight:-15 ~ 15,activatios:-500 ~ 1000);

2. 神经网络对噪音的适应性强,将数量化到一个更小的数集中并不会对整体的结果带来很大的影响;

3. 通过量化操作,可以有效提高点乘的计算效率。

三、tensorflow量化感知训练

tensorflow量化感知训练是一种伪量化的过程,它是在可识别的某些操作内嵌入伪量化节点(fake quantization nodes),用以统计训练时流经该节点数据的最大最小值,便于在使用TOCO转换tflite格式时量化使用并减少精度损失,其参与模型训练的前向推理过程令模型获得量化损失,但梯度更新需要在浮点下进行因而其并不参与反向传播过程。某些操作无法添加伪量化节点,这时候就需要人为的去统计某些操作的最大最小值,但如果统计不准那么将会带来较大的精度损失,因而需要较谨慎检查哪些操作无法添加伪量化节点。值得注意的是,伪量化节点的意义在于统计流经数据的最大最小值并参与前向传播提升精确度,但其在TOCO工具转换为量化模型后,其工作原理还是与训练后量化方式一致的。过程举例说明如下图:

tensorboard ckpt pb 模型的输出节点_网络模型压缩量化(mnist的CNN举例)

量化后的计算图

可识别的ReLU节点嵌入了对应的伪量化节点QuantizedRelu,统计了其流经数据的max和min值,但存在quantize和dequantize的过程,其输入输出依然是float浮点型。而当多个可识别的操作相邻时其嵌入伪量化节点的形式如图八左半部分所示,但多个quantize和dequantize连接时是可以相互抵消的,图八右半部分展示了精简的伪量化节点连接的过程。

tensorboard ckpt pb 模型的输出节点_网络模型压缩量化(mnist的CNN举例)

优化量化处理后的计算图,去掉冗余操作

(1)第一步,在训练图结构内添加伪量化节点quant_aware_train

# 记录预训练模型的所以变量

# 存储预训练模型的权值

# 增加全局变量集合

# 在重写图之前记录

per_trained_model_path = '/models/float_point/model.ckpt'

restore_dict = {}

reader = tf.train.NewCheckpointReader(per_trained_model_path)

for v in tf.global_variables():

tensor_name = v.name.split(':')[0]

if reader.has_tensor(tensor_name):

restore_dict[tensor_name] = v

# 重写图, 向训练图中加入伪量化操作

experimental_create_training_graph(input_graph=sess.graph,weight_bits=8,

一般是在loss之后optimizer之前添加create_training_graph关键函数,其将自动的帮我们在可识别的操作上嵌入伪量化节点,训练并保存模型后,模型图结构就会自动的存在伪量化节点及其统计的参数。tf.contrib.quantize.create_training_graph的参数input_graph表示训练的默认图层,quant_delay是指多少次迭代之后再进行量化,如果是已训练好进行微调量化的话,那么可以将quant_delay设为0。 activation_bits=8)。

# 保存伪量化训练后的模型

saver = tf.train.Saver()

saver.save(sess, './models/quant_aware_trained/model.ckpt')

tensorboard ckpt pb 模型的输出节点_网络模型压缩量化(mnist的CNN举例)
(2)第二步,重写推理图结构并保存为新的模型(freeze)

#为了模拟量化就地重绘 training input_graph。

experimental_create_eval_graph(input_graph=sess.graph, weight_bits=8, #固化图

const_graph = tf.graph_util.convert_variables_to_constants(

sess=sess,

input_graph_def=sess.graph.as_graph_def(),

output_node_names=['output_tensor/BiasAdd'])

# 序列化并将模型写成pb文件

with tf.gfile.GFile('./models/frozen.pb', "wb") as f:

f.write(const_graph.SerializeToString())

activation_bits=8)

推理和训练的伪量化图结构是存在较大差异的,而该操作正是要消除量化操作对batch normalization的影响,

tensorboard ckpt pb 模型的输出节点_网络模型压缩量化(mnist的CNN举例)
(3)第三步,转换模型为全量化模型convert_to_lite

利用tensorflow.lite.python.lite构建转换器,将pb文件转化成tflite文件。

converter = lite.TFLiteConverter.from_frozen_graph(

graph_def_file='/models/frozen.pb',

input_arrays=['input_tensor'],

output_arrays=['output_tensor/BiasAdd'],

input_shapes={'input_tensor': [1, 28, 28, 1]})

# 将pb文件转换成tflite文件

tflite_quantized_model = converter.convert()

with open('models/model.tflite', 'wb') as f:

f.write(tflite_quantized_model)

tensorboard ckpt pb 模型的输出节点_网络模型压缩量化(mnist的CNN举例)

至此,我们将得到一个经过Tensorflow量化感知训练后的全量化模型,

从每个步骤保存的文件,我们可以看出加入伪量化节点后保存的模型和原来的模型大小一样,约为14M,经过freeze后保存的pb文件为4.5M左右,减少为之前的25%左右,在经过lite转换变成tflite文件只有1.1M左右,又变成之前的1/4左右,并且精度几乎保持不变,在99%左右,实现了模型压缩的目的。

四、总结

代码主要使用keras完成,相比tensorflow,更加简洁,与tensorflow差别不大。其中要注意,不同版本会使用不同的函数和lite工具,本代码使用keras基于tf1.14版本。

第一步:在quant_aware_train.py的loss后,训练操作节点前加入tf.contrib.quantize experimental_create_training_graph,训练模型,保存ckpt文件

第二步:在freeze.py构建推理图,加tf.contrib.quantize.experimental_create_eval_graph,restore之前的保存的ckpt文件,冻结生成pb文件

第三步:利用tensorflow.lite.python.lite构建convert来将pb文件转换成tflite文件,变成其1/4大小。

虽然模型大小降低了许多,但是精准度还是需要进行测试来保证,之后还要会写一个基于flask框架、部署在本地的web应用,来比较最初保存的模型和量化后的模型在识别写字板上手写数字的识别的精准度。

引用自tensorflow中文社区中大佬的理解:

量化有两个途径:1.伪量化。具体展开:要生成两个graph,

一个用于training

,在compute gradients前使用create_training_graph,因为forward和backward都需要模拟量化。这个过程其实是找到需要量化的变量,插入fake quantization nodes 来记录 min-max range information。

另一个graph用于eval

,要在import_graph_def后,saver.restore前插入create_eval_graph。后面再freeze和tflite格式转换。

2. post-training quantization。无需伪量化,是toco拿过来一个pb文件,直接把toco的post-training quantization flag设成true完成的。但是注意,这个方法只quantize weights,而不是一个fully quantized model,并且在inference时,会把uint8的weights再转换回float32来做矩阵乘法。所以,这个方法其实依然相当于没做quantization。

它的存在应当只是便于用户转移model,计算时依然是个非量化model。

代码见于:

outman123/tensorflow​github.com

tensorboard ckpt pb 模型的输出节点_网络模型压缩量化(mnist的CNN举例)

参考:

tensorflowLite的量化使用问题

https://blog.csdn.net/lishanlu136/article/details/88872266

继续阅读