环境要求

  1. python3

  2. PyTorch >= 1.8.0

    总体使用流程

  3. 通过本工具得到训练好的float onnx模型，以及MNN模型压缩参数文件

  4. 通过MNN转换工具，输入这两个文件，得到最终的MNN稀疏模型（如果直接部署稀疏的float模型，而不叠加量化，那么转换时需要使用MNNConvert的 —weightQuantBits 8 参数进行转换，才会进行稀疏编码，否则模型大小将不变）

    支持的op，使用建议

  5. 目前支持torch.nn.Conv2d，torch.nn.Linear

  6. 优化器超参使用模型收敛阶段的超参，学习率可以在收敛阶段学习率的基础上再调小

    支持的剪枝算法

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

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

  9. TaylorFOChannelPruner：通道剪枝算法，此算法是将conv2d的一整个filter剪掉，因此此filter对应的输出通道会失效，剪枝完模型是一个规整的模型，不需要后端进行特殊加速，但为了将剪掉的filter从模型中真正去掉，需要使用MNNConverter进行转换，转换过程中会分析剪枝通道之间的依赖关系，并完成最终的转换。

    使用方法

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

  11. 创建Pruner对象，对原始模型进行转换，然后用转换之后的模型进行训练
  12. 在train过程中，调用pruner的do_pruning方法进行剪枝
  13. 导出MNN模型压缩参数文件，示例代码如下（关注其中pruner的用法）： ```python from mnncompress.pytorch.SNIP_level_pruner import SNIPLevelPruner from mnncompress.pytorch.SIMD_OC_pruner import SIMDOCPruner Pruner = SIMDOCPruner

你的模型代码

class Net(nn.Module): pass

model = Net()

加载已经训练好的模型

model.load_state_dict(torch.load(“ori_model.pt”)) model.to(device)

将模型进行转换，并使用转换后的模型进行训练，测试

更多配置请看API部分

pruner = SIMDOCPruner(model, total_pruning_iterations=1, sparsity=0.6, debug_info=False)

def train(model, data, optimizer, pruner) model.train() for d, t in data: optimizer.zero_grad() output = model(d) loss = F.nll_loss(output, t) loss.backward() optimizer.step()

  1.     # step之后调用pruner的剪枝方法
  2.     pruner.do_pruning()
  5. # 获取当前剪枝比例
  6. print(pruner.current_prune_ratios())

for epoch in range(1, epochs + 1): train(model, data, optimizer, pruner) test(model, data)

保存模型

model.eval() torch.save(model.state_dict(), “pruned_model.pt”) x = torch.randn(input_shape).to(device) torch.onnx.export(model, x, “pruned_model.onnx”)

保存MNN模型压缩参数文件，应在剪枝完毕之后进行，建议在保存模型时调用

pruner.save_compress_params(“compress_params.bin”, append=False)

  2. 5. 模型稀疏之后，可以进一步使用PyTorch训练量化工具进行量化，得到稀疏量化模型。也可以直接使用此float稀疏的模型进行推理，需要在MNN模型转换时指定剪枝得到的MNN模型压缩参数文件，并进行权值量化（才会进行稀疏编码）：
  3. ```bash
  4. mnnconvert --modelFile pruned_model.onnx  --MNNModel pruned_model.mnn --framework ONNX --bizCode MNNTest --compressionParamsFile compress_params.bin --weightQuantBits 8

相关API

SNIPLevelPruner / SIMDOCPruner / TaylorFOChannelPruner

  1. Pruner(model, sparsity=0.5, total_pruning_iterations=1, config_file=None, debug_info= False, 
  2.        prune_finetune_iterations=0, max_prune_ratio=0.99, 
  3.        align_channels=4  # only for 'TaylorFOChannelPruner'
  4.       ):
  1. 参数
  3. model：Module，pytorch模型
  5. sparsity：float，0~1，模型总体剪枝比例
  7. total_pruning_iterations：int，指定多少个iteration将指定比例的连接剪掉
  9. config_file：str，可先运行一次Pruner的do_pruning方法，得到搜索出来的剪枝比例yml文件之后，
  10.                   然后微调此yml文件中的各层剪枝比例，将修改后的文件名通过此参数传入，可控制各层剪枝比例
  12. debug_info：是否输出debug信息
  14. prune_finetune_iterations: int, 指定剪枝进行一步，finetune多少步，如指定为99，则剪枝一次，
  15.                                                                 finetune99次，在finetune的99次过程中剪枝mask不变
  17. max_prune_ratio: float, 指定各层最大剪枝比例
  19. align_channels: int, 仅通道剪枝算法TaylorFOChannelPruner有此参数，用于指定剪枝之后对齐的通道数
  1. 方法和属性
  3. do_pruning(result_file = "found_prune_ratios.yml")：进行一次剪枝，在训练optimizer.step之后调用，搜索到的各层剪枝比例将保存到此文件中
  5. current_prune_ratios()：返回当前各层剪枝比例
  7. save_compress_params(filename, append=False): 保存MNN模型压缩参数到filename文件
  8.                                               append 表示是否将剪枝参数信息append到传入的文件中，如果为false，则将新建或覆盖该文件，
  9.                                               append=True 表示剪枝之前还做了其他模型压缩操作，filename文件中已存在相关压缩信息