模型剪枝 - 3.5 简体中文

迭代剪枝

此方法包含两个阶段：模型分析和剪枝后的模型生成。完成模型分析后，分析结果会保存在名为 .vai/xxx.sens 的文件内。您可使用此文件对模型进行迭代剪枝。在迭代剪枝中，需逐渐对模型进行剪枝以达到目标稀疏度。这是通过使用迭代循环来达成的，此循环由剪枝步骤和微调步骤组成，并按步骤使用适度的剪枝率。如果尝试设置过高的剪枝率，则会导致步骤中精度损失且无法恢复。

1. 定义评估函数。该函数必须取一个模型作为其首个实参，并返回得分。

   def eval_fn(model, dataloader):
     top1 = AverageMeter('Acc@1', ':6.2f')
     model.eval()
     with torch.no_grad():
     for i, (images, targets) in enumerate(dataloader):
       images = images.cuda()
       targets = targets.cuda()
       outputs = model(images)
       acc1, _ = accuracy(outputs, targets, topk=(1, 5))
       top1.update(acc1[0], images.size(0))
     return top1.avg

2. 运行模型分析并获取剪枝后的模型。

   runner.ana(eval_fn, args=(val_loader,))
   
   model = pruning_runner.prune(removal_ratio=0.2)

模型分析只需执行一次即可。您可对模型进行迭代剪枝，而无需重新运行分析，因为针对特定剪枝率仅生成一个剪枝后的模型。通过剪枝获取的子网络一般都不够好，因为近似算法会根据分析结果来生成此独特剪枝模型。单步剪枝方法可以生成更好的子网络。

单步剪枝

此方法包含两个阶段：基于自适应 BN 搜索剪枝策略以及剪枝后的模型生成。搜索后，会生成名为 .vai/xxx.search 的文件，用于存储搜索结果（剪枝策略和对应的评估得分）。您可单步获取最终剪枝后模型。

num_subnet 可提供满足稀疏性要求（待确认）的候选子网络的目标数量。您可从这些候选结果中选择最佳子网络。该值越高，搜索耗时越长，但找到更好的子网络的可能性更高。

# Adaptive-BN-based searching for pruning strategy. 'calibration_fn' is a function for calibrating BN layer's statistics.
runner.search(gpus=['0'], calibration_fn=calibration_fn, calib_args=(val_loader,), eval_fn=eval_fn, eval_args=(val_loader,), num_subnet=1000, removal_ratio=0.7)

model = runner.prune(removal_ratio=0.7, index=None)

eval_fn 与迭代剪枝方法相同。以下代码示例中显示了用于实现自适应 BN 的 calibration_fn 函数。您应采用类似方式来定义自己的代码。

def calibration_fn(model, dataloader, number_forward=100):
  model.train()
  with torch.no_grad():
    for index, (images, target) in enumerate(dataloader):
      images = images.cuda()
      model(images)
    if index > number_forward:
      break

相比于迭代方法，单步剪枝方法具有几项优势：

• 生成的剪枝后模型通常更准确。满足要求的所有子网络都会进行评估。
• 工作流程更为简单，因为您可单步获取最终剪枝后的模型，无需迭代。
• 重新训练精简模型比重新训练稀疏模型更快。

单步剪枝有两大劣势：其一是随机生成的剪枝策略不可预测。另一项是必须对每个剪枝比率执行一次子网络搜索。