剪枝框架推荐第二弹，nni：不止是剪枝压缩工具

日了狗，csdn这帮搬运狗，搬运了也不标来源

首发 towardsdeeplearning 微信公众号和知乎专栏。

随着深度学习的发展，模型也越来越大。虽然精确度有了很大的提升，但较大的模型尺寸不仅会影响推理速度，还对部署的硬件要求较高。因此，模型压缩也是一个热门话题。

剪枝框架推荐第二弹，今天介绍下微软的开源的工具包：nni。

NNI (Neural Network Intelligence) 是一个轻量但强大的工具包，帮助用户自动的进行特征工程，神经网络架构搜索，超参调优以及模型压缩。

NNI 管理自动机器学习 (AutoML) 的 Experiment，调度运行由调优算法生成的 Trial 任务来找到最好的神经网络架构和/或超参，支持各种训练环境，如本机，远程服务器，OpenPAI，Kubeflow，基于 K8S 的 FrameworkController（如，AKS 等)， DLWorkspace (又称 DLTS), AML (Azure Machine Learning) 以及其它环境。

那么nni能在哪些场景下使用呢？官方说明了在如下的场景：

• 想要在自己的代码、模型中试验不同的自动机器学习算法。 想要在不同的环境中加速运行自动机器学习。
• 想要更容易实现或试验新的自动机器学习算法的研究员或数据科学家，包括：超参调优算法，神经网络搜索算法以及模型压缩算法。
• 在机器学习平台中支持自动机器学习。 NNI 提供命令行工具以及友好的 WebUI 来管理训练的 Experiment。 通过可扩展的 API，可定制自动机器学习算法和训练平台。 为了方便新用户，NNI 内置了最新的自动机器学习算法，并为流行的训练平台提供了开箱即用的支持。

回到剪枝算法，主流的模型压缩算法可分为两类，一类是压缩网络结构的剪枝算法，另一类是减小模型权重位数的量化算法。NNI 目前提供了 AGP、L1Filter、Slim、Lottery Ticket、FPGM、QAT、DoReFa 等9种模型压缩算法。用户也可根据需要，通过 NNI 的模型压缩接口来实现自己的压缩算法。

相关算法介绍及用法：https://aka.ms/nnizc

看下使用方法，我们以 FPGM Pruner 为例，在 PyTorch 代码中，我们可以通过config_list 来配置剪枝率和相关信息。

from nni.compression.torch import FPGMPruner
config_list = [{
    'sparsity': 0.5,
    'op_types': ['Conv2d']
}]
pruner = FPGMPruner(model, config_list)
pruner.compress()

然后在用户配置文件里配置下剪枝的类，FPGM Pruner，即可

PyTorch
..  autoclass:: nni.compression.torch.FPGMPruner

代码相对简单，还支持tensorflow和pytorch框架，基本能覆盖日常学习和工作的需求。

说下我对这个框架的看法，一个工具能否流行和良好的社区和交流环境有很大关系，nni在这方面相对做的好一些，开源于2018年，背靠微软，可能被抛弃的可能性和维护的积极性会高些。