Skip to content

SEARCH_USAGE.md

Latest commit

 

History

History
317 lines (253 loc) · 8.92 KB

SEARCH_USAGE.md

File metadata and controls

317 lines (253 loc) · 8.92 KB
Raw

统一 Search 脚本使用说明

概述

search.py 是一个统一的搜索脚本，整合了 k-means 量化、SAT/4Q 量化和线性量化三种方法，取代了原来的 search_kmeans.py 和 search_4Q.py。

主要特性

  • 统一接口：通过 --quant_method 参数切换量化方式
  • 自动配置：根据量化方法自动选择对应的环境和 DDPG 版本
  • 支持多种模型：MobileNet、MobileNetV2、ResNet、VGG 等
  • 完整流程：支持搜索（train）和导出（export）两个阶段

基本用法

1. k-means 量化搜索（替代 search_kmeans.py）

python search.py \
  --quant_method kmeans \
  --job train \
  --model mobilenet \
  --dataset imagenet \
  --dataset_root /path/to/imagenet \
  --ckpt_path /path/to/mobilenet_imagenet.pth.tar \
  --preserve_ratio_p 0.5 \
  --preserve_ratio_q 0.5 \
  --min_bit 1 \
  --max_bit 8 \
  --train_episode 600 \
  --warmup 100 \
  --gpu_id 0

2. SAT/4Q 量化搜索（替代 search_4Q.py）

python search.py \
  --quant_method sat \
  --job train \
  --model mobilenetv2 \
  --dataset imagenet \
  --dataset_root /path/to/imagenet \
  --ckpt_path /path/to/mobilenetv2_imagenet.pth.tar \
  --preserve_ratio_p 0.8 \
  --preserve_ratio_q 0.5 \
  --min_bit 2 \
  --max_bit 6 \
  --quant_type SAT \
  --train_episode 300 \
  --warmup 20 \
  --bsize 128 \
  --gpu_id 0

3. 线性量化搜索

python search.py \
  --quant_method linear \
  --job train \
  --model mobilenet \
  --dataset imagenet \
  --dataset_root /path/to/imagenet \
  --ckpt_path /path/to/mobilenet_imagenet.pth.tar \
  --preserve_ratio_p 0.5 \
  --preserve_ratio_q 0.5 \
  --train_episode 600 \
  --gpu_id 0

4. 导出剪枝后的模型

搜索完成后，使用找到的最优剪枝策略导出模型：

python search.py \
  --quant_method sat \
  --job export \
  --model mobilenetv2 \
  --dataset imagenet \
  --dataset_root /path/to/imagenet \
  --ckpt_path /path/to/mobilenetv2_imagenet.pth.tar \
  --ratios "1.0,0.75,1.0,0.88,..." \
  --export_path ../../save/APQ/export/mobilenetv2_pruned.pth.tar \
  --gpu_id 0

参数说明

核心参数

  • --quant_method: 量化方法选择

    • kmeans: k-means 聚类量化（动作空间：2 维）
    • sat 或 4q: SAT/PWLQ 量化（动作空间：3 维，分别控制剪枝、权重量化、激活量化）
    • linear: 标准线性量化（动作空间：2 维）

  • --job: 任务类型

    • train: RL 搜索最优剪枝和量化策略
    • export: 使用给定策略导出剪枝后的模型

模型与数据

  • --model: 模型名称
    • mobilenet, mobilenetv2
    • vgg
    • resnet18, resnet50
  • --dataset: 数据集（cifar10 / imagenet）
  • --dataset_root: 数据集路径
  • --ckpt_path: 预训练模型路径

剪枝参数

  • --preserve_ratio_p: 剪枝目标保留率（0-1）
  • --lbound: 每层最小保留率
  • --rbound: 每层最大保留率

量化参数

  • --preserve_ratio_q: 量化目标保留率
  • --min_bit: 最小量化比特数
  • --max_bit: 最大量化比特数
  • --quant_type: 量化类型（LinearQuant / SAT / TanhQuant）
  • --float_bit: 全精度浮点比特数（默认 32）

RL 训练参数

  • --train_episode: 训练轮数
    • k-means 推荐：600
    • SAT/4Q 推荐：300
  • --warmup: warmup 轮数（随机探索阶段）
    • k-means 推荐：100
    • SAT/4Q 推荐：20
  • --bsize: DDPG 训练的 batch size
    • k-means 推荐：64
    • SAT/4Q 推荐：128
  • --rmsize: 每层的 replay buffer 大小
  • --lr_c: Critic 学习率（默认 1e-3）
  • --lr_a: Actor 学习率（默认 1e-4）

导出参数

  • --ratios: 每层剪枝保留率，逗号分隔（例如："1.0,0.75,0.5,..."）
  • --channels: 剪枝后的通道数，逗号分隔（二选一，优先使用 ratios）
  • --export_path: 导出模型的保存路径

其他参数

  • --gpu_id: GPU 设备 ID
  • --seed: 随机种子
  • --suffix: 实验后缀名（用于区分不同实验）
  • --use_new_input: 使用新的输入特征

量化方法对比

方法 动作空间 DDPG 版本 环境 适用场景
kmeans 2 维 (p, q) ddpg_2 env_kmeans k-means 聚类量化
sat / 4q 3 维 (p, qw, qa) ddpg_3 env_SAT 权重+激活独立量化
linear 2 维 (p, q) ddpg_2 env (基础) 标准线性量化

完整工作流示例

场景：MobileNetV2 + ImageNet + SAT/4Q 量化

步骤 1：搜索最优策略

python search.py \
  --quant_method sat \
  --job train \
  --model mobilenetv2 \
  --dataset imagenet \
  --dataset_root /data/imagenet \
  --ckpt_path /ckpt/mobilenetv2_imagenet.pth.tar \
  --preserve_ratio_p 0.8 \
  --preserve_ratio_q 0.5 \
  --min_bit 2 \
  --max_bit 6 \
  --train_episode 300 \
  --suffix exp1 \
  --gpu_id 0

输出：

  • 日志目录：../../save/APQ/mobilenetv2_imagenet_rp0.8_rq0.5_sat_search_exp1/
  • 最优策略会记录在 log.txt 中：
    • best prune policy: [...]
    • best quant policy: [...]

步骤 2：导出剪枝模型

从日志中找到 best prune policy（例如 [1.0, 0.75, 1.0, ...]），用于导出：

python search.py \
  --quant_method sat \
  --job export \
  --model mobilenetv2 \
  --dataset imagenet \
  --dataset_root /data/imagenet \
  --ckpt_path /ckpt/mobilenetv2_imagenet.pth.tar \
  --ratios "1.0,0.75,1.0,0.88,1.0,..." \
  --export_path ../../save/APQ/export/mobilenetv2_0.8flops_pruned.pth.tar \
  --gpu_id 0

输出：

  • 导出的剪枝模型：mobilenetv2_0.8flops_pruned.pth.tar
  • 终端会打印： 
    => Channels after pruning: [3, 24, 16, 96, 24, 104, ...]
=> Copy-paste for config update:
[3, 24, 16, 96, 24, 104, ...]

步骤 3：更新模型配置

使用 prune_helper.py 更新 mobilenetv2.py 中的 pruned_channels：

python prune_helper.py \
  --model mobilenetv2 \
  --profile 0.8flops \
  --channels "[3, 24, 16, 96, 24, 104, ...]"

步骤 4：微调

使用统一的 finetune.py 微调：

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

步骤 5：评估

python eval.py \
  --model mobilenetv2_0.8flops \
  --dataset imagenet \
  --data_root /data/imagenet \
  --ckpt_path ./logs/mobilenetv2_0.8flops_imagenet_finetune_sat-runXX/ckpt.best.pth.tar

与原脚本的对应关系

原脚本 新脚本使用方式
search_kmeans.py search.py --quant_method kmeans
search_4Q.py search.py --quant_method sat 或 --quant_method 4q

常见问题

Q: 如何选择合适的 quant_method？
A:

  • 如果你想用 k-means 聚类量化，选 kmeans
  • 如果你想权重和激活分别量化（更精细），选 sat 或 4q
  • 如果你想简单的线性量化，选 linear

Q: 搜索需要多长时间？
A:

  • k-means（600 episodes）：通常需要几个小时到一天
  • SAT/4Q（300 episodes）：通常需要数小时
  • 具体取决于模型大小、数据集和 GPU 性能

Q: 导出时 ratios 和 channels 的区别？
A:

  • ratios：每层的保留比例（0-1），会自动计算 channels
  • channels：直接指定每层剪枝后的通道数
  • 两者二选一，优先使用 ratios

Q: 如何恢复中断的搜索？
A: 目前脚本会在训练过程中自动保存中间模型（每 1/3 episodes），但暂不支持从中断点恢复。建议：

  1. 适当减少 train_episode
  2. 使用 --suffix 区分不同实验
  3. 监控 tensorboard 实时查看进度

Q: 为什么 SAT/4Q 的 warmup 比 k-means 少？
A: SAT/4Q 的动作空间是 3 维，探索空间更大，因此使用较少的 warmup 可以更快进入有效训练。k-means 是 2 维动作空间，需要更多 warmup 来充分探索。

TensorBoard 监控

训练过程中可以使用 TensorBoard 实时监控：

tensorboard --logdir ../../save/APQ/

关键指标：

  • reward/last: 最近一次 episode 的奖励
  • reward/best: 历史最佳奖励
  • info/accuracy_p: 剪枝后精度
  • info/accuracy_q: 量化后精度
  • info/compress_ratio: 压缩率
  • preserve_rate_prune/X: 每层剪枝保留率
  • preserve_rate_w/X, preserve_rate_a/X: 每层权重/激活量化比特（仅 SAT/4Q）

技巧和建议

  1. 从小模型开始：先在 CIFAR-10 上验证流程，再转到 ImageNet
  2. 合理设置目标：preserve_ratio_p 和 preserve_ratio_q 不要设得太低，建议从 0.5-0.8 开始
  3. 保存日志：每次实验都用 --suffix 标记，便于后续对比
  4. 监控训练：通过 TensorBoard 实时查看 reward 曲线，如果长时间不收敛可以提前停止
  5. GPU 显存：如果显存不足，可以减小 --data_bsize