环境要求

  1. 支持TensorFlow1.X，在TensorFlow1.X最后一个版本1.15.3上测试通过；TensorFlow2.0暂未测试

    总体使用流程

  2. 通过本工具得到训练好的float pb模型，以及量化参数文件

  3. 通过MNN转换工具，输入这两个文件，得到最终的量化模型

    支持的op，使用建议

  4. 建议从训练好的float模型进行finetune

  5. 目前支持 Conv2D, DepthwiseConv2dNative，带参数MatMul（EMA，LSQ量化算法支持）
  6. 使用OAQ量化算法需要量化矩阵乘 MatMul ，请将其利用 reshape 和 Conv2D（1 * 1 卷积） 组合实现
  7. 优化器超参使用模型收敛阶段的超参，学习率可以适当调节一下

  8. 模型的第一层或者前面几层对精度影响较大，可以尝试跳过不进行量化

    支持的训练量化算法

  9. EMAQuantizer: naive QAT算法，scale采用滑动平均更新

  10. LSQQuantizer: 改进QAT算法，scale会在网络训练中学习更新

  11. OAQQuantizer: 改进LSQ算法，int8计算中累加到int16而不是通常的累加到int32，在低端手机上可以获得比普通量化30%~50%的加速，低端设备上加速更明显

    使用方法

    从checkpoint finetune

  12. 搭建好网络之后，在session运行之前，构造 Quantizer 对象，指定相关配置项，构造此对象的同时会向前面搭建好的网络中插入fake quant节点，用于实现训练量化功能，示例代码： ```python from mnncompress.tensorflow import EMAQuantizer, strip_MNN_QAT_ops, get_qat_state_placeholder, MNN_QAT_variables from mnncompress.tensorflow import LSQQuantizer, strip_MNN_QAT_ops, get_qat_state_placeholder, MNN_QAT_variables from mnncompress.tensorflow import OAQQuantizer, strip_MNN_QAT_ops, get_qat_state_placeholder, MNN_QAT_variables QATQuantizer = EMAQuantizer # 也可以改为使用LSQQuantizer或者OAQQuantizer

first, 构建前向网络模型

build_model()

second, 插入fake quant节点，注意此句位置应在构建optimizer之前

更多配置详见API部分

quantizer = QATQuantizer(tf.get_default_graph())

构建optimizer，反向网络

opt = …

只恢复原模型中的变量

saver = tf.train.Saver(var_list=[v for v in tf.all_variables() if v not in MNN_QAT_variables()])

sess = tf.Session()

初始化所有变量，包括MNN训练量化中用到的变量

sess.run(tf.global_variables_initializer())

给原模型中的变量用checkpoint赋值

saver.restore(sess, ‘save/model.ckpt’)

注意：该saver只维护了原来的模型的变量，要保存插入fake quant的模型，需要新建一个saver

new_saver = tf.train.Saver()

  2. 2. 正常训练，注意，训练和测试过程中需要给quantizer内部的 **MNN_QAT_is_training** 这个 **placeholder** 赋值，训练时赋值True，测试时赋值False；**注意添加第1行、第7行代码：**
  3. ```python
  4. quantizer.init(sess)
  5. quant_state = quantizer.is_training # 或者 quant_state = get_qat_state_placeholder()
  7. for data in dataset:
  8.     feed_dict = {x: batch_x, labels: batch_y, quant_state: True}
  9.     sess.run(train_step, feed_dict=feed_dict)
  10.     quantizer.update(sess)

  1. 测试过程中保存训练得到的量化参数文件，示例代码：

    1. # 如果先进行了剪枝训练，那么需要将剪枝中得到的 "compress_params.bin" 文件在下方传入，并设置 append=True
    2. quantizer.save_compress_params("compress_params.bin", sess, append=False)

  2. 训练完之后，我们得到了一些checkpoint文件和对应的量化参数文件，现在我们要来得到最终推理用的frozen pb了，示例代码： ```python from mnncompress.tensorflow import EMAQuantizer, strip_MNN_QAT_ops

保存frozen pb之前去掉（绕过）插入的量化节点，传入 strip_MNN_QAT_ops 的graph参数

也可以是从checkpoint读取出来的graph

strip_MNN_QAT_ops(sess.graph)

output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=output_names) f = open(“frozen.pb”, ‘wb’) f.write(output_graph_def.SerializeToString())

  2. 5. 使用MNN模型转换工具，输入frozen pb和量化参数文件，得到最终的量化模型，示例：
  3. ```shell
  4. mnnconvert --modelFile frozen.pb  --MNNModel quant_model.mnn --framework TF --bizCode AliNNTest --compressionParamsFile compress_params.bin

测试训练量化结果需注意的地方

如果你是训练和测试分开两个文件，测试时用python代码重新构图，在构图时给bn的is_training又直接赋值为python bool False，然后加载训练的checkpoint进行测试，那你重新构建的图里面是没有训练量化节点的，此时需要在你构图之后，创建session之前，恢复变量之前，你再新建一个和训练时候一样的Quantizer对象（复制粘贴），向测试图中插入训练量化节点即可，然后根据训练时候保存的checkpoint恢复训练得到的变量，就可以了，当然测试的时候也需要给 quant_state 赋值为False，即可。

另外注意，有的使用者在训练和测试的时候给模型加了不同的scope，这时训练的到的compress_params.bin就不能用了，因为tensor名字对不上，此时，只需要在测试之后转pb之前加入下面两句就可以了：

  1. quantizer.save_compress_params("compress_params.bin", sess)
  2. strip_MNN_QAT_ops(sess.graph)
  4. output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=output_names)
  5. f = open("frozen.pb", 'wb')
  6. f.write(output_graph_def.SerializeToString())

相关API

EMAQuantizer/OAQQuantizer/LSQQuantizer

  1. Quantizer(graph=None, is_training = None, skip_quant_layers = [], skip_quant_flag = [], bits = 8, 
  2.           min_clamp_value = 0, # 此参数仅OAQQuantizer有
  3.           debug_info = False, retain_sparsity=False)
  1. 参数
  3. graph: 需要进行训练量化的模型图
  4. is_training: None or Python bool，为None时，Quantizer的is_training属性是一个placeholder，用于控制量化状态，需要在feed_dict中额外传值；
  5.                             为True或False时，则和用户代码中的训练或测试flag一样，is_training是一个constant，不需要在feed_dict中额外赋值
  6. skip_quant_layers: [str], 指定图中跳过量化的op的名字
  7. skip_quant_flag: [str]，跳过op名字中含有指定字符串的op的量化，例如想要跳过模型中整个decoder1的量化，这些op的名字中都含有"decoder1"，那么在此参数中指定["decoder1"]即可
  8. bits: int, 指定量化比特数，目前仅支持8bit量化
  9. min_clamp_value: int, OAQQuantizer特有参数，用于指定最小的量化范围，例如最低量化为6bit，则此数值设置为2^(6-1) - 1 = 31
  10. debug_info: bool，是否输出debug信息
  12. retain_sparsity: bool, 是否保持量化原始float模型的稀疏结构不变，剪枝之后叠加量化时，需要将此参数设为True

为获得相关层的名字，我们提供了一个帮助函数：

  1. import mnncompress.tensorflow.helpers as helper
  3. helper.get_op_name_of_type(graph, ['Conv2D', 'DepthwiseConv2dNative', 'MatMul'])
  1. 方法和属性
  3. is_training: property，返回Quantizer内部的 MNN_QAT_is_training 这个 placeholder
  5. strip_qat_ops(): 保存froze pb前去掉（绕过）插入的训练量化节点
  7. save_compress_params(filename, sess, append=False): 保存量化参数信息到filename文件，sess为维护训练量化图的session, 
  8.                                                     append 表示是否将量化参数信息append到传入的proto文件中，如果为false，则将新建或覆盖该文件，
  9.                                                     append=True 表示量化之前还做了其他模型压缩操作，如剪枝，filename文件中已存在相关压缩信息
  10. init(sess): 初始化Quantizer内部信息
  11. update(sess): 更新内部需要的信息

strip_MNN_QAT_ops

  1. strip_MNN_QAT_ops(graph)
  1. 此函数用于保存froze pb前去掉（绕过）插入的训练量化节点
  3. 如果训练和测试分开来做，意味着测试时是读取训练得到的checkpoint然后进行测试的
  5. 此时图中已经有了训练量化的fake quant节点，但无法使用Quantizer维护的原图信息进行
  6. 保存froze pb前去掉（绕过）插入的训练量化节点 的功能
  8. 此时可以使用该函数来达到去掉（绕过）fake quant节点的目的

get_qat_state_placeholder

  1. get_qat_state_placeholder(graph=None)
  1. 用于在无quantizer对象的情况下获取MNN训练量化工具内部标志训练或者测试的placeholder

MNN_QAT_variables

  1. MNN_QAT_variables()
  1. 用于获取MNN训练量化特有的量化variable集合