APQ：剪枝 + 量化 自动搜索与微调流程说明

本目录实现了针对 CNN（MobileNet/MobileNetV2/ResNet 等）的 剪枝 + 量化（k-means / 4Q-SAT）自动搜索与微调。
下面梳理从预训练模型开始，到搜索策略、导出模型、修改 config、微调与评估的完整流程。

---

1. 代码结构速览

• 顶层脚本（统一版本 - 推荐使用）

  • search.py：统一搜索脚本（整合 search_kmeans.py / search_4Q.py 功能，通过 --quant_method 切换）
  • finetune.py：统一微调脚本（整合 finetune_linearquant.py / finetune_SAT.py 功能，通过 --quant_backend 切换）
  • eval.py：对给定模型进行评估
  • prune_helper.py：根据剪枝后的通道数，自动更新模型 config（mobilenetv2.py / resent_p.py）

• 顶层脚本（原始版本 - 保留用于参考）

  • search_kmeans.py：剪枝 + k-means 量化 的 RL 搜索
  • search_4Q.py：剪枝 + 4Q（权重/激活分开）量化 的 RL 搜索
  • finetune_linearquant.py：线性量化微调
  • finetune_SAT.py：SAT 量化微调

• 库代码

  • lib/env/：剪枝+量化环境定义（pruning_quantize_env*.py、rewards.py）
  • lib/rl/：RL 算法（DDPG 等）
  • lib/utils/：量化模块、数据加载、通用工具
  • models/：各类模型和 profile（剪枝配置）

---

2. 快速开始（推荐使用统一脚本）

2.1 使用统一搜索脚本 search.py

优势：一个脚本支持所有量化方法（k-means / SAT / 线性量化），参数统一，易于对比实验。

基本用法：

# k-means 量化搜索
python search.py --quant_method kmeans --job train --model mobilenet ...

# SAT/4Q 量化搜索
python search.py --quant_method sat --job train --model mobilenetv2 ...

# 导出模型
python search.py --quant_method sat --job export --ratios "1.0,0.75,..." ...

详细文档：参见 SEARCH_USAGE.md

2.2 使用统一微调脚本 finetune.py

优势：一个脚本支持所有量化后端（linear / sat / pwlq），参数统一。

基本用法：

# 线性量化微调
python finetune.py --quant_backend linear --model mobilenet_0.5flops ...

# SAT 量化微调
python finetune.py --quant_backend sat --model mobilenetv2_0.8flops ...

详细文档：参见 FINETUNE_USAGE.md

2.3 完整工作流（推荐）

# 1. 搜索最优策略
python search.py --quant_method sat --job train --model mobilenetv2 \
  --dataset imagenet --preserve_ratio_p 0.8 --preserve_ratio_q 0.5 ...

# 2. 导出剪枝模型
python search.py --quant_method sat --job export \
  --ratios "1.0,0.75,..." --export_path exported.pth.tar ...

# 3. 更新模型配置
python prune_helper.py --model mobilenetv2 --profile 0.8flops \
  --channels "[3, 24, 16, ...]"

# 4. 微调
python finetune.py --quant_backend sat --model mobilenetv2_0.8flops \
  --ckpt_path exported.pth.tar --quant_strategy 8 8 6 6 ...

# 5. 评估
python eval.py --model mobilenetv2_0.8flops --ckpt_path best.pth.tar ...

---

3. 原始脚本使用说明（保留用于参考）

以下是原始 search_kmeans.py 和 search_4Q.py 的使用说明。推荐使用上面的统一脚本 search.py。

3.1 k-means 量化搜索流程（search_kmeans.py）

3.1.1 搜索（train）

以 MobileNet + ImageNet 为例：

cd APQ

python search_kmeans.py \
  --job train \
  --model mobilenet \
  --dataset imagenet \
  --dataset_root /path/to/imagenet \
  --ckpt_path /path/to/mobilenet_imagenet_pretrain.pth.tar \
  --preserve_ratio_p 0.5 \
  --preserve_ratio_q 0.5 \
  --min_bit 1 --max_bit 8 \
  --train_episode 600 \
  --gpu_id 0

搜索过程中：

• 环境 Pruning_Quatize_Env（k-means 版）给出状态和 reward
• DDPG 输出剪枝动作 action_p 和量化动作 action_q
• env.best_strategy_p / best_strategy_q 记录了当前最优的剪枝/量化策略

3.1.2 导出剪枝+量化模型（export）

导出时只需要 确定剪枝 ratios 或 channels，脚本内部会根据它们生成最终模型：

python search_kmeans.py \
  --job export \
  --model mobilenet \
  --dataset imagenet \
  --dataset_root /path/to/imagenet \
  --ckpt_path /path/to/mobilenet_imagenet_pretrain.pth.tar \
  --ratios "1.0,0.75,0.75,0.6875,..." \
  --export_path ../../save/APQ/export/mobilenet_imagenet_kmeans_pruned.pth.tar \
  --gpu_id 0

export_model 内部会打印：

• Original model channels
• Pruning with ratios
• Channels after pruning（剪枝后的通道数，可用于更新模型 config）

---

3.2 4Q / SAT 量化搜索流程（search_4Q.py）

3.2.1 搜索（train）

4Q 搜索同时考虑：

• 剪枝（preserve_ratio_p）
• 权重量化 bit（action_qw）
• 激活量化 bit（action_qa）

示例命令：

python search_4Q.py \
  --job train \
  --model mobilenetv2 \
  --dataset imagenet \
  --dataset_root /path/to/imagenet \
  --ckpt_path /path/to/mobilenetv2_imagenet_pretrain.pth.tar \
  --preserve_ratio_p 0.8 \
  --preserve_ratio_q 0.5 \
  --min_bit 2 --max_bit 6 \
  --quant_type SAT \
  --train_episode 300 \
  --gpu_id 0

训练结束后：

• env.best_strategy_p：每层剪枝保留率
• env.best_strategy_q：每层 [w_bit, a_bit]（权重/激活 bit 配置）

3.2.2 导出剪枝模型（export）

python search_4Q.py \
  --job export \
  --model mobilenetv2 \
  --dataset imagenet \
  --dataset_root /path/to/imagenet \
  --ckpt_path /path/to/mobilenetv2_imagenet_pretrain.pth.tar \
  --ratios "1.0,0.75,1.0,1.0,..." \
  --export_path ../../save/APQ/export/mobilenetv2_imagenet_4Q_pruned.pth.tar \
  --gpu_id 0

同样会打印 Channels after pruning，可用于后续修改 mobilenetv2 的 pruned_channels。

---

4. 使用 prune_helper.py 自动更新模型 config

在导出阶段，search_*.py 会打印形如：

=> Channels after pruning: [3, 24, 16, 96, 24, 104, ..., 304, 144]

你可以直接用这串通道数更新 models/ 中的配置，而无需手动改代码。

4.1 更新 MobileNetV2 的 pruned_channels

假设针对 ImageNet 的 4Q 搜索得到了完整通道列表，并希望更新 profile='0.8flops'：

cd APQ

python prune_helper.py \
  --model mobilenetv2 \
  --profile 0.8flops \
  --channels "[3, 24, 16, 96, 24, 104, 24, 112, 32, 184, 32, 176, 32, 184, 64, 288, 64, 344, 64, 280, 64, 376, 96, 528, 96, 408, 96, 472, 160, 616, 160, 840, 160, 904, 304, 144]"

脚本会在 models/mobilenetv2.py 中找到：

• elif profile == '0.8flops': 对应分支
• 替换其中的 pruned_channels = [...]。

4.2 更新 ResNet 的 cfg_18 / cfg_50

如果有一串新的通道列表想用于 ResNet-18/50：

# 更新 ResNet-50 的 cfg_50
python prune_helper.py \
  --model resnet50 \
  --channels "[32, 32, 40, 32, 256, 32, 24, 256, ...]"

# 更新 ResNet-18 的 cfg_18
python prune_helper.py \
  --model resnet18 \
  --channels "[64, 24, 64, 24, 64, 56, ...]"

脚本会修改 models/resent_p.py 中对应的 cfg_50 / cfg_18。

---

5. 统一微调脚本 finetune.py

finetune.py 合并了 finetune_linearquant.py 和 finetune_SAT.py 的功能，通过 --quant_backend 选择量化后端：

• linear：标准线性/均匀量化（对应原 quantize_modules）
• sat / pwlq：SAT / PWLQ 量化（对应 quantize_modules_PWLQ）

5.1 线性量化微调（替代 finetune_linearquant.py）

python finetune.py \
  --quant_backend linear \
  --model mobilenet_0.5flops \
  --dataset cifar10 \
  --data_root /path/to/data \
  --ckpt_path ../../save/APQ/export/mobilenet_cifar10_test.pth.tar \
  --quant_strategy 8 -1 8 8 8 8 8 7 7 8 8 8 8 7 8 7 8 7 8 8 8 8 8 8 8 1 8 8 \
  --lr 0.05 \
  --n_epoch 150 \
  --lr_type exp \
  --batch_size 128 \
  --n_worker 8

5.2 SAT / 4Q 量化微调（替代 finetune_SAT.py）

python finetune.py \
  --quant_backend sat \
  --model mobilenetv2_0.8flops \
  --dataset imagenet \
  --data_root /path/to/imagenet \
  --ckpt_path ../../save/APQ/export/mobilenetv2_imagenet_4Q_pruned.pth.tar \
  --quant_strategy 8 8 6 6 6 6 5 6 6 5 ... \
  --lr 0.1 \
  --n_epoch 50 \
  --lr_type cos \
  --batch_size 256 \
  --n_worker 32 \
  --n_gpu 2

说明：

• quant_strategy 是扁平列表：每两个数字为一层的 [w_bit, a_bit]
• 长度需与可量化层数一致，脚本会自动检查

5.3 仅评估导出模型

python finetune.py \
  --quant_backend sat \
  --model mobilenetv2_0.8flops \
  --dataset imagenet \
  --data_root /path/to/imagenet \
  --ckpt_path ../../save/APQ/export/mobilenetv2_imagenet_4Q_pruned.pth.tar \
  --eval

---

6. 整体推荐工作流总结

以 MobileNetV2 + ImageNet + 4Q 为例：

1. 搜索剪枝 + 量化策略

   • 运行 search_4Q.py --job train ...
   • 在 log / TensorBoard 中查看 best_strategy_p（剪枝）和 best_strategy_q（bit）

2. 导出剪枝模型

   • 运行 search_4Q.py --job export --ratios "..." --export_path ...
   • 记录输出中的 Channels after pruning: [...]

3. 更新模型 config

   • 使用 prune_helper.py，把 Channels after pruning 写入 mobilenetv2.py 对应 profile 的 pruned_channels

4. 根据策略微调

   • 使用 finetune.py --quant_backend sat，传入：
     • --model mobilenetv2_0.8flops
     • --ckpt_path 指向导出的 pruned ckpt
     • --quant_strategy 由 best_strategy_q 展开而来

5. 评估与对比

   • 用 finetune.py --eval 或 eval.py 对最终模型精度和 FLOPs、参数量进行评估。

Acknowledgement

We would like to thank Prof. Bing Li for providing guidance and maintaining related research resources available at https://github.com/libingict. This work is built upon and inspired by prior open-source efforts from the MIT HAN Lab, in particular the following repositories:

• AMC: https://github.com/mit-han-lab/amc
• HAQ: https://github.com/mit-han-lab/haq

We gratefully acknowledge the authors and contributors of these projects for making their code publicly available.