ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

Hessian感知的混合精度后训练量化:学习型图像压缩部署实战

Hessian感知的混合精度后训练量化:学习型图像压缩部署实战 做学习型图像压缩Learned Image Compression, LIC模型部署时相信不少同学遇到过同样的尴尬模型在 GPU 上率失真表现很漂亮一旦要转成端侧或服务端推理模型内存和时延双双告急模型体积也偏大。压缩模型本身就是“算力换码率”网络里的卷积、注意力、超先验模块一个比一个重实际业务场景不可能一直用 FP32 全精度推理。于是量化就成了一个绕不开的话题但直接套用通用量化方案LIC 模型的码率波动和重建质量下降往往很严重。本文围绕一篇典型的量化方向工作HAMP-LICHessian-Aware Mixed-Precision Post-Training Quantization for Learned Image Compression展开梳理它背后的核心思路如何用 Hessian 信息衡量 LIC 模型中不同层对量化的敏感度又如何基于敏感度做混合精度比特分配以及最终怎么落到 Post-Training Quantization后训练量化的工程流程里。本文适合三类读者正在做图像/视频压缩模型部署的算法工程师。想了解混合精度量化原理、但不清楚如何衡量层敏感度的同学。准备复现或改进量化压缩论文的研究者。读完你会掌握Hessian 感知敏感度分析的基本原理、混合精度比特分配的常用搜索思路、以及一套可以照着改的 PyTorch 代码骨架。1. 背景与动机学习型图像压缩模型为何需要量化1.1 学习型图像压缩LIC是什么先对齐一下概念。传统图像压缩基于手工设计的变换、量化和熵编码例如 JPEG、WebP、HEVC/H.265、AV1 等。而学习型图像压缩是用神经网络替代上述手工模块典型结构是“自编码器 超先验Hyperprior 熵模型 上下文模型”。一个经典的 LIC 模型会包含主编码器/主解码器完成图像特征的提取和重建。超编码器/超解码器对主特征的标准差scale进行估计生成熵模型参数。上下文模型利用已解码的邻域信息进一步优化概率估计。熵编码模块基于概率分布做算术编码。论文里常提到的 Minnen、Ballé、Cheng 等模型核心结构都围绕这套框架展开。这类模型把压缩问题变成了“率失真优化”问题训练目标是[ L R \lambda \cdot D ]其中 (R) 是估计码率(D) 是重建失真如 MSE 或 MS-SSIM(\lambda) 控制两者的平衡。1.2 部署中的精度与存储矛盾虽然 LIC 模型在压缩效率上不断刷新纪录但部署时问题很直接模型参数动辄几十 MB端侧存储压力大。编码器/解码器需要跑多次卷积和注意力计算时延高。超先验模块中的自回归上下文计算是串行的影响端到端吞吐。量化是把浮点权重和激活从 FP32 转为低位宽表示如 INT8、INT4、甚至 INT2从而减少模型体积、降低内存带宽、加速推理。可一旦把 LIC 模型整体压到 INT8 以下码率和重建质量会明显退化原因是压缩模型对特征的统计特性非常敏感尤其是超先验分支里的参数一旦精度损失概率估计不准熵编码的码率就会异常放大。1.3 HAMP-LIC 解决什么问题HAMP-LIC 的核心出发点可以用一句话概括不是所有层都需要高精度也不是所有层都能承受低位宽。用 Hessian 信息判断每一层“对量化有多敏感”再据此分配不同的量化比特数让整体精度损失在满足模型体积目标的前提下尽可能小。这是一个典型的混合精度后训练量化问题。它不像 QATQuantization-Aware Training那样需要重新训练整个模型而是基于少量校准数据分析模型各层对量化的敏感度然后搜索一组最优的每层比特配置。由于不用端到端重训它在部署加速场景下更实用。2. 后训练量化PTQ基础2.1 量化与反量化先回顾最基础的量化公式。把一个浮点张量 (x) 量化到 (b) bit 整数[ q \text{round}\left( \frac{x}{s} z \right) ]反量化为[ \hat{x} s \cdot (q - z) ]其中(s) 是缩放因子scale由浮点范围和量化范围决定。(z) 是零点zero point用于对齐浮点零和量化整数零。round 操作引入了量化误差。量化误差的来源主要是两个舍入误差和截断误差。前者由 round 导致后者由浮点范围超出量化可表示范围导致。范围截断和精度损失之间存在权衡范围取大了截断少但舍入误差大范围取小了舍入误差小但容易截断。2.2 PTQ 与 QAT 的对比维度PTQ后训练量化QAT量化感知训练是否需要重训不需要或仅需少量微调需要完整训练/微调流程校准数据量通常几百到几千张即可依赖训练集成本高部署成本低适合快速上线高需要训练资源和时间精度表现高位宽时接近原模型低位宽时通常更稳灵活性适合混合精度快速搜索每跑一组精度都要重训代价大HAMP-LIC 选择 PTQ正是因为 LIC 模型训练成本高收敛慢不希望为了部署再去重训一遍。通过 PTQ 加少量校准数据就能完成从 FP32 到混合精度整型模型的转换。2.3 为什么选择 PTQ在真实业务里模型迭代往往很快压缩模型的超参数例如 (\lambda)一变网络结构可能不变但权重分布就变了。如果每次修改都要做 QAT整个迭代周期会被拉得很长。PTQ 的价值在于只用少量图片校准。不需要反向传播重新优化整个网络。可以快速评估“当前模型量化到多少 bit 可接受”。当然PTQ 在极低位宽如 2-bit 权重下精度下降会更明显但混合精度可以缓解这个问题敏感层给 8-bit不敏感层给 4-bit 甚至 2-bit整体平均 bit 控制在目标范围附近。3. Hessian-Aware 敏感度分析3.1 为什么不同层需要不同精度混合精度的前提是网络各层对量化的承受能力不同。先举个例子。假设有两个模块模块 A 输出经过强归一化权重数值范围很集中量化误差会被后续结构吸收模块 B 的输出直接参与熵参数估计微小的偏移会导致概率分布偏斜进而影响码率。对模块 A 用 4-bit 可能完全没有问题对模块 B 用 4-bit 则可能让码率暴涨 20% 以上。所以统一给所有层分配相同 bit 数要么浪费存储要么牺牲精度。那么问题来了怎么量化“敏感度”直观想法是直接看量化前后损失函数的变化量。比如[ \Delta L L(\hat{\theta}) - L(\theta) ]其中 (\hat{\theta}) 是量化后的权重。但直接计算需要逐层尝试不同 bit 组合组合空间呈指数爆炸。因此需要一个更高效的敏感度度量。3.2 Hessian 与量化误差的关系从泰勒展开看损失函数关于权重的二阶近似是[ L(\theta \delta) \approx L(\theta) g^T \delta \frac{1}{2} \delta^T H \delta ]其中 (g) 是梯度(H) 是 Hessian 矩阵。在模型已经收敛的情况下梯度 (g) 接近 0那么量化带来的损失变化主要由二阶项决定[ \Delta L \approx \frac{1}{2} \delta^T H \delta ]所以Hessian 矩阵反映了权重扰动量化误差对损失的放大程度。Hessian 越大说明该层对参数扰动越敏感需要更高 bit 的量化。在混合精度量化中Hessian 信息有两种常见用法用 Hessian 的迹Trace表示整体敏感度。用 Hessian 对角线Diagonal逐参数估计敏感度。用 Hessian 与量化误差的二次型 ( \delta^T H \delta ) 估计该层量化后的损失增量。对于 LIC 模型损失函数是率失真损失也就是[ L R \lambda \cdot D ]对量化噪声更敏感的层其扰动会同时影响码率项和失真项。HAMP-LIC 的做法就是通过计算率失真损失关于各层权重的 Hessian找出“哪些层量化后会让 R 和 D 同时恶化的更多”。3.3 实际工程中的 Hessian 近似直接计算完整 Hessian 矩阵不现实。一个 3×3 卷积如果输入输出通道都是 256权重张量大小为 (3\times3\times256\times256)即约 59 万参数Hessian 会是 59 万 × 59 万的矩阵无法存储。所以工程上通常会做近似Hessian 对角线近似忽略参数之间的相关性只保留对角线元素。Fisher 信息矩阵近似用梯度的外积期望近似 Hessian即 (F \approx \frac{1}{N}\sum \nabla L \nabla L^T)。Fisher 信息在某些条件下与 Hessian 期望一致。Kronecker 分解把 Hessian 近似成两个小矩阵的 Kronecker 积例如 KFAC。HAMP-LIC 这类方法里通常采用“少量校准数据 多次反向传播”来估计 Hessian 对角线或 Fisher 信息。核心代码思路后面会详细给出。这里要特别注意Hessian 的估计需要模型处于合适的模式如 BatchNorm 要用校准集统计校准数据的分布要和实际业务数据尽量一致否则敏感度排序会失真。4. 混合精度量化与比特分配4.1 统一精度 vs 混合精度统一精度Uniform Precision就是所有层都用同一种 bit 宽度比如全部 INT8 或全部 INT4。实现简单部署框架支持度高但不是最优解。混合精度Mixed Precision允许每层使用不同 bit 数。例如编码器第一层卷积INT8超编码器中间层INT6解码器输出层INT4上下文模型INT8目标是在总模型大小或平均 bit 数约束下最小化量化后的率失真损失增量。混合精度的难点不在“概念”而在“搜索”。假设网络有 50 个量化单元每个单元可选 ({2,4,6,8}) 四种 bit组合数是 (4^{50})暴力枚举不可能。4.2 比特分配的搜索策略常见的搜索策略有几种贪心搜索Greedy Search从所有层都是最高精度开始每一步尝试把某一层降低一档 bit选择“损失增量除以 bit 节省量”最小的层重复直到满足体积约束。优点是实现简单效果好缺点是可能陷于局部最优。动态规划DP把比特分配建模成背包问题。每层看作一个物品每个 bit 档位看作该物品的一种“体积-收益”组合用背包 DP 求最优解。优点是能找到全局最优缺点是当网络层数多、候选 bit 档位多时状态数很大。启发式搜索Evolution / Bayesian Optimization适合搜索空间极大、目标函数非凸的情况但耗时较长。HAMP-LIC 这类方法更偏向用“敏感度指标 贪心/DP”的组合策略。先通过 Hessian 信息给每层打敏感度分然后在一个统一约束下做最优分配。4.3 HAMP-LIC 的整体流程整体流程可以拆成以下步骤准备校准时数据集通常是从验证集或训练集中抽样的几百张图像。加载 FP32 预训练 LIC 模型。对每一层或每一个量化单元计算 Hessian 信息。计算每层在候选 bit 档位下的量化误差估计。构建目标函数最小化率失真损失增量约束是平均 bit 数或模型大小。执行贪心/DP 搜索得到每层最优 bit 配置。按配置完成 PTQ输出量化模型。用测试集验证码率和重建质量的退化情况必要时对少数层做修正。5. 代码实现与配置示例下面用 PyTorch 给出一个核心思路的实现骨架。注意这是为了帮助理解算法流程写的示例代码不是 HAMP-LIC 的官方实现实际复现时需要结合你使用的 LIC 模型结构和 PyTorch 版本调整。5.1 Hessian 对角线估计假设我们有一个 LIC 模型输入是一批图像输出是率失真损失。我们想估计某个卷积层权重的 Hessian 对角线。这里用 Fisher 近似对每个校准样本计算损失对权重的梯度然后对梯度的平方求平均。这比直接调用torch.autograd.functional.hessian更高效也更适合大模型。# 文件路径hessian_estimate.py import torch from torch.autograd import grad def estimate_fisher_diag(model, calib_loader, loss_fn, target_module): 估计某个模块权重的 Fisher 对角线Hessian 的常用近似。 参数: model: FP32 LIC 模型 calib_loader: 校准数据 DataLoader loss_fn: 返回率失真损失的函数入参为 (model, images) target_module: 要分析的 nn.Module 返回: diag: 与 target_module.weight 形状相同的张量表示逐参数敏感度 model.eval() diag None count 0 for images in calib_loader: images images.cuda() loss loss_fn(model, images) # 只对目标模块权重的梯度感兴趣 grads grad(loss, target_module.weight, retain_graphFalse, create_graphFalse) g grads[0] g2 g.detach() ** 2 if diag is None: diag torch.zeros_like(g2) diag g2 count 1 diag diag / max(count, 1) return diag这里用梯度的平方平均近似二阶信息本质是 Fisher 信息。严格来说 Fisher 信息与 Hessian 的期望并不完全相等但很多混合精度量化工作里Fisher 对角线被证明是 Hessian 对角线非常有效的替代计算成本低得多。如果你确实想看到更接近 Hessian 对角线的高阶信息可以尝试torch.autograd.functional.hessian但要注意该 API 对显存和计算时间的消耗非常大仅适合小规模验证。对于实际 LIC 模型建议用 Fisher 近似。5.2 敏感度排序与比特分配有了每层敏感度之后如何把比特分配变成一个可执行的搜索过程下面给出一个贪心分配的示例。假设我们预先定义候选 bit 档位例如[2, 4, 6, 8]目标平均 bit 数为target_avg_bit。# 文件路径bit_allocation.py import copy def greedy_bit_allocation(layers_sensitivity, layers_weight_size, candidate_bits(2, 4, 6, 8), target_avg_bit6.0): 基于敏感度做贪心比特分配。 参数: layers_sensitivity: dictkey 为层名value 为该层敏感度标量如 trace layers_weight_size: dictkey 为层名value 为该层权重参数量 candidate_bits: 可选比特档位 target_avg_bit: 目标平均比特数 返回: bit_config: dictkey 为层名value 为分配的比特数 # 初始化所有层为最高 bit max_bit max(candidate_bits) bit_config {name: max_bit for name in layers_sensitivity} total_params sum(layers_weight_size.values()) current_bits sum(bit_config[n] * layers_weight_size[n] for n in bit_config) target_bits target_avg_bit * total_params while current_bits target_bits: best_loss_ratio None best_layer None best_next_bit None for name in bit_config: current_bit bit_config[name] # 找到当前层的下一档低 bit lower_bits [b for b in candidate_bits if b current_bit] if not lower_bits: continue next_bit max(lower_bits) # 敏感度越高降低 bit 的代价越大 # 所以选择“敏感度 / 节省的 bit 量”最小的层 bit_saved (current_bit - next_bit) * layers_weight_size[name] cost layers_sensitivity[name] * (current_bit - next_bit) / bit_saved if best_loss_ratio is None or cost best_loss_ratio: best_loss_ratio cost best_layer name best_next_bit next_bit if best_layer is None: break bit_config[best_layer] best_next_bit current_bits - (max_bit - best_next_bit) * layers_weight_size[best_layer] return bit_config这段代码的核心思想是“每一步降低一层 bit尽量优先降低敏感度低、节省 bit 多的层”。它本质上是一个贪心策略虽然不保证全局最优但在实际部署中效率很高得到的配置通常已经接近 DP 方案。5.3 配置文件示例实际项目中量化配置最好不要硬编码在代码里而是放进 YAML 或 JSON 配置文件方便实验管理和复现。下面给出一份简单的量化配置示例# 文件路径configs/hampl_lic_config.yaml model: name: minnen_hyperprior checkpoint: ./checkpoints/fp32_mse_lambda_0.01.pth quantization: backend: ptq # 后训练量化 calibration_images: 256 # 校准图片数量 data_path: ./data/calibration batch_size: 4 image_size: 256 target_avg_bit: 5.0 # 目标平均比特数 candidate_bits: [2, 4, 6, 8] # 候选比特档位 weight_quant_type: per_channel # 权重按通道量化 activation_quant_type: per_tensor search: method: greedy # greedy / dp / random max_iterations: 100 eval: test_dataset: kodak metric: [psnr, ms_ssim, bd_rate]配置文件的好处是你换一组候选 bit、换一个目标平均 bit 数不需要改代码只需要改 YAML。这一点在实验比较多的时候非常省事。6. 实验设计与验证6.1 数据集与评价指标学习型图像压缩领域常用的评测数据集主要有Kodak24 / Kodak经典的 24 张标准测试图像分辨率固定为 512×768。CLICChallenge on Learned Image Compression包含专业和移动端图像。Tecnick100 张高清自然图像分辨率较高。DIV2K常用于训练或作为附加验证集。评价指标一般有两类重建质量指标PSNR、MS-SSIM。压缩效率指标BD-Rate、BD-PSNR。BD-Rate 表示在相同重建质量下码率的平均节省百分比负值表示比基准更省码率。量化模型和 FP32 模型比较时BD-Rate 越接近 0说明量化对压缩性能的影响越小。HAMP-LIC 的实验中一般会报告统一精度与混合精度在相同平均 bit 数下的 BD-Rate 对比以及不同目标 bit 下的 PSNR 曲线。6.2 对比基线在验证混合精度量化方法时应该至少做这几组对比FP32 基线未量化模型作为性能上界。统一 INT8常见部署方案说明标准量化的效果。统一 INT4验证低位宽统一量化的损失。MSE 敏感度混合精度用简单 MSE 重建误差排序做混合精度作为对比。Hessian/Fisher 敏感度混合精度本文方法。通过对比你可以看出“Hessian-aware 敏感度分析”相对于“简单按层大小或 MSE 排序”到底带来了多少提升。6.3 实验结果解读要点实验报告中通常会有两类结论Hessian 敏感度排序与人工经验一致例如编码器首层、熵模型相关层敏感度更高被分配到更高 bit。混合精度优于统一精度在相同总参数 bit 数下混合精度方案的 BD-Rate 损失更小。在你自己复现时不要只看最终数字还要分析哪些层被分配了高位宽是否符合直觉校准图像数量和分布是否影响了敏感度排序量化后码率上升主要集中在哪个模块7. 常见问题与排查思路混合精度 PTQ 在实际跑的时候会遇到不少问题。下面整理一些高频现象和排查方向。问题现象常见原因解决思路量化后码率明显上升熵模型/超先验相关层被分配到过低 bit概率估计失真检查熵模型相关卷积是否被分配低位宽强制提升其 bit 数PSNR 下降严重但码率变化不大主解码器层量化噪声直接体现在重建图像上提高解码器敏感层 bit或对激活做 per-tensor 改为 per-channel校准集上误差很小测试集上误差大校准数据分布与测试分布不一致过拟合到校准集扩充校准集尽量覆盖不同场景和分辨率Hessian/Fisher 估计不稳定校准批大小过小梯度噪声大加大 batch size或多次采样取平均贪心搜索得到的 bit 配置不符合预期敏感度标量计算公式不合适尝试不同的聚合方式trace、max eigen、平均对角线模型显存不足Hessian 估计需要二次反向传播改用 Fisher 近似或对校准数据分批处理部署框架不支持某些混合 bit硬件/推理库只支持 8bit/4bit 统一量化把混合精度结果映射到最近的可支持档位其中最常见的坑是“熵模型量化导致码率爆炸”。因为 LIC 模型的比特消耗是由熵模型给出的概率分布决定的只要概率估计有偏差实际熵编码后的码率就会增加。在 HAMP-LIC 的框架下这意味着超编码器、超解码器、上下文模型相关的卷积层通常是敏感度最高的层应该分配更高 bit甚至保持 FP16 或 INT8。8. 量化部署最佳实践与工程建议8.1 校准确数据集校准集是 PTQ 中最容易被低估的部分。几点建议不要只从训练集随机抽图而是混合训练集、验证集和少量业务真实图像。校准集数量建议在 200~1000 张之间。太少会导致统计量不稳定太多会拖慢 Hessian 估计。校准图像的分辨率要和实际部署场景一致特别是 LIC 模型对输入分辨率敏感。8.2 量化粒度选择权重量化建议优先使用per-channel因为不同输出通道的权重范围差异很大per-tensor 量化会让范围小的通道被截断或精度下降。激活量化如果硬件支持也尽量使用 per-channel如果不支持 per-channel 激活量化可以在感知敏感度分析时把激活量化误差也纳入 Hessian 估计中而不仅仅分析卷积权重。HAMP-LIC 文章中更关注权重量化带来的性能退化但实际部署中激活量化往往也需要混合精度或相应补偿。8.3 与部署框架对接无论使用 TensorRT、OpenVINO、ONNX Runtime 还是自研推理引擎都有一些通用注意点先确认推理框架对混合精度的支持粒度。有些框架只支持整个卷积层统一 bit有些支持通道粒度。权重从 PyTorch 导出时要确保量化参数scale、zero point能被 ONNX/推理引擎正确解析。部署前最好在 C 侧写一个针对量化模型的回归测试逐层比较输出和 PyTorch 参考结果的误差误差超过阈值的层及时调整 bit 配置。8.4 安全与可维护性这部分容易被忽略但在生产环境很重要量化模型的 bit 配置、校准集、FP32 权重、评估指标要一起归档方便追溯。对生产环境的模型做变更时先在测试集上做完整评估确认 BD-Rate 退化在可接受范围后再发布。涉及模型文件下载、部署、更新时遵循团队发布流程不要在未验证的情况下直接替换线上模型。如果后续要调整某个 bit 配置不要手动改模型文件里的数字而是通过代码和配置文件重新生成保证可复现。8.5 性能优化方向混合精度量化的收益不止是模型体积变小还可能有推理时延收益。但要注意混合精度的加速效果取决于后端算子是否真的对每种 bit 都有优化实现。有些框架把 4bit 算子做了特殊优化有些则会把 4bit 反量化回 FP16 再计算后者可能反而更慢。如果目标是端侧实时编码/解码优先关注超先验和上下文模型路径上的算子延迟而不是只盯着总体模型体积。可以考虑把部分 2bit 或 4bit 层合并为统一的低比特算子减少 kernel launch 开销。9. 下一步可以继续深入的方向如果这篇文章让你对混合精度量化产生了兴趣建议按下面顺序继续学习先实现一个基线 PTQ把一个简单的 LIC 模型整体量化到 INT8跑通部署链路。然后加入敏感度分析用 Fisher 对角线估计每层敏感度观察哪些层被分到高 bit。再实现混合精度搜索从贪心算法开始再尝试 DP 或贝叶斯搜索。最后结合 QAT/蒸馏如果混合精度 PTQ 在极端低 bit 下仍然不满足精度要求可以用少量蒸馏微调敏感层形成 PTQ轻量微调的混合方案。HAMP-LIC 的价值在于它给出了一个完整的思路链路从 LIC 模型的率失真损失出发用 Hessian 信息描述量化扰动的影响再通过混合精度分配实现精度和体积的平衡。这套思路不仅适用于图像压缩对视频压缩模型、生成模型、扩散模型的量化也有参考意义。实际动手时可以先用小模型验证 Hessian 估计和搜索流程的正确性再迁移到完整 LIC 模型上。量化没有万能方案不同模型结构、不同数据集、不同后端硬件都会影响最终效果但 Hessian-aware 敏感度分析这个思路基本是通用的第一步。
返回列表