YOLOV7量化第二步: 模型标定

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()  # 如果没有校准器，我们只需重新启用量化功能。

    # 在此阶段，校准过程完成，模型已经准备好以量化的状态进行更高效的运行。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228