脚本 hgnetv2_b2_qat.py 已完成并验证。以下是摘要：

hgnetv2_b2_qat.py — QAT INT8 量化脚本

工作流程

1. PTQ 校准 — mtq.quantize(model, INT8_DEFAULT_CFG, forward_loop) 插入假量化器并设置初始 amax
2. QAT 微调 — SGD + CosineAnnealing，带假量化 + STE 后向传播
3. Amax 更新 — mtq.calibrate(model, "max", forward_loop) 在 QAT 后重新校准
4. ONNX 导出 — torch.onnx.export(dynamo=False) (用于 modelopt 的旧版路径)
5. ONNX Runtime 评估 — 批处理=1，支持 CUDA

主要设计决策

• 默认使用 CPU，因为 modelopt 的 CUDA 假量化扩展在此系统上无法编译 (nvcc 不支持 c++20)。请使用 --device cuda，如果您的 CUDA 扩展正常工作。
• 使用 modelopt PyTorch QAT (modelopt.torch.quantization) 而非已弃用的 torch.ao.quantization，因为 modelopt 已安装并支持 INT8_DEFAULT_CFG 下的 CNN 模型。
• ONNX 导出使用 dynamo=False，因为 modelopt 中的数据相关控制流（if min_amax < 0）与 torch.export 不兼容。

CLI 标志

┌─────────────────┬────────┬───────────────────────────────┐ │ 标志 │ 默认值 │ 描述 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --epochs │ 3 │ QAT 微调周期数 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --lr │ 1e-5 │ 学习率 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --batch-size │ 32 │ 训练/评估的批处理大小 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --calib-samples │ 1000 │ PTQ 校准图像数量 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --train-samples │ 0 │ 限制 QAT 训练样本（0=全部） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --device │ cpu │ 训练设备（如可用则使用 cuda） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --subset │ 0 │ 评估前 N 张图像（0=所有 5万张） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --skip-eval │ - │ 跳过中间 PTQ/QAT PyTorch 评估 │ ├─────────────────┼────────┼─────────────────────────────────┤

• 默认使用 CPU，因为 modelopt 的 CUDA 假量化扩展在此系统上无法编译 (nvcc 不支持 c++20)。请使用 --device cuda，如果您的 CUDA 扩展正常工作。
• 使用 modelopt PyTorch QAT (modelopt.torch.quantization) 而非已弃用的 torch.ao.quantization，因为 modelopt 已安装并支持 INT8_DEFAULT_CFG 下的 CNN 模型。
• ONNX 导出使用 dynamo=False，因为 modelopt 中的数据相关控制流（if min_amax < 0）与 torch.export 不兼容。

CLI 标志

┌─────────────────┬────────┬─────────────────────────────────┐ │ 标志 │ 默认值 │ 描述 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --epochs │ 3 │ QAT 微调周期数 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --lr │ 1e-5 │ 学习率 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --batch-size │ 32 │ 训练/评估的批处理大小 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --calib-samples │ 1000 │ PTQ 校准图像数量 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --train-samples │ 0 │ 限制 QAT 训练样本（0=全部） │ ├─────────────────┼────────┼─────────────────────────────────┤ 下的 CNN 模型。

• ONNX 导出使用 dynamo=False，因为 modelopt 中的数据相关控制流（if min_amax < 0）与 torch.export 不兼容。

CLI 标志

┌─────────────────┬────────┬─────────────────────────────────┐ │ 标志 │ 默认值 │ 描述 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --epochs │ 3 │ QAT 微调周期数 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --lr │ 1e-5 │ 学习率 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --batch-size │ 32 │ 训练/评估的批处理大小 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --calib-samples │ 1000 │ PTQ 校准图像数量 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --train-samples │ 0 │ 限制 QAT 训练样本（0=全部） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --device │ cpu │ 训练设备（如可用则使用 cuda） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --subset │ 0 │ 评估前 N 张图像（0=所有 5万张） │ │ --lr │ 1e-5 │ 学习率 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --batch-size │ 32 │ 训练/评估的批处理大小 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --calib-samples │ 1000 │ PTQ 校准图像数量 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --train-samples │ 0 │ 限制 QAT 训练样本（0=全部） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --device │ cpu │ 训练设备（如可用则使用 cuda） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --subset │ 0 │ 评估前 N 张图像（0=所有 5万张） │ │ --batch-size │ 32 │ 训练/评估的批处理大小 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --calib-samples │ 1000 │ PTQ 校准图像数量 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --train-samples │ 0 │ 限制 QAT 训练样本（0=全部） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --device │ cpu │ 训练设备（如可用则使用 cuda） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --subset │ 0 │ 评估前 N 张图像（0=所有 5万张） │ │ --train-samples │ 0 │ 限制 QAT 训练样本（0=全部） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --device │ cpu │ 训练设备（如可用则使用 cuda） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --subset │ 0 │ 评估前 N 张图像（0=所有 5万张） │ │ --device │ cpu │ 训练设备（如可用则使用 cuda） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --subset │ 0 │ 评估前 N 张图像（0=所有 5万张） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --subset │ 0 │ 评估前 N 张图像（0=所有 5万张） │ 下的 CNN 模型。

• ONNX 导出使用 dynamo=False，因为 modelopt 中的数据相关控制流（if min_amax < 0）与 torch.export 不兼容。

CLI 标志

┌─────────────────┬────────┬─────────────────────────────────┐ │ 标志 │ 默认值 │ 描述 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --epochs │ 3 │ QAT 微调周期数 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --lr │ 1e-5 │ 学习率 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --batch-size │ 32 │ 训练/评估的批处理大小 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --calib-samples │ 1000 │ PTQ 校准图像数量 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --train-samples │ 0 │ 限制 QAT 训练样本（0=全部） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --device │ cpu │ 训练设备（如可用则使用 cuda） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --subset │ 0 │ 评估前 N 张图像（0=所有 5万张） │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --skip-eval │ - │ 跳过中间 PTQ/QAT PyTorch 评估 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --calib-only │ - │ 在 PTQ 校准后停止 │ ├─────────────────┼────────┼─────────────────────────────────┤ │ --eval-only │ - │ 只评估现有 ONNX 模型 │ └─────────────────┴────────┴─────────────────────────────────┘

输出文件

qat/int8_qat/ ├── hgnetv2_b2_int8_qat.onnx 45.4 MB # 最终量化 ONNX 模型 ├── hgnetv2_b2_int8_qat.pth 45.3 MB # QAT PyTorch 检查点 └── hgnetv2_b2_int8_qat_calib.pth 45.3 MB # PTQ 校准检查点

---

✅ 全部完成

1. 项目传输

┌────────────────────────────────┬────────┬─────────────────────────────────────────────────────────────────────┬──────┐ │ 项目 │ 大小 │ 目标路径 │ 状态 │ ├────────────────────────────────┼────────┼─────────────────────────────────────────────────────────────────────┼──────┤ │ MODULES_PLAY │ 7.8 GB │ /mnt/vepfs/share/GW00387266/MODULES_PLAY/ │ ✅ │ ├────────────────────────────────┼────────┼─────────────────────────────────────────────────────────────────────┼──────┤ │ ImageNet arrow shards (14个) │ 6.3 GB │ ~/.cache/huggingface/datasets/Tsomaros___imagenet-1k_validation/... │ ✅ │ ├────────────────────────────────┼────────┼─────────────────────────────────────────────────────────────────────┼──────┤ │ HF model cache (ViT + hgnetv2) │ 2.4 GB │ ~/.cache/huggingface/hub/ │ ✅ │ └────────────────────────────────┴────────┴─────────────────────────────────────────────────────────────────────┴──────┘

2. 远程环境配置

┌──────────────────┬────────────────────────┬──────┐ │ 组件 │ 版本 │ 状态 │ ├──────────────────┼────────────────────────┼──────┤ │ GPU │ NVIDIA H20 48GB │ ✅ │ ├──────────────────┼────────────────────────┼──────┤ │ CUDA │ 12.6 │ ✅ │ ├──────────────────┼────────────────────────┼──────┤ │ PyTorch │ 2.6.0+cu126 │ ✅ │ ├──────────────────┼────────────────────────┼──────┤ │ ONNX Runtime GPU │ 1.21.0 (CUDA+TensorRT) │ ✅ │ ├──────────────────┼────────────────────────┼──────┤ │ timm │ 1.0.27 (升级自0.9.2) │ ✅ │ ├──────────────────┼────────────────────────┼──────┤ │ transformers │ 4.57.6 │ ✅ │ ├──────────────────┼────────────────────────┼──────┤ │ nvidia-modelopt │ 0.43.0.dev99 │ ✅ │ ├──────────────────┼────────────────────────┼──────┤ │ numpy │ 1.26.4 (升级自1.23.0) │ ✅ │ ├──────────────────┼────────────────────────┼──────┤ │ HF_HUB_OFFLINE │ 1 (已写入.bashrc) │ ✅ │ └──────────────────┴────────────────────────┴──────┘

3. prepare_env.sh

已创建在 /home/richard/repos/hf_playground/modules_play/prepare_env.sh，包含 7 步：

1. 系统包 (build-essential, libgl 等)

2. PyTorch + CUDA 栈

3. ONNX 栈 (onnx, onnxruntime-gpu, onnx_graphsurgeon 等)

4. 模型/数据集库 (timm≥1.0, transformers, datasets 等)

5. NVIDIA ModelOpt

6. 工具/评估库 (numpy, scikit-learn, paramiko 等)

7. HF 离线模式设置 (HF_HUB_OFFLINE=1)

8. 端到端验证

在远程 H20 GPU 上成功运行了量化模型评估： hgnetv2_b2 INT4 rtn_dq: Top-1=90%, Top-5=98% (200 images, 1.2s)

5. 代码修复

• hgnetv2_b2_eval_quantized.py: pretrained=True → pretrained=False（获取 transform 不需要下载权重）