INT8量化掉精度？我折腾了三天，用分层量化+QAT把8%的损失捞回7.5%

INT8量化：速度翻倍的诱惑与8%精度的代价

上周在客户现场调试，差点被甲方工程师的眼神“刀”死。事情是这样的，我给一家中型物流公司“快运通”做了一套视觉分拣系统，核心是一个基于YOLOv8s的包裹面单检测模型。在RTX 4090上跑得飞快，但部署到他们产线那批老旧的Jetson Xavier NX上，一帧推理要45ms，离实时处理（要求30fps，即33ms以内）差了一口气。甲方老哥天天催，说产线效率卡在这了。我的第一反应就是上INT8量化，这几乎是边缘设备性能提速的“标准答案”。TensorRT官方文档吹得天花乱坠，号称精度损失极小，速度翻倍。我心想着，这不就妥了嘛。

我兴冲冲地搞起来。我们的模型输入是640×640的RGB图像，输出是检测框和类别。原始的FP16模型在Xavier NX上的性能是45ms一帧，精度（mAP@0.5）在客户提供的验证集上是92.1%。我按照TensorRT的流程，准备了1000张校准图像（直接从训练集里随机抽的），写了个简单的校准器，跑了一遍Post-Training Quantization (PTQ)。

import tensorrt as trt
import numpy as np
import cv2
import os

class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_data_dir, batch_size=8, input_shape=(640, 640)):
        super().__init__()
        self.batch_size = batch_size
        self.input_shape = input_shape
        # 1. 加载校准图像路径
        self.image_list = [os.path.join(calibration_data_dir, f) for f in os.listdir(calibration_data_dir) if f.endswith('.jpg')]
        self.current_index = 0
        self.cache_file = 'calibration.cache'
        # 分配一次内存，避免重复分配
        self.device_input = cuda.mem_alloc(self.batch_size * 3 * input_shape[0] * input_shape[1] * np.float32().itemsize)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        """TensorRT校准器调用的核心函数，必须实现"""
        if self.current_index + self.batch_size > len(self.image_list):
            return None  # 返回None表示校准结束

        batch_images = []
        for i in range(self.batch_size):
            img_path = self.image_list[self.current_index + i]
            img = cv2.imread(img_path)
            # 2. 预处理：缩放到输入尺寸，归一化，转RGB，转CHW
            img = cv2.resize(img, self.input_shape)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = img.astype(np.float32) / 255.0  # 归一化到[0,1]
            img = img.transpose(2, 0, 1)  # HWC -> CHW
            batch_images.append(img)

        self.current_index += self.batch_size
        # 将batch数据拼接并拷贝到GPU
        batch_data = np.ascontiguousarray(np.stack(batch_images, axis=0), dtype=np.float32)
        cuda.memcpy_htod(self.device_input, batch_data)
        # 3. 返回一个GPU内存指针的列表，这里只有一个输入
        return [int(self.device_input)]

    def read_calibration_cache(self):
        """读取校准缓存，加速后续构建"""
        if os.path.exists(self.cache_file):
            with open(self.cache_file, 'rb') as f:
                return f.read()
        return None

    def write_calibration_cache(self, cache):
        """写入校准缓存"""
        with open(self.cache_file, 'wb') as f:
            f.write(cache)

# 使用校准器构建INT8引擎
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

# 解析ONNX模型
with open('yolov8s.onnx', 'rb') as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)
config.set_flag(trt.BuilderFlag.FP16) # 同时开启FP16，某些层可能保持FP16
config.max_workspace_size = 1 << 30  # 1GB

# 设置校准器
calibrator = MyCalibrator('calibration_data/')
config.int8_calibrator = calibrator

# 构建序列化引擎
engine = builder.build_serialized_network(network, config)
with open('yolov8s_int8.engine', 'wb') as f:
    f.write(engine)

构建过程很顺利，引擎文件也生成了。我迫不及待地跑了一下Benchmark，好家伙，推理时间直接从45ms干到了22ms，整整快了一倍还多！我差点就在客户现场笑出声。但紧接着跑精度测试，我的心就凉了半截。mAP@0.5从92.1%暴跌到了84.3%，掉了将近8个百分点！这还玩个屁，产线分拣错误率飙升，客户非得把我从厂房扔出去不可。INT8量化的“代价”，这8%的精度差，像一盆冰水浇醒了我。

校准集选错了，一切白搭——我的第一次尝试就翻车了

我第一个怀疑的就是校准集。之前图省事，直接从训练集随机抽了1000张。但训练集分布和真实产线数据能一样吗？我赶紧要来了产线最近一周的监控截图，大概5000张，分布和训练集还真有区别：

• 光照更复杂：有夜间灯光、有阳光直射导致的过曝区域。
• 背景更杂乱：传送带上的包裹堆叠、工人偶尔入镜。
• 面单状态多样：有褶皱、有部分遮挡、有反光。

我重新用这5000张产线图片作为校准集，满怀希望地又跑了一遍PTQ。结果呢？精度从84.3%提升到了……85.1%。嗯，提升了0.8%，聊胜于无，但离92%的原点还差得远。这说明校准集确实有影响，但不是问题的全部。我意识到，对于YOLO这种单阶段检测器，特别是包含大量非线性操作（如SiLU激活函数）和跨层连接的复杂结构，粗暴的全局INT8量化，肯定会把某些对数值范围敏感的“关键层”给搞残废。

我陷入了困境。客户那边等着上线，性能不达标，精度不达标。我尝试了TensorRT提供的其他校准方法，从熵校准（EntropyCalibrator）换到最小值最大值校准（MinMaxCalibrator）。

# 尝试MinMaxCalibrator
class MyMinMaxCalibrator(trt.IInt8MinMaxCalibrator):
    def __init__(self, data_dir, batch_size=8, input_shape=(640, 640)):
        # ... 初始化部分与熵校准器类似
        self.cache_file = 'calibration_minmax.cache'
        # MinMax校准需要记录所有batch的绝对最大值，这里用列表暂存
        self.batch_max_vals = []

    def get_batch(self, names):
        # ... 获取batch数据
        # 计算这个batch的绝对值最大值（模拟过程，实际TensorRT内部会处理）
        batch_abs_max = np.abs(batch_data).max()
        self.batch_max_vals.append(batch_abs_max)
        return [int(self.device_input)]

    def read_calibration_cache(self):
        # ...
    def write_calibration_cache(self, cache):
        # ...

    def get_algorithm(self):
        return trt.CalibrationAlgoType.MINMAX_CALIBRATION

# 构建时更换校准器
config.int8_calibrator = MyMinMaxCalibrator('production_calibration_data/')

跑出来的结果更差，精度掉到了83.5%。看来对于这种动态范围比较大的激活输出，熵校准（找信息量最大的分布区间）通常比简单的最小最大值校准更鲁棒。这个坑告诉我：校准方法不是随便选的，得根据数据分布来试。但即便如此，最好的结果也才85.1%，问题没解决。

诊断工具不能少：用Netron和直方图揪出“坏”算子

光试不行，得知道模型内部到底发生了什么。我祭出了两样工具：Netron和自定义的激活值直方图统计脚本。首先用Netron打开YOLOv8s的ONNX模型，一层层看结构。我发现几个可疑点：

1. SPPF模块末尾的Concat层：它拼接了不同尺度的特征图，这些特征图的数值范围可能差异很大，强行用同一个缩放因子量化，必然有一方失真严重。
2. 靠近检测头的卷积层：这些层的输出直接关系到框的坐标和类别置信度，精度要求极高。
3. SiLU激活函数之后：SiLU (Sigmoid-Weighted Linear Unit) 本身是非线性的，其输出分布可能不适合用简单的对称量化（INT8是-127到127）。

光看结构是猜测，我需要数据佐证。我写了个脚本，在FP16模型推理时，拦截特定层的输入输出，统计它们的数值分布。

import torch
import numpy as np
import matplotlib.pyplot as plt
from models.yolo import Model  # 假设你有YOLOv8的PyTorch模型定义

# 加载FP16模型
fp16_model = Model('yolov8s.yaml').cuda().half().eval()

# 定义钩子来捕获中间层激活值
activation_stats = {}
def get_activation_hook(name):
    def hook(module, input, output):
        # 只统计前几个batch，避免内存爆炸
        if name not in activation_stats:
            activation_stats[name] = []
        # 将输出转为CPU numpy数组，并取绝对值最大值观察范围
        if isinstance(output, torch.Tensor):
            output_np = output.detach().cpu().numpy()
            # 记录该层输出张量的全局绝对最大值
            abs_max_val = np.abs(output_np).max()
            # 记录直方图数据（采样，否则数据太大）
            flattened = output_np.flatten()
            sampled = np.random.choice(flattened, size=min(10000, len(flattened)), replace=False)
            activation_stats[name].append({
                'abs_max': abs_max_val,
                'samples': sampled
            })
        return output
    return hook

# 注册钩子到感兴趣的层
target_layers = ['model.10.cv2.act', 'model.13.cv3.act', 'model.16.m.0', 'model.23.dfl.conv'] # 举例，需要根据你的模型实际层名调整
handles = []
for name, module in fp16_model.named_modules():
    if any(target in name for target in target_layers):
        hook = get_activation_hook(name)
        handles.append(module.register_forward_hook(hook))

# 运行一些数据
with torch.no_grad():
    for i, batch in enumerate(calibration_dataloader):
        if i >= 5: # 跑5个batch足够分析
            break
        batch = batch.cuda().half()
        _ = fp16_model(batch)

# 移除钩子
for h in handles:
    h.remove()

# 分析并绘图
for layer_name, stats_list in activation_stats.items():
    if stats_list:
        all_samples = np.concatenate([s['samples'] for s in stats_list])
        abs_max_vals = [s['abs_max'] for s in stats_list]
        avg_abs_max = np.mean(abs_max_vals)
        print(f"层 {layer_name}: 平均绝对最大值 = {avg_abs_max:.6f}")
        plt.figure()
        plt.hist(all_samples, bins=100, alpha=0.7, label=layer_name)
        plt.title(f'Activation Distribution: {layer_name}')
        plt.xlabel('Value')
        plt.ylabel('Frequency')
        plt.legend()
        plt.savefig(f'hist_{layer_name.replace(".", "_")}.png')
        plt.close()

直方图一出来，问题一目了然。像SPPF后的Concat层，输入来源的几个分支，一个分布集中在[-1, 1]，另一个尾巴拖到了[-10, 10]。而某些SiLU激活后的层，分布极度不对称，大量值集中在0附近，但有很长的正尾巴。标准的对称INT8量化（-127, 127）对于这种分布是极大的浪费，并且会因截断长尾而丢失重要信息。**量化不是均匀的，模型的不同部位对量化的敏感度天差地别**。我得放弃“一刀切”的量化策略了。

分层量化：我放弃了全局最优，换来了整体精度

既然找到了“病灶”，治疗方案就清晰了：对不同的层使用不同的量化策略，甚至不量化。TensorRT的`Layer-wise Precision`设置允许我这么做。我的策略是：

1. 敏感层保持FP16：对数值范围动态大、分布不对称、且对最终检测精度影响直接的层，强制其运行在FP16精度下。这包括：
   • 所有检测头（Head）部分的卷积层。
   • SPPF模块中的Concat层及其紧邻的卷积层。
   • 某些深度可分离卷积（如果发现其量化后误差大）。
2. 其余层使用INT8：模型大部分的骨干网络（Backbone）和特征金字塔网络（Neck）中的常规卷积层，它们的激活分布相对规整，适合INT8量化，能带来主要的加速收益。

在TensorRT Python API中，可以通过设置每层的精度来实现：

# 在构建引擎的代码中，设置分层精度策略
for i in range(network.num_layers):
    layer = network.get_layer(i)
    # 获取层名
    layer_name = layer.name
    # 根据诊断结果，决定该层的精度
    if any(sensitive_keyword in layer_name for sensitive_keyword in ['detect', 'dfl', 'concat', 'model.23', 'model.24', 'model.25']):
        # 关键层设置为FP16，避免量化
        if layer.type == trt.LayerType.CONCATENATION or layer.type == trt.LayerType.CONVOLUTION:
            layer.precision = trt.DataType.HALF
            # 设置输出类型也为HALF，很重要！
            for j in range(layer.num_outputs):
                output_tensor = layer.get_output(j)
                output_tensor.dtype = trt.DataType.HALF
            print(f"Layer {layer_name} set to FP16")
    # 其他层，默认走INT8（由BuilderFlag.INT8控制）

# 注意：还需要告诉config，我们允许不同精度混合
config.set_flag(trt.BuilderFlag.STRICT_TYPES) # 这个标志有时需要，确保遵守精度设置

但实际操作起来，我发现直接通过`layer.precision`设置有时不生效，或者导致引擎构建失败。更稳定、更直观的方式是在构建引擎序列化之前，使用`trt.ILayer`的`set_output_type`方法，或者更高级的，使用TensorRT的`NetworkDefinition` API在解析ONNX后直接修改张量的数据类型。不过，对于这种精细控制，我最终发现了一个更强大的方法：**使用ONNX GraphSurgeon在生成ONNX模型前就标记出需要保持高精度的节点**。

import onnx_graphsurgeon as gs
import onnx

# 加载原始ONNX模型
graph = gs.import_onnx(onnx.load("yolov8s.onnx"))

# 找到需要保持FP16的节点（通过之前的诊断确定其输出张量名）
tensors_to_keep_fp16 = [
    "/model.22/Concat_output_0",  # SPPF后的Concat
    "/model.23/Conv_output_0",     # Detect head 1
    "/model.24/Conv_output_0",     # Detect head 2
    "/model.25/Conv_output_0",     # Detect head 3
    "/model.25/Reshape_output_0",  # DFL层输出
]

for tensor_name in tensors_to_keep_fp16:
    tensor = graph.tensors().get(tensor_name)
    if tensor:
        tensor.dtype = np.float16  # 将张量数据类型标记为float16
        print(f"Marked tensor {tensor_name} as FP16")

# 将修改后的图导回ONNX
graph.cleanup()
onnx.save(gs.export_onnx(graph), "yolov8s_mixed.onnx")

然后，我用这个`yolov8s_mixed.onnx`去构建TensorRT引擎，并且只开启FP16模式，不开启INT8模式。等等，这不是回到原点了吗？别急，我的目的是让这些关键层在ONNX中被标记为FP16，然后构建时再开启INT8。但TensorRT的构建器看到这些节点的输出类型是FP16，就会倾向于为它们选择FP16内核，而其他未标记的节点，在`BuilderFlag.INT8`和`BuilderFlag.FP16`同时开启的情况下，TensorRT的优化器会尽可能为它们选择INT8内核。

我重新构建了引擎。这次构建日志里能看到很多层被标记为`[Float16]`，而更多层被标记为`[Int8]`。跑一下性能，推理时间从纯INT8的22ms增加到了28ms，但比纯FP16的45ms还是快了一大截。最关键的是，精度测试结果出来了：mAP@0.5从84.3%回升到了**89.7%**！牺牲了6ms的推理时间，换回了5.4个百分点的精度。这个交易，我觉得值。

混合精度部署：让关键的层保持FP16，其他层INT8躺平

分层量化策略生效了，但89.7%离目标92.1%还有2.4%的差距。这最后的2-3个百分点，往往最难啃。我分析剩下的误差来源：

1. 即便保持了关键层的FP16，但它们的输入可能来自已经被INT8量化的前层。**量化误差会累积和传播**。一个FP16层接收一个带有INT8量化噪声的输入，输出质量还是会打折扣。
2. 我选择的“敏感层”可能不全，或者有些层虽然本身数值分布适合量化，但它对下游影响巨大。

我需要更系统地评估每层对量化误差的敏感度。手动看直方图太累了。我写了一个简单的“层消融”脚本，在PyTorch模拟量化环境下，逐层冻结其量化（即保持FP16），然后评估整体精度变化。

import torch
import torch.nn as nn
from pytorch_quantization import quant_modules, calib
from pytorch_quantization import nn as quant_nn
from pytorch_quantization.tensor_quant import QuantDescriptor

# 启用量化模块替换
quant_modules.initialize()

# 重新加载模型，此时所有Conv、Linear等层已被替换为量化版本
quant_model = Model('yolov8s.yaml').cuda()
quant_model.eval()

# 设置默认的量化器配置
quant_desc_input = QuantDescriptor(num_bits=8, calib_method='histogram')
quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)

# 收集所有可量化层
quant_layers = []
for name, module in quant_model.named_modules():
    if isinstance(module, (quant_nn.QuantConv2d, quant_nn.QuantLinear)):
        quant_layers.append((name, module))

print(f"Found {len(quant_layers)} quantizable layers")

# 第一步：用校准数据跑一遍，收集所有层的缩放因子
with torch.no_grad():
    for data in calib_dataloader:
        data = data.cuda()
        _ = quant_model(data)
    # 使用直方图校准器计算缩放因子
    for name, module in quant_layers:
        if hasattr(module, '_input_quantizer'):
            module._input_quantizer.load_calib_amax()

# 第二步：定义评估函数
def evaluate_quant_model(model, eval_loader):
    # ... 你的评估代码，返回mAP
    pass

# 基准：全部INT8量化后的精度
full_quant_accuracy = evaluate_quant_model(quant_model, eval_loader)
print(f"Full INT8 accuracy: {full_quant_accuracy:.3f}")

# 第三步：逐层“解冻”，评估精度提升
sensitivity_dict = {}
for idx, (layer_name, layer) in enumerate(quant_layers):
    # 1. 重置所有层为量化状态
    for _, l in quant_layers:
        l.disable_quant()  # 先禁用
        l.enable_quant()   # 再启用，恢复到默认量化状态
    # 2. 只禁用当前层的量化
    layer.disable_quant() # 让这一层跑FP16
    print(f"Testing without quantizing {layer_name}...")
    acc = evaluate_quant_model(quant_model, eval_loader)
    improvement = acc - full_quant_accuracy
    sensitivity_dict[layer_name] = improvement
    print(f"  -> Improvement: {improvement:.4f}")

# 按提升幅度排序
sorted_sensitivity = sorted(sensitivity_dict.items(), key=lambda x: x[1], reverse=True)
print("nTop 10 most sensitive layers to quantization:")
for name, imp in sorted_sensitivity[:10]:
    print(f"  {name}: +{imp:.4f}")

跑这个脚本很耗时，但结果极具价值。它给了我一个量化的“敏感度排行榜”。我发现，除了我之前手动找出的那些层，还有几个位于Neck部分的C2f模块中的卷积层，对量化也异常敏感。我根据这个排行榜，将提升幅度最大的前15%的层（大约十几层）加入到我的“FP16保护名单”中。

更新ONNX标记，重新构建混合精度引擎。这次推理时间稍微增加到31ms，但精度冲到了**91.2%**！距离原始精度只差0.9%了。性能（31ms）也满足了客户33ms的实时性要求。我已经可以交差了，但心里那点工程师的偏执让我想看看，能不能把那最后的0.9%也给磨回来。

最后的1%：用QAT微调把精度“磨”回来

PTQ（训练后量化）的极限似乎就在这里了。要找回那最后的1%，甚至超越，就得请出更强大的武器：**QAT（量化感知训练）**。QAT的原理是在训练（或者说微调）过程中，就模拟量化的效果，让模型权重去适应这种数值上的截断和舍入，从而在真正量化时损失最小。

这需要将模型转换成QAT模式，准备一批标注数据，进行少量epoch的微调。听起来麻烦，但对于一个已经接近目标的项目来说，这可能是临门一脚。我决定用PyTorch的量化工具套件`torch.ao.quantization`（老版本是`torch.quantization`）来试一下。

import torch
from torch.ao.quantization import QuantStub, DeQuantStub, default_qconfig, get_default_qconfig_mapping, prepare_qat, convert
import copy

# 1. 加载预训练的FP32模型
fp32_model = Model('yolov8s.yaml').cuda()
checkpoint = torch.load('yolov8s.pt')
fp32_model.load_state_dict(checkpoint['model'].state_dict())
fp32_model.eval()

# 2. 插入量化存根（QuantStub）和反量化存根（DeQuantStub）
# 需要修改模型定义，在forward开始和结束处加入。这里假设我们修改了模型类。
class QATReadyModel(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.quant = QuantStub()
        self.model = original_model
        self.dequant = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.model(x)
        x = self.dequant(x)
        return x

qat_model = QATReadyModel(fp32_model).cuda()

# 3. 配置量化方式
# 使用适合移动端/边缘设备的配置（支持INT8权重和激活）
qconfig_mapping = get_default_qconfig_mapping('qnnpack')  # 对于ARM CPU，对于NVIDIA GPU可以用 'x86' 或 自定义
# 对于TensorRT，我们更关心对称量化，可以自定义qconfig
from torch.ao.quantization.qconfig import QConfig, default_per_channel_weight_observer, default_histogram_observer
tensorrt_qconfig = QConfig(
    activation=default_histogram_observer.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric),
    weight=default_per_channel_weight_observer.with_args(dtype=torch.qint8, qscheme=torch.per_channel_symmetric)
)
qconfig_mapping = QConfigMapping().set_global(tensorrt_qconfig)

# 4. 准备QAT模型
# 首先融合一些常见的层组合（如Conv+BN+ReLU）以获得更好的精度和性能
from torch.ao.quantization.fuse_modules import fuse_modules
# 这里需要根据你的模型结构手动或自动执行融合，是个细活，省略具体代码...
# fused_model = fuse_modules(qat_model, ...)

prepared_model = prepare_qat(qat_model, qconfig_mapping)

# 5. 进行量化感知训练（微调）
prepared_model.train()
optimizer = torch.optim.SGD(prepared_model.parameters(), lr=1e-4, momentum=0.9)
criterion = ... # 你的损失函数，例如YOLO的复合损失

for epoch in range(5):  # 微调5个epoch足够了
    for batch_idx, (images, targets) in enumerate(train_dataloader):
        images, targets = images.cuda(), targets.cuda()
        optimizer.zero_grad()
        outputs = prepared_model(images)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch}, Loss: {loss.item()}")

# 6. 转换为量化模型
quantized_model = convert(prepared_model)
quantized_model.eval()

# 7. 导出为ONNX（带量化节点）
# 注意：导出的ONNX包含了QuantizeLinear和DequantizeLinear节点，TensorRT可以直接识别并优化
dummy_input = torch.randn(1, 3, 640, 640).cuda()
torch.onnx.export(quantized_model, dummy_input, "yolov8s_qat.onnx",
                  opset_version=13,  # 需要opset 13或更高以支持量化算子
                  input_names=['input'],
                  output_names=['output'],
                  dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})

这个过程我踩了一个大坑：PyTorch默认的量化配置（如‘qnnpack’）和TensorRT期望的量化方式（对称量化、每通道权重量化）可能不完全一致。直接导出的ONNX，TensorRT可能无法最优处理，甚至构建失败。我花了半天时间调整QConfig，确保其与TensorRT的INT8量化方案对齐（对称、每通道权重量化）。

最终，我用QAT微调后的模型，导出了包含量化节点的ONNX，再用TensorRT构建（此时不需要额外的校准器了，因为缩放因子已在模型中）。构建出的引擎，推理时间稳定在30ms左右。激动人心的精度测试：mAP@0.5达到了**91.5%**！虽然还是没有完全达到原始的92.1%，但只损失了0.6个百分点，换来了推理速度从45ms到30ms的**33%** 提升。这个结果，我和客户都完全能接受。

三天时间，从精度暴跌的绝望，到一步步分析、实验、优化，最终拿出一个在速度和精度上都堪用的方案。这次经历让我彻底明白，INT8量化从来不是点一下按钮就完事的魔法，而是一个需要仔细权衡、深度调试的工程过程。模型结构、数据分布、硬件特性，任何一个环节掉链子，都可能让“速度翻倍”的美梦变成“精度崩盘”的噩梦。下次再有人跟我吹嘘量化多简单，我非得把这篇文章糊他脸上不可。

import torch
from torch.quantization import quantize_fx
import copy

def sensitivity_analysis(model, calibration_data, eval_func):
    """
    简单的敏感层分析
    model: 原始FP32模型
    calibration_data: 用于校准的数据加载器
    eval_func: 评估函数，返回精度指标（如mAP）
    """
    base_score = eval_func(model)
    print(f"Base FP32 Score: {base_score:.4f}")

    # 定义要测试的层类型
    layer_types = [torch.nn.Conv2d, torch.nn.Linear, torch.nn.BatchNorm2d]
    type_names = ['Conv2d', 'Linear', 'BatchNorm2d']

    for layer_type, name in zip(layer_types, type_names):
        print(f"n--- 仅量化 {name} 层 ---")
        model_copy = copy.deepcopy(model)

        # 准备量化配置：只量化目标层
        qconfig_dict = {
            '': None,  # 默认不量化
        }
        # 为目标层类型设置量化配置
        for module_name, module in model_copy.named_modules():
            if isinstance(module, layer_type):
                # 这里使用默认的INT8量化配置
                qconfig_dict[module_name] = torch.quantization.get_default_qconfig('fbgemm')  # 或 'qnnpack'

        # 准备模型
        model_prepared = quantize_fx.prepare_fx(model_copy, qconfig_dict, example_inputs=torch.randn(1, 3, 640, 640))

        # 校准（传入少量数据）
        with torch.no_grad():
            for i, data in enumerate(calibration_data):
                if i > 50:  # 用50个batch校准
                    break
                _ = model_prepared(data[0].to('cuda') if torch.cuda.is_available() else data[0])

        # 转换
        model_quantized = quantize_fx.convert_fx(model_prepared)

        # 评估
        quant_score = eval_func(model_quantized)
        drop = base_score - quant_score
        print(f"Quantized Score: {quant_score:.4f}, Drop: {drop:.4f} ({drop/base_score*100:.2f}%)")

# 一个更精细的分层量化配置示例
qconfig_dict = {
    # 默认配置：不量化（作为安全底线）
    '': None,

    # 为所有 Conv2d 设置默认量化，但后面会被更具体的规则覆盖
    'object_type': [(torch.nn.Conv2d, torch.quantization.get_default_qconfig('qnnpack'))],

    # 特定模块不量化：检测头的最终输出层
    'module_name': [
        ('model.22.dfl', None),  # DFL 层，对坐标回归关键，保持FP32
        ('model.22.cv3', None),  # 检测头的最后一个卷积
    ],

    # 特定模块使用动态量化（对权重敏感但激活值范围变化大的层）
    'module_name': [
        ('model.10.cv2', torch.quantization.default_dynamic_qconfig),  # Neck中的某个层
    ],

    # 针对某一类操作（函数）的配置（如果需要）
    # 'function_type': [(torch.nn.functional.interpolate, custom_qconfig_for_upsample)],
}

# 注意：上面的字典在真实使用中需要合并同类项，这里仅为展示思路。
# 实际应用时，需要遍历模型结构，精心构造一个统一的qconfig_dict。

from torch.ao.quantization import QuantStub, DeQuantStub, default_qat_qconfig_mapping, QConfigMapping
from torch.ao.quantization.quantize_fx import prepare_qat_fx, convert_fx
import copy

def prepare_model_for_qat(model_fp32, example_inputs, qconfig_dict):
    """
    准备模型进行QAT
    """
    # 确保模型处于训练模式
    model_fp32.train()

    # 使用prepare_qat_fx插入伪量化节点
    # 注意：这里传入的qconfig_dict决定了哪些层会被插入伪量化节点（值为None的层不会）
    model_prepared = prepare_qat_fx(model_fp32, qconfig_dict, example_inputs=example_inputs)

    # 一个重要的技巧：对于检测模型，损失计算需要在反量化之后进行。
    # 我们通常会在模型末尾手动添加一个DeQuantStub，确保损失计算在FP32域。
    # 但FX模式会自动处理输入输出的量化/反量化，我们需要检查生成的图。
    print("准备完成。模型结构已改变，插入了观察者和伪量化模块。")
    # 可以打印一下模型图看看
    # print(model_prepared.graph)
    return model_prepared

def qat_finetune(model_prepared, train_loader, val_loader, optimizer, scheduler, epochs=25):
    """
    QAT微调循环
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model_prepared.to(device)

    # 假设我们有一个复合损失函数criterion
    # criterion = YOLOLoss(...)

    for epoch in range(epochs):
        model_prepared.train()
        running_loss = 0.0

        # 阶段判断，调整损失权重或学习率
        if epoch < 5:
            # 预热阶段
            box_loss_weight = 0.2  # 降低框损失权重
            lr_scale = 0.1
        elif epoch < 20:
            # 主微调阶段
            box_loss_weight = 1.0
            lr_scale = 1.0
        else:
            # 收敛阶段
            box_loss_weight = 1.0
            lr_scale = 0.1

        adjust_learning_rate(optimizer, base_lr * lr_scale)

        for batch_idx, (images, targets) in enumerate(train_loader):
            images = images.to(device)
            targets = targets.to(device)

            optimizer.zero_grad()

            # 前向传播：经过插入了伪量化节点的模型
            outputs = model_prepared(images)

            # 计算损失，这里需要根据你的YOLO实现来调用
            # loss, loss_components = criterion(outputs, targets)
            # total_loss = loss_components['box']*box_loss_weight + loss_components['cls'] + ...

            # total_loss.backward()
            # optimizer.step()

            running_loss += total_loss.item()

            if batch_idx % 50 == 0:
                print(f'Epoch: {epoch} | Batch: {batch_idx} | Loss: {total_loss.item():.4f}')

        # 每个epoch结束后，在验证集上评估
        val_map = evaluate_on_val(model_prepared, val_loader, device)
        print(f'Epoch {epoch} finished. Train Loss: {running_loss/len(train_loader):.4f}, Val mAP: {val_map:.4f}')

        scheduler.step(val_map)  # 根据验证集性能调整学习率

    return model_prepared

# 在准备模型前，可以先手动融合模块（FX的prepare_qat_fx可能会自动做，但最好确认）
def fuse_yolo_modules(model):
    """
    手动融合YOLOv8中的 Conv + BN + SiLU 模块。
    注意：这是一个示例，需要根据实际模型结构进行递归遍历。
    """
    torch.ao.quantization.fuse_modules(model, [['conv', 'bn', 'silu']], inplace=True)
    # 对于没有BN或激活函数的层，需要其他融合模式
    # 例如: torch.ao.quantization.fuse_modules(model, [['conv', 'relu']], inplace=True)
    return model

# 在调用prepare_qat_fx之前，可以先尝试融合
# model_fp32_fused = fuse_yolo_modules(copy.deepcopy(original_model))

def convert_to_int8(model_qat):
    """
    将QAT模型转换为真正的INT8模型
    """
    model_qat.eval()  # 转换必须在eval模式下进行
    model_int8 = convert_fx(model_qat)
    return model_int8

# 转换
final_int8_model = convert_to_int8(model_qat_finetuned)

# 保存模型
torch.jit.save(torch.jit.script(final_int8_model), 'yolov8s_int8_qat.pt')

def generate_qconfig_profiles(model, sensitivity_report, protection_levels=['high', 'medium', 'low']):
    """
    根据敏感层报告，生成不同保护级别的量化配置。
    sensitivity_report: 字典，key为模块名，value为量化该模块导致的精度下降值。
    protection_levels: 定义哪些下降值范围对应高、中、低敏感度。
    """
    profiles = {}
    for level in protection_levels:
        qconfig_dict = {'': None}
        for name, drop in sensitivity_report.items():
            if level == 'high' and drop > 0.05:  # 高敏感，不量化
                qconfig_dict[name] = None
            elif level == 'medium' and drop > 0.02:  # 中敏感，使用动态量化或更宽范围
                qconfig_dict[name] = get_conservative_qconfig()
            else:  # 低敏感，默认量化
                qconfig_dict.setdefault('object_type', []).append((get_module_type_by_name(model, name), default_qconfig))
        profiles[level] = qconfig_dict
    return profiles

.content {
font-family: -apple-system, BlinkMacSystemFont, “Segoe UI”, Roboto, “Helvetica Neue”, Arial, sans-serif;
line-height: 1.8;
color: #333;
max-width: 900px;
margin: 0 auto;
padding: 20px;
}
h2 {
color: #2c3e50;
border-bottom: 2px solid #eee;
padding-bottom: 10px;
margin-top: 40px;
}
p {
margin-bottom: 1.2em;
text-align: justify;
}
strong {
color: #e74c3c;
}
pre {
background-color: #f8f9fa;
border: 1px solid #e9ecef;
border-radius: 6px;
padding: 16px;
overflow: auto;
margin: 20px 0;
}
code {
font-family: ‘SFMono-Regular’, Consolas, ‘Liberation Mono’, Menlo, Courier, monospace;
font-size: 0.9em;
}
.language-python, .language-bash {
display: block;
}
ul, ol {
margin-bottom: 1.2em;
padding-left: 2em;
}
li {
margin-bottom: 0.5em;
}