HAMP-LIC：Hessian感知的混合精度量化，解决图像压缩模型部署难题-平芜编程栈

图像压缩模型近年来在率失真性能上已经明显超越传统编码标准，但“能跑实验”和“能落地部署”之间还存在一条不小的鸿沟。端到端模型的编码端、解码端都包含大量浮点算子，即使训练好的模型精度很高，一旦直接做后训练量化，潜在特征分布和熵模型的异常敏感常常让压缩性能快速崩塌。这里的关键矛盾在于：不是每个层对量化误差都同样敏感，统一设低比特会牺牲质量，统一设高比特又达不到压缩和加速的目的。HAMP-LIC（Hessian-Aware Mixed-Precision Post-Training Quantization for Learned Image Compression）就是这样一个被设计用来解决该问题的思路：在海量敏感度信息引导下，对学习图像压缩模型做混合精度后训练量化，让不同的模块拿到各自能承受的位宽，在尽量不重训的前提下保住率失真性能。

这篇文章会围绕 HAMP-LIC 的实际价值展开：它到底解决什么问题、方法流程如何拆解、要用哪些工具复现或验证、实验时该看哪些指标、以及最常见的坑在哪里。如果你正在做学习图像压缩的模型压缩，或者想把 LIC 模型推到边缘设备，这篇文章可以直接收藏。

1. 核心能力速览

能力项	说明
项目名称	HAMP-LIC: Hessian-Aware Mixed-Precision Post-Training Quantization for Learned Image Compression
项目类型	研究型模型压缩方法，面向学习图像压缩（LIC）的后训练量化方案
核心解决方向	后训练量化场景下，混合精度位宽分配带来的率失真性能退化
关键机制	Hessian 矩阵 / 二阶敏感度估计，混合精度位宽搜索，无需端到端重训
目标模型	自编码器结构、超先验结构等常见学习图像压缩模型
典型评估指标	PSNR、MS-SSIM、BD-Rate、编码/解码耗时、模型体积
运行框架	Python + PyTorch，可基于 compressAI 生态进行复现验证
启动方式	脚本式实验，通常为 train / evaluate 模式，不涉及 WebUI
是否支持 CPU	量化与评估可以跑 CPU，但 Hessian 估计和模型推理明显更吃 GPU
显存占用	依赖模型大小、批大小和 Hessian 估计方式，需以实际环境测试为准
API 接口	论文理论不涉及，研究者可自行封装为推理服务
批量任务	可批量处理测试图像集，需自行编写数据循环
适合读者	图像压缩方向研究生、边缘端 AI 部署工程师、量化算法工程师

从表里可以直观看到，HAMP-LIC 不是一个拿来即用的部署工具，它更像一套“如何为 LIC 模型找到合理低比特位宽分配”的方法论。要跑通这套流程，需要自己准备模型权重、校准数据集、量化配置和评价脚本。

2. 为什么学习图像压缩模型需要后训练量化

学习图像压缩模型的典型流程是：编码器把输入图像映射到潜在表示，潜在表示经过量化后，再交给熵编码器压缩成比特流。解码端拿到比特流后，先通过熵解码恢复量化后的潜在表示，再用解码器重建图像。整个过程可以端到端训练，最终优化的损失函数是码率和失真的加权和：

[ L = R(\hat{y}) + \lambda \cdot D(x, \hat{x}) ]

其中 (R(\hat{y})) 是熵模型对量化潜在表示的码率估计，(D) 是重建图像与原始图像之间的失真，(\lambda) 控制码率和失真之间的权衡。

关键在量化这一步。训练时，模型通常使用直通估计器来近似量化操作的梯度，让网络学会适应量化的噪声。但到了部署阶段，很多工程团队不会执行完整的量化感知训练，因为那需要准备训练数据、重放训练流程、重新调超参，成本非常高。因此，后训练量化（Post-Training Quantization，PTQ）成为更现实的选项：在模型已经训练好的前提下，用一小部分校准数据统计激活分布，把权重和激活从浮点转成低比特整数。

但 LIC 模型做 PTQ 比普通分类网络更难。原因主要有三个：

潜在表示分布高度不均匀，不同图像、不同通道的数值范围差异大，统一量化会导致较大的信息损失。
熵模型与编码器、解码器强耦合，量化误差会传导到概率估计上，再传导回码率估计，最终影响的是“码率-失真”整体表现，而不只是单张图片的像素误差。
模型的不同层对量化误差的容忍度不同。有的层稍微量化就掉点严重，有的层即使压到 4bit 也几乎无感。

因此，如果仅仅给整个模型设一个统一的位宽，结果往往是在“模型太大”和“质量损失过多”之间二选一。HAMP-LIC 的核心动机，就是打破这种“一刀切”的量化方式，用敏感度来驱动位宽分配。

3. HAMP-LIC 方法拆解

从标题可以拆出三个关键词：Hessian-Aware、Mixed-Precision、Post-Training Quantization。这三者共同构成 HAMP-LIC 的方法骨架。

3.1 后训练量化：避免大规模重训

HAMP-LIC 的工作前提是模型已经完成率失真优化训练。后续的量化过程不再改变模型内部参数，只需要在少量校准数据上统计量化范围，并通过敏感度分析来决定每个模块应该用多少位宽。这种设定对工程场景很重要，因为重训一个 LIC 模型往往需要多卡训练数天，而 PTQ 的校准时间通常在分钟到小时级别。

PTQ 中还有一个前置问题：对权重做量化，还是对激活做量化，还是两者都做。在图像压缩模型里，权重和激活都会影响最终码流。潜在表示量化这一环本身就在模拟编码器端的行为，而激活量化会直接影响重建质量。实践中通常要分别统计，HAMP-LIC 的敏感度估计也应当覆盖编码器、解码器、超先验等不同子网络。

3.2 Hessian 感知：用二阶信息量化敏感度

那么怎么判断某个层到底能不能压到低比特？直接用一层层的误差去试，成本高且不精确。更可靠的做法是看损失函数对某一层输出的二阶导数，也就是 Hessian。这个逻辑在传统模型压缩里已经验证过：损失函数对参数的曲率越大，说明参数越敏感，量化误差带来的影响越难恢复。

在 HAMP-LIC 这类方法里，通常用 Hessian 的迹或对角近似来表示每个量化单元的敏感度。设有一层参数 (w)，量化带来的扰动为 (\Delta w)，损失函数的变化可以近似为：

[ \Delta L \approx \frac{\partial L}{\partial w} \cdot \Delta w + \frac{1}{2} \Delta w^T H \Delta w ]

在模型已经收敛的情况下，一阶项接近零，二阶项成为主导。所以，Hessian 矩阵 (H) 越大，说明该层对参数扰动的放大作用越强，量化误差越危险。反过来，Hessian 较小的层可以分配更低的位宽。

实际计算中，直接展开完整 Hessian 矩阵是不现实的。常用的替代方案是 Hutchinson 估计，通过随机向量 (v) 估计 Hessian 的迹：

[ \mathrm{Tr}(H) \approx \frac{1}{N} \sum_{i=1}^{N} v_i^T H v_i ]

PyTorch 下可以通过两次自动求导算出向量-海森乘积（Hessian vector product，HVP），不需要显式构造矩阵。

3.3 混合精度：敏感度决定位宽

拿到各层/各模块的敏感度后，接下来就是混合精度位宽分配。候选位宽通常是 4bit、6bit、8bit 这样的整数宽度，也可以包含 12bit 或 16bit 来保证关键路径不丢精度。分配的目标是：在总位宽预算或总模型体积约束下，最小化整体的量化误差。

这里面有个平衡：低敏感度层用低比特，高敏感度层用高比特。寻优方式既可以做成贪心，也可以建模成整数优化问题。贪心策略通常是先给所有层一个基础位宽，比如 8bit，然后计算把每层降到 6bit 的代价，优先把“代价最小”的层降下来，直到满足目标体积或平均位宽要求。

在实现上，混合精度带来一个额外的兼容性问题：如果不同层位宽不同，最终部署时就要为每种位宽准备独立的算子实现。硬件上可能支持也可能不支持。因此，实验阶段跑通不代表部署阶段一定能跑通，这一点在做 HAMP-LIC 验证时要提前意识。

3.4 整体流程

把上面的三个模块串起来，HAMP-LIC 的完整实验流程可以概括为：

加载已训练的学习图像压缩模型。
准备少量校准图像，计算率失真损失。
对候选量化层做 Hessian 迹估计，得到敏感度排序。
根据敏感度与位宽预算，搜索混合精度配置。
按配置对模型各层做后训练量化。
在验证集上评估 PSNR、MS-SSIM、BD-Rate，并观察编码时间与模型大小变化。

这套流程中真正耗时的部分是 Hessian 估计。如果模型很大，校准数据很多，建议控制迭代次数或先只在部分层上估计。

4. 复现与验证环境准备

HAMP-LIC 的官方完整代码目前不一定有公开实现，复现时需要依据方法的理论描述自行搭建。比较稳妥的路线是在 compressAI 生态上做二次开发，因为 compressAI 提供了大量已训练的学习图像压缩模型权重和评估工具。

4.1 基础环境

建议使用 Linux 环境，Python 3.10 以上，CUDA 显卡驱动按 PyTorch 版本选择。这里给出一份通用环境清单：

conda create -n hamp-lic python=3.10 -y conda activate hamp-lic pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install compressai==1.3.6 pip install scipy pandas tensorboard matplotlib tqdm

实际安装时，--index-url里的 CUDA 版本号要按照本机驱动和 PyTorch 支持情况修改。如果本机只有 CPU，也可以安装 CPU 版 PyTorch 做小规模验证计算，但 Hessian 估计会比较慢。

4.2 测试数据

图像压缩方向常用 Kodak 数据集做验证，包含 24 张 768x512 的图像。CLIC 数据集则更大，适合做完整训练和严格评测。如果只是用来验证量化和位宽搜索流程，可以准备几十张不同类型的自然图像作为校准集，再准备若干张图像作为验证集。

校准集不需要很多，通常 16 到 64 张即可。关键是覆盖不同场景：风景、建筑、人物、文本区域，这样量化范围统计才更稳健。

4.3 模型选择

compressAI 提供了多类预训练模型，例如bmshj2018-factorized、cheng2020-anchor等。HAMP-LIC 这类方法的设计目标通常是通用 LIC 模型结构，但复现时建议从单一模型开始，比如：

from compressai.zoo import cheng2020_anchor # quality 参数控制码率-失真权衡点 net = cheng2020_anchor(quality=4, pretrained=True).eval()

量化实验前，先跑通模型的原始推理，记录浮点基线的 PSNR 和 MS-SSIM，作为后续比较的基准。

5. 功能测试与效果验证

HAMP-LIC 这类研究型方法，验证的重点不是“能不能出图”，而是“量化后率失真性能退化多少”。按照下面的流程，可以在自己的环境里跑出一套可比较的量化评估结果。

5.1 浮点基线评估

写一个评估脚本，对验证集逐张执行net(image)，得到压缩后的重建图像、码率，再计算 PSNR 和 MS-SSIM。这个结果就是浮点基线，后续所有量化评估都要和它对齐。

import torch from torchvision.transforms.functional import to_tensor from compressai.zoo import cheng2020_anchor device = "cuda" if torch.cuda.is_available() else "cpu" net = cheng2020_anchor(quality=4, pretrained=True).eval().to(device) def evaluate_image(net, img_tensor, device): img_tensor = img_tensor.unsqueeze(0).to(device) with torch.no_grad(): out = net(img_tensor) # out["x_hat"] 是重建图像，out["likelihoods"] 用于估计码率 return out

记录每张图像的比特率（bits per pixel，bpp）和 PSNR 值，保存到表格。

5.2 校准数据构建

校准数据要从原始图像中裁剪出小块，避免整图过大导致显存溢出。常见的做法是把图像切成分辨率 256x256 或 512x512 的块，然后转成模型输入格式：

from PIL import Image import torch from torchvision.transforms.functional import to_tensor def build_calib_loader(image_paths, patch_size=256, batch_size=1): images = [] for path in image_paths: img = Image.open(path).convert("RGB") img = img.resize((patch_size, patch_size)) images.append(to_tensor(img)) data = torch.stack(images) loader = torch.utils.data.DataLoader(data, batch_size=batch_size, shuffle=False) return loader

batch_size不建议设置太大，因为 Hessian 估计需要保留计算图，显存占用比普通推理高很多。

5.3 Hessian 敏感度估计

Hessian 迹估计可以通过自动求导实现，核心是计算损失对选定参数的梯度，再反向计算向量-海森乘积。代码思路是这样的：

def estimate_hessian_trace(model, loader, loss_fn, target_params, n_iter=10): traces = {name: 0.0 for name, _ in target_params} for batch in loader: batch = batch.to(device) model.zero_grad() out = model(batch) loss = loss_fn(out) # 率失真损失 grads = torch.autograd.grad(loss, [p for _, p in target_params], create_graph=True) for (name, p), grad in zip(target_params, grads): trace_est = 0.0 for _ in range(n_iter): v = torch.randn_like(grad) hvp = torch.autograd.grad(grad, p, grad_outputs=v, retain_graph=True)[0] trace_est += (v * hvp).sum().item() traces[name] += trace_est / n_iter model.zero_grad() return traces

这里的关键细节是create_graph=True和retain_graph=True，否则无法做二次求导。损失函数要包含码率和失真两部分，如果只使用 MSE 失真，敏感度无法反映对码率的影响。

5.4 混合精度位宽搜索

拿到敏感度后，先给所有层设置一个参考位宽，比如 8bit，再基于敏感度从低到高排序，逐层尝试降到 6bit 或 4bit。判断依据是每一步量化带来的损失增量。这是一个典型的贪心近似，实现复杂度低，效果已经比统一位宽好很多。

def greedy_bitwidth_search(sensitivity, base_bits=8, candidates=(8, 6, 4), target_avg=6.5): layers = sorted(sensitivity.items(), key=lambda x: x[1]) # 先全部按 base_bits 量化，再对敏感度低的层降比特 assignment = {name: base_bits for name, _ in layers} # 根据目标平均位宽计算可降低的总 bit 数 total_reduction = sum(base_bits for _ in layers) - target_avg * len(layers) for name, _ in layers: for bits in sorted(candidates, reverse=True): if bits >= assignment[name]: continue reduction = assignment[name] - bits if total_reduction >= reduction: assignment[name] = bits total_reduction -= reduction if total_reduction <= 0: break if total_reduction <= 0: break return assignment

注意，这个示例代码对应的是“各层内部统一量化位宽”的设定。如果需要做到通道级混合精度，实现会更复杂，但搜索逻辑本质相同：敏感度低的通道，用更低的位宽。

5.5 量化与评估

量化位宽分配完成后，需要对模型执行实际量化。最直接的做法是把浮点权重映射到目标位宽的整数量化范围，并在推理时反量化回浮点，用于模拟部署效果。模拟部署时激活值也按相同策略做量化统计：

import torch def quantize_tensor(x, bits): qmax = 2 ** (bits - 1) - 1 scale = x.abs().max() / qmax x_q = torch.clamp(torch.round(x / scale), -qmax, qmax) return x_q * scale

这只是最基础的对称量化模板。HAMP-LIC 工程的实现会根据每层统计的数值范围选择更精细的量化参数。评估方式不变，仍然计算 PSNR、MS-SSIM 和 bpp。

如果量化后的 PSNR 相对浮点基线下降小于 0.2dB，说明量化策略基本可行；如果下降超过 0.5dB，就需要检查敏感度估计是否准确，或者某些层是否被分配了过低的位宽。

5.6 判断量化是否成功的标准

在量化实验里，判断标准不能只看单张图是否清晰。至少要满足三个条件：

码率与量化前相比没有明显恶化，如果 bpp 出现异常升高，说明量化破坏了熵模型的概率估计。
PSNR / MS-SSIM 下降幅度在任务可接受范围内。
模型体积和推理耗时确实因为低比特量化而下降，否则“混合精度”的实际收益就不存在。

6. 显存占用与性能观察方法

HAMP-LIC 的显存占用主要来自三个阶段：普通推理、Hessian 迹估计、量化后模拟推理。其中 Hessian 迹估计阶段的显存压力最大，因为需要保留梯度图和中间激活。

普通推理阶段，可以用nvidia-smi实时观察显存占用。以 compressAI 的cheng2020_anchor模型为例，在 256x256 输入分辨率下，浮点推理通常只需要 1 到 3GB 显存，具体取决于批大小。Hessian 估计阶段，显存占用可能显著上升，因为每个 batch 都要保留多个梯度图，如果出现 OOM，优先减小批大小到 1，或者减小校准图像分辨率。

CPU 推理也是可行的，但速度会明显下降。Hessian 估计在 CPU 上跑尤其慢，因为涉及多次反向传播。如果只是验证 bit 分配是否合理，可以把 Hessian 估计的迭代次数降到 3 到 5 次。

批量评估时，建议把测试数据按分辨率分组处理，避免不同分辨率图像在同一个 batch 里触发布局不匹配问题。多进程加载数据时，也要注意显存是否被多个 DataLoader worker 放大。实际测试时可以把num_workers设为 0 或 2，优先保证显存稳定。

另外要观察量化前后编码时间的变化。低比特量化在某些硬件上会加速，但在 PyTorch 模拟阶段，由于需要执行缩放和反量化操作，耗时可能反而增加。正式部署时，需要依赖推理引擎和硬件算子库来获得真实加速。

7. 常见问题与排查方法

问题现象	可能原因	排查方式	解决方案
量化后 PSNR 掉得厉害	高敏感层被分配了过低比特位宽	输出各层敏感度和位宽对照表	提高该层位宽，或改用更精细的通道级敏感度估计
bpp 明显升高	量化破坏了熵模型的概率分布	对比量化前后似然输出	熵模型相关层保持较高位宽，不做激进量化
Hessian 估计显存溢出	批大小过大、输入分辨率过高、梯度图太多	观察显存占用峰值	减小批大小到 1，降低输入分辨率，减少 Hessian 迭代次数
量化模拟结果不稳定	校准数据覆盖不足，量化范围统计不准	检查校准图像类别和数量	增加不同类型图像作为校准集，或使用更稳健的离群点处理
推理速度反而变慢	PyTorch 模拟量化没有走到低比特算子	检查是否使用真正的量化后端	实验阶段接受模拟结果，部署阶段再接入推理引擎
不同层位宽不同，部署代码写起来很复杂	混合精度带来算子碎片化	检查硬件是否支持混合位宽	先做层级混合精度，验证收益后再考虑通道级混合精度
原始浮点模型和量化模型输出尺寸不一致	模型结构被改动，或量化过程中不小心改了 forward 逻辑	打印每层输出 shape	逐层检查模型结构，量化包装类不要改动原 forward 逻辑
BD-Rate 计算报错	测试点数和码率点不足	检查是否测试了多个 quality 或 lambda 点	至少测试 3 到 4 个不同质量档位

8. 合规边界与部署最佳实践

HAMP-LIC 的使用边界要提前说明清楚。学习图像压缩模型本身可能受到不同开源协议约束，比如 compressAI 中的预训练权重有自己的授权方式。复现和二次开发时，要确认模型权重、测试数据集的使用条件。尤其在实际工程中，如果压缩的是人物照片、证件、医疗影像等敏感数据，需要评估数据隐私合规要求，不能因为“本地跑了一下”就忽略数据安全。

对于把 HAMP-LIC 方法应用到实际产品的团队，更稳妥的工程化路径是：

先用公共数据集做浮点基线和量化对比，确定方法是否有效。
评估目标硬件是否支持混合精度算子，避免量化策略在部署阶段被推翻。
对高敏感度的层保留较高位宽，确保率失真性能。
在正式上线前用真实业务数据做一轮交叉验证，确认量化模型对业务数据分布不敏感。
如果是商业产品或开源项目，注意模型权重、代码、数据集的许可证兼容性。

这样做的好处是能在“量化精度”和“部署可行性”之间找到落地解，而不是只在论文理想设定下好看。

9. 总结与下一步

HAMP-LIC 最值得关注的不是某层具体设 4bit 还是 8bit，而是它把“量化位宽选择”从经验调参变成了有明确理论依据的优化问题。用 Hessian 信息描述层敏感度，再结合混合精度分配，可以在不重训模型的情况下缓解后训练量化在图像压缩模型上的性能退化。

如果想自己验证这套方案，我建议从下面几步开始：

找一个 compressAI 预训练模型，先跑通浮点基线。
准备 20 到 50 张图像的校准集，测试 Hessian 估计流程是否稳定。
用贪心搜索做一层混合精度分配，对比统一 6bit 量化的 PSNR 和 bpp。
如果 Hessian 估计显存占用过高，先尝试只对解码器部分做敏感度估计，把流程跑通后再扩大范围。

最容易踩的坑有三个：一是 Hessian 估计时忘了retain_graph=True，导致二次求导失败；二是校准集过小，量化范围统计不稳定；三是只看 PSNR，不看 bpp，导致熵模型被破坏的隐患被掩盖。

下一步可以关注的方向是把敏感度估计从“层级别”细化到“通道级别”，以及把量化策略放到真实推理引擎上做时延和显存验证。还可以尝试把 HAMP-LIC 的 Hessian 估计方法推广到其他生成模型上，比如扩散模型或超分模型。整个方向的核心逻辑是一致的：先搞清楚哪里对量化敏感，再把比特花在最需要的地方。