环境要求

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

    大致算法原理

    使用权值mask梯度信息，权值mask梯度的绝对值为权值的重要性表示，全局对所有权值的重要性大小进行排序，将指定比例的重要性较小的权值剪掉，参考论文 SNIP: Single-shot Network Pruning based on Connection Sensitivity

    支持的op，使用建议

  2. 目前支持Conv2D，带可学习参数MatMul需要使用reshape+1*1 Conv2D卷积的组合来实现

    支持的剪枝算法

  3. SNIPLevelPruner：最细粒度的随机剪枝算法，稀疏单位为单个权值。一般需要剪枝比例达到80%以上才有加速，此方法主要进行压缩。

  4. SIMDOCPruner：稀疏单位为1*4的剪枝算法，一般需要30%以上剪枝比例才能加速，精度比通道剪枝好。

  5. TaylorFOChannelPruner：通道剪枝算法，以卷积filter为单位进行剪枝，剪枝完仍然是整齐的模型，无须框架底层特殊实现，但需要用MNN转换工具进行转换才能获得最终剪枝模型。

    使用方法

  6. 前提：一个预先训练好的模型checkpoint（也可以从头开始训练稀疏模型）

  7. 首先恢复出已经训练好的checkpoint，示例代码： ```python import tensorflow as tf from mnncompress.tensorflow import SNIPLevelPruner, SIMDOCPruner, TaylorFOChannelPruner from mnncompress.tensorflow import SensitivityAnalyzer

构图

…

恢复checkpoint

sess = tf.Session() sess.run(tf.global_variables_initializer()) saver = tf.train.Saver() saver.restore(sess, ‘original_model/model.ckpt’)

  2. 3. 获取网络中的（gradient，variable）pair
  3. ```python
  4. # 模型中用到的优化器，例如，optimizer = tf.train.AdamOptimizer(1e-4)
  5. optimizer = ...
  7. # 获取（gradient，variable）pair
  8. gradients_vars = optimizer.compute_gradients(loss)
  10. # 模型更新op
  11. train_step = optimizer.apply_gradients(gradients_vars, global_step=tf.train.get_or_create_global_step())
  1. 计算模型中各网络连接的敏感度，并保存为npy文件。主要计算过程是喂几个batch数据进网络，计算得到各参数的梯度： ```python from mnncompress.tensorflow import SensitivityAnalyzer

analyzer = SensitivityAnalyzer(gradients_vars) for i in range(10): batch_x, batch_y = mnist.train.next_batch(50) batch_x = np.reshape(batch_x,(-1, 28, 28,1))

  1. feed_dict={x_PH: batch_x, labels: batch_y}
  3. # 计算各参数的梯度
  4. gv_numpy = sess.run(gradients_vars, feed_dict=feed_dict)
  5. analyzer.analyze(gv_numpy)

analyzer.save_sensitivity_npy(“model_sens.npy”)

  2. 5. 构建Pruner，示例代码：
  3. ```python
  4. target_sparsity = 0.6
  5. total_prune_iterations = 1
  7. # 可选SNIPLevelPruner, SIMDOCPruner, TaylorFOChannelPruner等算法，更多配置参数参见下方api文档
  8. pruner = SNIPLevelPruner("model_sens.npy", sess.graph, target_sparsity, total_prune_iterations, debug_info=True)

  1. 正常训练，注意，训练时每次反向更新之后，执行一下剪枝步骤即可（第6行）（测试时不需要），这样保存下来的模型就已经是稀疏的了，示例代码：

    1. for data in training_dataset:
    2.  feed_dict = {x: batch_x, labels: batch_y}
    3.  sess.run([train_step], feed_dict=feed_dict)
    5.  # 反向更新之后，执行剪枝，测试时不需要
    6.  pruner.do_pruning(sess)

  2. 测试过程中保存一下剪枝相关的参数文件：

    1. pruner.save_compress_params("compress_params.bin", sess)

  3. 训练完成保存frozen pb即可得到最终稀疏pb模型，示例代码：

    1. output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=output_names)
    2. f = open("frozen.pb", 'wb')
    3. f.write(output_graph_def.SerializeToString())

  4. 模型稀疏之后，可以进一步使用TF训练量化工具进行量化，得到稀疏量化模型。也可以直接使用此float稀疏的模型进行推理，需要在MNN模型转换时指定第6步中的bin文件，并进行权值量化（才会进行稀疏编码）：

    1. mnnconvert --modelFile frozen.pb  --MNNModel quant_model.mnn --framework TF --bizCode AliNNTest --compressionParams compress_params.bin --weightQuantBits 8

    相关API

    SNIPLevelPruner/SIMDOCPruner/TaylorFOChannelPruner

    1. Pruner(sensitivity_npy_file, graph=None, sparsity=0.5, total_pruning_iterations=1, config_file=None, debug_info= False, 
    2.     prune_finetune_iterations=0, max_prune_ratio=0.99, 
    3.     prune_ratio_result_file="算法名_found_prune_ratios.yml", 
    4.     align_channels=4 # TaylorFOChannelPruner特有参数
    5.    )

    ```python 参数

sensitivity_npy_file: 敏感度分析得到的npy文件

graph: tensorflow计算图，如为None，则默认使用tf.get_default_graph()来获取计算图

sparsity: float，0~1， 模型的总体剪枝比例

total_pruning_iterations: int， 多少个iteration内将指定比例的参数稀疏掉

config_file: str，可先运行一次Pruner的do_pruning方法，得到搜索出来的剪枝比例yml文件， 然后微调此yml文件中的各层剪枝比例，将修改后的文件名通过此参数传入，可控制各层剪枝比例

debug_info: bool，是否输出debug信息

prune_finetune_iterations: int，剪枝一次，finetune此参数指定步数，然后再剪枝一次。 如果total_pruning_iterations设为10，prune_finetune_iterations设为99，则1000个iteration之后才执行完毕剪枝过程

max_prune_ratio: 0~1，控制各层最大的剪枝比例

prune_ratio_result_file: str，指定搜索得到的剪枝比例保存的文件名

align_channels: int，指定通道剪枝之后对齐的通道数

  1. ```python
  2. 方法和属性
  4. do_pruning(sess): 训练时每更新一次模型参数之后，执行剪枝步骤，测试时不需要
  6. save_compress_params(filename, sess, append=False): 保存剪枝参数信息到filename文件，sess为维护训练量化图的session, 
  7.                                                     append 表示是否将剪枝参数信息append到传入的proto文件中，如果为false，则将新建或覆盖该文件，
  8.                                                     append=True 表示剪枝之前还做了其他模型压缩操作，filename文件中已存在相关压缩信息