2.模型标定

当然可以，模型量化中的标定（calibration）是一个关键过程，它主要确保在降低计算精度以减少模型大小和提高推理速度的同时，不会显著损害模型的准确性。现在，我将根据您提供的步骤解释这一过程。

1. 收集网络层的输入/输出信息

首先，我们需要通过运行模型（使用标定数据集，而不是训练数据或测试数据）来收集关于每层的输入和输出的信息。这个数据集应该是多样化的，以便涵盖到可能的各种情况。

在这一步中，模型是在推理模式下运行的，所有层的输出都被记录下来。这通常是通过修改模型的代码来实现的，以便在每个层之后捕获并存储激活的分布。这些数据将用于下一步中的统计分析。

具体实施时，这一步可能涉及编写一个循环，该循环遍历标定数据集的每个样本，并逐一通过模型。在每一层，您需要捕获并可能临时存储输入和输出数据（通常是张量的形式）。

def collect_stats(model, data_loader, device, num_batch=200):
    model.eval()  # 将模型设置为评估（推理）模式。这在PyTorch中很重要，因为某些层（如Dropout和BatchNorm）在训练和评估时有不同的行为。

    # 开启校准器
    for name, module in model.named_modules():  # 遍历模型中的所有模块。`named_modules()`方法提供了一个迭代器，按层次结构列出模型的所有模块及其名称。
        if isinstance(module, quant_nn.TensorQuantizer):  # 检查当前模块是否为TensorQuantizer类型，即我们想要量化的特定类型的层。
            if module._calibrator is not None:  # 如果此层配备了校准器。
                module.disable_quant()  # 禁用量化。这意味着层将正常（未量化）运行，使校准器能够收集必要的统计数据。
                module.enable_calib()  # 启用校准。这使得校准器开始在此层的操作期间收集数据。
            else:
                module.disable()  # 如果没有校准器，简单地禁用量化功能，但不进行数据收集。

    # 在此阶段，模型准备好接收数据，并通过处理未量化的数据来进行校准。
    
    # test
    with torch.no_grad():  # 关闭自动求导系统。这在进行推理时是有用的，因为它减少了内存使用量，加速了计算，而且我们不需要进行反向传播。
        for i, datas in enumerate(data_loader):  # 遍历数据加载器。数据加载器将提供批量的数据，通常用于训练或评估。
            imgs = datas[0].to(device, non_blocking=True).float()/255.0  # 获取图像数据，转换为适当的设备（例如GPU），并将其类型转换为float。除以255是常见的归一化技术，用于将像素值缩放到0到1的范围。
            model(imgs)  # 用当前批次的图像数据执行模型推理。

            if i >= num_batch:  # 如果我们已经处理了指定数量的批次，则停止迭代。
                break

    # 关闭校准器
    for name, module in model.named_modules():  # 再次遍历所有模块，就像我们之前做的那样。
        if isinstance(module, quant_nn.TensorQuantizer):  # 对于TensorQuantizer类型的模块。
            if module._calibrator is not None:  # 如果有校准器。
                module.enable_quant()  # 重新启用量化。现在，校准器已经收集了足够的统计数据，我们可以再次量化层的操作。
                module.disable_calib()  # 禁用校准。数据收集已经完成，因此我们关闭校准器。
            else:
                module.enable()  # 如果没有校准器，我们只需重新启用量化功能。

    # 在此阶段，校准过程完成，模型已经准备好以量化的状态进行更高效的运行。

2. 计算动态范围和比例因子

一旦我们收集了各层的激活数据，接下来的步骤是分析这些数据来确定量化参数，即动态范围（也称为量化范围）和比例因子（scale）。

• 动态范围是指在量化过程中，张量数据可以扩展到的范围。它是原始数据的最大值和最小值之间的差值。这个范围很重要，因为我们希望我们的量化表示能够覆盖可能的所有值，从而避免饱和和信息丢失。

这个过程中，method: A string. One of [‘entropy’, ‘mse’, ‘percentile’] 我们有三种办法，这个实际上要在做实验的时候看哪一个精度更高，这个就是看map值计算的区别

def compute_amax(model, device, **kwargs):
    # 遍历模型中的所有模块，`model.named_modules()`方法提供了一个迭代器，包含模型中所有模块的名称和模块本身。
    for name, module in model.named_modules():
        # 检查当前模块是否为TensorQuantizer的实例，这是处理量化的部分。
        if isinstance(module, quant_nn.TensorQuantizer):
            # （这里的print语句已被注释掉，如果取消注释，它将打印当前处理的模块的名称。）
            # print(name)

            # 检查当前的量化模块是否具有校准器。
            if module._calibrator is not None:
                # 如果该模块的校准器是MaxCalibrator的实例（一种特定类型的校准器）...
                if isinstance(module._calibrator, calib.MaxCalibrator):
                    # ...则调用load_calib_amax()方法，该方法计算并加载适当的'amax'值，它是量化过程中用于缩放的最大激活值。
                    module.load_calib_amax()
                else:
                    # ...如果校准器不是MaxCalibrator，我们仍然调用load_calib_amax方法，但是可以传递额外的关键字参数。
                    # 这些参数可能会影响'amax'值的计算。
                    module.load_calib_amax(**kwargs)  # ['entropy', 'mse', 'percentile']   这里有三个计算方法，实际过程中要看哪一个比较准，再考虑用哪一个
                # 将计算出的'amax'值（现在存储在模块的'_amax'属性中）转移到指定的设备上。
                # 这确保了与模型数据在同一设备上的'amax'值，这对于后续的计算步骤（如训练或推理）至关重要。
                module._amax = module._amax.to(device)

Scanning '/app/dataset/coco2017/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupted: 100%|███████████████████████| 5000/5000 [00:00<?, ?it/s]
Origin pth_Model map: 
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100%|█████████████████████████████████████| 625/625 [00:49<00:00, 12.61it/s]
                 all        5000       36781       0.717       0.626       0.675       0.454
Fusing layers... 
RepConv.fuse_repvgg_block
RepConv.fuse_repvgg_block
RepConv.fuse_repvgg_block
IDetect.fuse
QDQ auto init map: 
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100%|█████████████████████████████████████| 625/625 [00:39<00:00, 15.78it/s]
                 all        5000       36781       0.718       0.627       0.676       0.455
Calibrate Model map: 
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100%|█████████████████████████████████████| 625/625 [01:42<00:00,  6.09it/s]
                 all        5000       36781        0.73       0.618       0.674       0.454

2.3 完整代码

import torch
from pytorch_quantization import quant_modules 
from models.yolo import Model
from pytorch_quantization.nn.modules import _utils as quant_nn_utils
from pytorch_quantization import calib
import sys
import re
import yaml
import os
os.chdir("/app/bob/yolov7_QAT/yolov7")



def load_yolov7_model(weight, device="cpu"):
    ckpt = torch.load(weight, map_location=device)                                  # 加载模型，模型参数在哪个设备上
    model = Model("cfg/training/yolov7.yaml", ch=3, nc=80).to(device)               # 跟yolov7的结构，这里没有包含参数
    state_dict = ckpt["model"].float().state_dict()                                 # 从加载的权重中提取模型的状态字典(state_dict), 包含了模型全部的参数，包括卷积权重等
    model.load_state_dict(state_dict, strict=False)                                 # 把提取出来的参数放到yolov7的结构中
    return model                                                                    # 返回正确权重和参数的模型

import collections
from utils.datasets import create_dataloader
def prepare_dataset(cocodir, batch_size=8):   
    dataloader = create_dataloader(         # 这里的参数是跟官网的是一样的
        f"{cocodir}/val2017.txt",
        imgsz=640,
        batch_size=batch_size,
        opt=collections.namedtuple("Opt", "single_cls")(False),  # collections.namedtuple("Opt", "single_cls")(False)
        augment=False, hyp=None, rect=True, cache=False, stride=32, pad=0.5, image_weights=False,
        )[0]
    return dataloader

import test as test 
from pathlib import Path
import os
def evaluate_coco(model, loader, save_dir='.', conf_thres=0.001, iou_thres=0.65):
    
    if save_dir and os.path.dirname(save_dir) != "":
        os.makedirs(os.path.dirname(save_dir), exist_ok=True)

    return test.test(
        "./data/coco.yaml",
        save_dir=Path(save_dir),
        conf_thres=conf_thres,
        iou_thres=iou_thres,
        model=model,
        dataloader=loader,
        is_coco=True,
        plots=False,
        half_precision=True,
        save_json=False
    )[0][3]   
    
      
from pytorch_quantization import nn as quant_nn
from pytorch_quantization.tensor_quant import QuantDescriptor
from absl import logging as quant_logging
# intput QuantDescriptor: Max ==> Histogram
def initialize():
    quant_desc_input = QuantDescriptor(calib_method="histogram")   # "max" 
    quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
    quant_nn.QuantMaxPool2d.set_default_quant_desc_input(quant_desc_input)
    quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
    quant_logging.set_verbosity(quant_logging.ERROR) 
    
def prepare_model(weight, device):
    # quant_modules.initialize()   # 自动加载qdq节点
    initialize()                 # intput QuantDescriptor: Max ==> Histogram
    model = load_yolov7_model(weight, device)
    model.float()
    model.eval()
    with torch.no_grad():
        model.fuse()  # conv bn 进行层的合并, 加速
    return model

# 执行量化替换
def transfer_torch_to_quantization(nn_instance, quant_mudule):
    quant_instance = quant_mudule.__new__(quant_mudule)
    for k, val in vars(nn_instance).items():
        setattr(quant_instance, k, val)

    def __init__(self):
        # 返回两个QuantDescriptor的实例    self.__class__是quant_instance的类, EX: QuantConv2d
        quant_desc_input, quant_desc_weight = quant_nn_utils.pop_quant_desc_in_kwargs(self.__class__)
        if isinstance(self, quant_nn_utils.QuantInputMixin):
            self.init_quantizer(quant_desc_input)
            if isinstance(self._input_quantizer._calibrator, calib.HistogramCalibrator):
                self._input_quantizer._calibrator._torch_hist = True
        else:
            self.init_quantizer(quant_desc_input, quant_desc_weight)

            if isinstance(self._input_quantizer._calibrator, calib.HistogramCalibrator):
                self._input_quantizer._calibrator._torch_hist = True
                self._weight_quantizer._calibrator._torch_hist = True

    __init__(quant_instance)
    return quant_instance

def quantization_ignore_match(ignore_layer, path):
    if ignore_layer is None:
        return False
    if isinstance(ignore_layer, str) or isinstance(ignore_layer, list):
        if isinstance(ignore_layer, str):
            ignore_layer = [ignore_layer]
        if path in ignore_layer:
            return True
        for item in ignore_layer:
            if re.match(item, path):  
                return True  
    return False

# 递归函数
def torch_module_find_quant_module(module, module_dict, ignore_layer, prefix=''):
    for name in module._modules:
        submodule = module._modules[name]
        path =  name if prefix == '' else prefix + '.' + name
        torch_module_find_quant_module(submodule, module_dict, ignore_layer, prefix=path)

        submodule_id = id(type(submodule))
        if submodule_id in module_dict:
            ignored = quantization_ignore_match(ignore_layer, path)
            if ignored:
                print(f"Quantization : {path} has ignored.")
                continue
            # 转换
            module._modules[name] = transfer_torch_to_quantization(submodule, module_dict[submodule_id])

# 用量化模型替换
def replace_to_quantization_model(model, ignore_layer=None):
    """
    这里构建的module_dict里面的元素是一个映射的关系, 例如torch.nn -> quant_nn.QuantConv2d, 一共是15个, 跟DEFAULT_QUANT_MAP对齐
    """
    module_dict = {}
    for entry in quant_modules._DEFAULT_QUANT_MAP:           # 构建module_dict, 把DEFAULT_QUANT_MAP填充
        module = getattr(entry.orig_mod, entry.mod_name)     # 提取的原始的模块，从torch.nn中获取conv2d这个字符串
        module_dict[id(module)] = entry.replace_mod          # 使用替换的模块
    torch_module_find_quant_module(model, module_dict, ignore_layer)
    
def collect_stats(model, data_loader, device, num_batch=200):
    model.eval()  # 将模型设置为评估（推理）模式。这在PyTorch中很重要，因为某些层（如Dropout和BatchNorm）在训练和评估时有不同的行为。

    # 开启校准器
    for name, module in model.named_modules():  # 遍历模型中的所有模块。`named_modules()`方法提供了一个迭代器，按层次结构列出模型的所有模块及其名称。
        if isinstance(module, quant_nn.TensorQuantizer):  # 检查当前模块是否为TensorQuantizer类型，即我们想要量化的特定类型的层。
            if module._calibrator is not None:  # 如果此层配备了校准器。
                module.disable_quant()  # 禁用量化。这意味着层将正常（未量化）运行，使校准器能够收集必要的统计数据。
                module.enable_calib()  # 启用校准。这使得校准器开始在此层的操作期间收集数据。
            else:
                module.disable()  # 如果没有校准器，简单地禁用量化功能，但不进行数据收集。

    # 在此阶段，模型准备好接收数据，并通过处理未量化的数据来进行校准。
    
    # test
    with torch.no_grad():  # 关闭自动求导系统。这在进行推理时是有用的，因为它减少了内存使用量，加速了计算，而且我们不需要进行反向传播。
        for i, datas in enumerate(data_loader):  # 遍历数据加载器。数据加载器将提供批量的数据，通常用于训练或评估。
            imgs = datas[0].to(device, non_blocking=True).float()/255.0  # 获取图像数据，转换为适当的设备（例如GPU），并将其类型转换为float。除以255是常见的归一化技术，用于将像素值缩放到0到1的范围。
            model(imgs)  # 用当前批次的图像数据执行模型推理。

            if i >= num_batch:  # 如果我们已经处理了指定数量的批次，则停止迭代。
                break

    # 关闭校准器
    for name, module in model.named_modules():  # 再次遍历所有模块，就像我们之前做的那样。
        if isinstance(module, quant_nn.TensorQuantizer):  # 对于TensorQuantizer类型的模块。
            if module._calibrator is not None:  # 如果有校准器。
                module.enable_quant()  # 重新启用量化。现在，校准器已经收集了足够的统计数据，我们可以再次量化层的操作。
                module.disable_calib()  # 禁用校准。数据收集已经完成，因此我们关闭校准器。
            else:
                module.enable()  # 如果没有校准器，我们只需重新启用量化功能。

    # 在此阶段，校准过程完成，模型已经准备好以量化的状态进行更高效的运行。

def compute_amax(model, device, **kwargs):
    # 遍历模型中的所有模块，`model.named_modules()`方法提供了一个迭代器，包含模型中所有模块的名称和模块本身。
    for name, module in model.named_modules():
        # 检查当前模块是否为TensorQuantizer的实例，这是处理量化的部分。
        if isinstance(module, quant_nn.TensorQuantizer):
            # （这里的print语句已被注释掉，如果取消注释，它将打印当前处理的模块的名称。）
            # print(name)

            # 检查当前的量化模块是否具有校准器。
            if module._calibrator is not None:
                # 如果该模块的校准器是MaxCalibrator的实例（一种特定类型的校准器）...
                if isinstance(module._calibrator, calib.MaxCalibrator):
                    # ...则调用load_calib_amax()方法，该方法计算并加载适当的'amax'值，它是量化过程中用于缩放的最大激活值。
                    module.load_calib_amax()
                else:
                    # ...如果校准器不是MaxCalibrator，我们仍然调用load_calib_amax方法，但是可以传递额外的关键字参数。
                    # 这些参数可能会影响'amax'值的计算。
                    module.load_calib_amax(**kwargs)  # ['entropy', 'mse', 'percentile']   这里有三个计算方法，实际过程中要看哪一个比较准，再考虑用哪一个
                # 将计算出的'amax'值（现在存储在模块的'_amax'属性中）转移到指定的设备上。
                # 这确保了与模型数据在同一设备上的'amax'值，这对于后续的计算步骤（如训练或推理）至关重要。
                module._amax = module._amax.to(device)

def calibrate_model(model, dataloader, device):
    # 收集信息
    collect_stats(model, dataloader, device)
    # 获取动态范围,计算amax值,scale值
    compute_amax(model, device, method='mse')

if __name__ == "__main__":
    weight = "./yolov7.pt"
    cocodir = "/app/dataset/coco2017"    #../dataset/coco2017
    device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
    
    # load最初版本的模型
    pth_model = load_yolov7_model(weight=weight, device=device)
    # print(pth_model)
    dataloader = prepare_dataset(cocodir=cocodir, )
    print("Origin pth_Model map: ")
    ap = evaluate_coco(pth_model, dataloader)
    
    # 加载自动插入QDQ节点的模型
    # print("Before prepare_model")
    model = prepare_model(weight=weight, device=device)
    # print("After prepare_model")
    print("QDQ auto init map: ")
    qdq_auto_ap = evaluate_coco(model, dataloader)

    # print("Before replace_to_quantization_model")
    replace_to_quantization_model(model)
    # print("After replace_to_quantization_model")

    # print("Before calibrate_model")
    calibrate_model(model, dataloader, device)
    # print("After calibrate_model")
    print("Calibrate Model map: ")
    cali_ap = evaluate_coco(model, dataloader)