Microsoft NNI模型剪枝快速入门指南

Microsoft NNI模型剪枝快速入门指南

nni An open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning. 项目地址: https://gitcode.com/gh_mirrors/nn/nni

什么是模型剪枝

模型剪枝是一种通过减少模型权重或中间状态大小来降低模型计算量和存储需求的技术。在深度神经网络中，许多参数对最终输出的贡献很小，剪枝就是识别并移除这些冗余参数的过程。

Microsoft NNI（Neural Network Intelligence）提供了强大的模型剪枝功能，支持多种剪枝策略：

1. 预训练模型 → 剪枝 → 微调剪枝后模型
2. 训练过程中剪枝（剪枝感知训练）→ 微调剪枝后模型
3. 直接剪枝 → 从头训练剪枝后模型

准备工作

在开始剪枝前，我们需要准备一个基础模型。本教程使用一个简单的卷积神经网络在MNIST数据集上进行预训练。

import torch
import torch.nn.functional as F
from torch.optim import SGD
from nni_assets.compression.mnist_model import TorchModel, trainer, evaluator, device

# 初始化模型
model = TorchModel().to(device)
print(model)  # 查看模型结构

# 定义优化器和损失函数
optimizer = SGD(model.parameters(), 1e-2)
criterion = F.nll_loss

# 预训练3个epoch
for epoch in range(3):
    trainer(model, optimizer, criterion)
    evaluator(model)

模型结构输出显示这是一个典型的CNN，包含卷积层、全连接层和ReLU激活函数。预训练后模型在测试集上达到约95%的准确率。

配置剪枝策略

NNI提供了多种剪枝算法，这里我们使用L1NormPruner，它基于权重的L1范数进行剪枝。

首先需要定义config_list，指定剪枝的规则：

config_list = [{
    'op_types': ['Linear', 'Conv2d'],  # 剪枝线性层和卷积层
    'exclude_op_names': ['fc3'],       # 排除名为fc3的层
    'sparse_ratio': 0.5                # 目标稀疏率为50%
}]

这个配置表示：

• 对所有Linear和Conv2d类型的层进行剪枝
• 排除名为'fc3'的最后一层（通常输出层不剪枝）
• 每层的稀疏率目标为50%

执行剪枝

初始化剪枝器并应用剪枝：

from nni.compression.pruning import L1NormPruner

pruner = L1NormPruner(model, config_list)
print(model)  # 查看被包装后的模型

# 生成剪枝掩码
_, masks = pruner.compress()

# 查看各层实际稀疏率
for name, mask in masks.items():
    print(name, 'sparsity:', '{:.2}'.format(mask['weight'].sum() / mask['weight'].numel()))

此时模型尚未真正变小，只是标记了哪些权重应该被保留。可以看到各层的稀疏率接近配置的50%。

模型加速

要使模型真正变小，需要进行模型加速：

pruner.unwrap_model()  # 解包模型

from nni.compression.speedup import ModelSpeedup

# 使用随机输入加速模型
dummy_input = torch.rand(3, 1, 28, 28).to(device)
ModelSpeedup(model, dummy_input, masks).speedup_model()

print(model)  # 查看加速后的模型

加速后模型结构显示各层通道数明显减少：

• conv1从6通道减到3通道
• conv2从16通道减到8通道
• 全连接层维度也相应减小

微调剪枝后模型

剪枝后的模型性能通常会下降，需要通过微调恢复准确率：

# 重新初始化优化器（因为模型参数已改变）
optimizer = SGD(model.parameters(), 1e-2)

# 微调3个epoch
for epoch in range(3):
    trainer(model, optimizer, criterion)

剪枝效果分析

通过这个过程，我们实现了：

1. 模型参数大幅减少（约50%稀疏率）
2. 计算量降低（各层通道数减少）
3. 通过微调保持模型准确率

实际应用中，可以根据需求调整稀疏率和微调epoch数，在模型大小和性能间取得平衡。

进阶建议

1. 尝试不同的剪枝算法（如FPGMPruner、TaylorFOWeightFilterPruner等）
2. 使用渐进式剪枝策略，逐步提高稀疏率
3. 结合量化技术进一步压缩模型
4. 在更复杂的数据集（如CIFAR-10/100）上验证效果

NNI提供了丰富的剪枝和模型压缩工具，可以帮助开发者高效地优化神经网络模型。

nni An open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning. 项目地址: https://gitcode.com/gh_mirrors/nn/nni