YOLOV7量化第一步: 准备模型

1. 准备模型手动/自动插入QDQ节点

1.1 准备容器

准备容器的时候有可能出现共享内存的问题, 这里在启动容器的时候设置一下,

# 增加了里面的共享内存
docker run --gpus all -it --shm-size=64g --name env_pyt_1.12 -v $(pwd):/app nvcr.io/nvidia/pytorch:22.03-py3 ```

美式的做法是

.
├── __pycache__
├── my_ptq.py
├── ptq.py
├── qat.py
├── quantize.py
├── rules.py
├── yolov7
└── yolov7.pt

实际做的时候会有路径的问题，最简单的办法就是直接去yolov7文件夹里面写

1.2 加载数据集然后进行一个简单的验证

这里面的create DataLoader是使用的是yolov7里面自带的， 下面是官方的create_dataloader代码

def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False,
                      rank=-1, world_size=1, workers=8, image_weights=False, quad=False, prefix=''):
    # Make sure only the first process in DDP process the dataset first, and the following others can use the cache
    with torch_distributed_zero_first(rank):
        dataset = LoadImagesAndLabels(path, imgsz, batch_size,
                                      augment=augment,  # augment images
                                      hyp=hyp,  # augmentation hyperparameters
                                      rect=rect,  # rectangular training
                                      cache_images=cache,
                                      single_cls=opt.single_cls,
                                      stride=int(stride),
                                      pad=pad,
                                      image_weights=image_weights,
                                      prefix=prefix)

    batch_size = min(batch_size, len(dataset))
    nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, workers])  # number of workers
    sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None
    loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader
    # Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader()
    dataloader = loader(dataset,
                        batch_size=batch_size,
                        num_workers=nw,
                        sampler=sampler,
                        pin_memory=True,
                        collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn)
    return dataloader, dataset

下面是我们自己的代码
注意这里面的op是一个对象，这里用collections.namedtuple("Opt", "single_cls")(False)
是为了设置Opt这个对象里面的"single_cls"是False, 它里面还有其他的参数, 这里因为yolov7里面返回的是两个，第一个是data loader, 第二个是dataset

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),
        augment=False, hyp=None, rect=True, cache=False, stride=32, pad=0.5, image_weights=False
        )[0]
    return dataloader

下面是验证coco的函数, 这里就是把plot关了我们不用可视化的内容，然后官方test的内容第一个和第四个是mp和map, 这两个是我们要的参数，这里的map是0.5-0.95的
在这里插入图片描述

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]

import torch
from pytorch_quantization import quant_modules 
from models.yolo import Model
import sys


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),
        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]


if __name__ == "__main__":
    weight = "../yolov7.pt"
    device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
    
    model = load_yolov7_model(weight=weight, device=device)
    
    cocodir = "/app/dataset/coco2017"    #../dataset/coco2017
    
    dataloader = prepare_dataset(cocodir=cocodir, )
    
    ap = evaluate_coco(model, dataloader)

首先，让我们看一下代码的输出：

Scanning '/app/dataset/coco2017/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupted: 100%|██████████████████████| 5000/5000 [00:00<?, ?it/s]
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100%|████████████████████████████████████| 625/625 [00:40<00:00, 15.29it/s]
                 all        5000       36781       0.718       0.626       0.676       0.452

这个输出提供了您模型在COCO 2017验证数据集上的性能评估。我将为您逐个解释关键指标：

Scanning…:
- 这一行显示了数据加载的状态。4952张图片和标签被找到，48张图片缺失，没有空标签，也没有损坏的标签。
Class, Images, Labels:
- Class: 分类的名称。在这里，“all”表示考虑所有分类。
- Images: 总共测试的图像数量。这里是5000。
- Labels: 在这5000张图片中，总共有36781个对象标签。
P (Precision): 这是模型的精确度，值为0.718。这意味着当模型预测某个对象存在时，它有71.8%的概率是正确的。
R (Recall): 这是模型的召回率，值为0.626。这意味着模型能够正确检测到62.6%的实际存在的对象。
mAP@.5: 这是在IoU（Intersection over Union）阈值为0.5时的平均精度。IoU是预测的边界框和真实边界框之间的交集与并集之比。值为0.676表示在IoU=0.5时，模型的平均精度为67.6%。
mAP@.5:.95: 这是在IoU从0.5到0.95的范围内的平均精度（以0.05为步长）。在COCO数据集上，这是一个非常重要的指标，因为它考虑了不同IoU阈值下的性能。值为0.452表示平均精度为45.2%。

简单来说，mAP（mean Average Precision）是评估物体检测模型性能的关键指标。它考虑了模型的精确度和召回率，并为每个类别计算AP（平均精度），然后取所有类别的AP的平均值得到mAP。

总体上，这些指标为您提供了模型在COCO 2017验证数据集上的性能概览。您可以使用这些指标来比较不同模型或相同模型的不同版本的性能。

1.3 准备自动插入QDQ节点模型

def prepare_model(weight, device):
    quant_modules.initialize()   # 自动加载qdq节点
    # initialize()
    model = load_yolov7_model(weight, device)
    model.float()
    model.eval()
    with torch.no_grad():
        model.fuse()  # conv bn 进行层的合并, 加速
    return model
    
if __name__ == "__main__":
    weight = "../yolov7.pt"
    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)
    
    # 加载自动插入QDQ节点的模型
    model = prepare_model(weight=weight, device=device)
    print(model)

这里对比了一下自动插入的QDQ节点的模型, 通过打印的方式可以看到这里确实是有自动插入的节点

      (cv3): Conv(
        (conv): QuantConv2d(
          512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)
          (_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
          (_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)

下面是自动插入QDQ节点计算出来的map, 可以看出来其实是跟直接计算coco的是一样的，

Scanning '/app/dataset/coco2017/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupted: 100%|██████████████████████████████████████████████████| 5000/5000 [00:00<?, ?it/s]
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100%|████████████████████████████████████████████████████████████████| 625/625 [01:03<00:00,  9.78it/s]
                 all        5000       36781       0.722       0.622       0.675       0.452

1.4 手动插入QDQ节点

1.4.1. 定义模块映射

首先，我们有一个预定义的_DEFAULT_QUANT_MAP列表，它为原始PyTorch模块（如torch.nn.Conv2d）和它们的量化版本（如quant_nn.QuantConv2d）提供了映射关系。

1.4.2. 构建ID到量化模块的字典

在replace_to_quantization_model函数中，我们遍历_DEFAULT_QUANT_MAP列表并使用getattr函数动态地从torch.nn或其他模块中获取对应的类。然后，我们为这些类创建一个ID到量化版本的映射，并将其存储在module_dict字典中。

1.4.3. 递归地遍历模型的子模块

接下来，我们使用torch_module_find_quant_module函数递归地遍历整个模型的子模块。对于每个子模块，我们查看它的ID是否出现在module_dict字典中，以确定是否需要进行量化替换。

1.4.4. 检查忽略列表

在决定是否替换子模块之前，我们首先检查它是否出现在ignore_layer列表中。如果在这个列表中，那么我们不对它进行量化替换。这是通过quantization_ignore_match函数实现的。

1.4.5. 执行量化替换

如果一个子模块需要被量化并且不在忽略列表中，我们使用transfer_torch_to_quantization函数执行替换。这个函数首先为量化模块创建一个新的实例，然后复制原始模块的所有属性到新的量化模块。接着，它使用内部的__init__方法来完成量化初始化过程，包括设置量化描述符和校准器。

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]     
    
def prepare_model(weight, device):
    # quant_modules.initialize()   # 自动加载qdq节点
    # initialize()
    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)
    
if __name__ == "__main__":
    weight = "./yolov7.pt"
    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)
    
    # 加载自动插入QDQ节点的模型
    model = prepare_model(weight=weight, device=device)
    replace_to_quantization_model(model)
    print(model)
    
    cocodir = "/app/dataset/coco2017"    #../dataset/coco2017
    
    dataloader = prepare_dataset(cocodir=cocodir, )
    
    ap = evaluate_coco(pth_model, dataloader)

因为这里没有训练，所以计算出来的map也不会有太大的变化

Scanning '/app/dataset/coco2017/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupted: 100%|█| 5000/5000 [00:00<?, 
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100%|█████████| 625/625 [00:41<00:00, 15.18it/s]
                 all        5000       36781       0.724       0.621       0.675       0.453

同样这个得量化也可以看出来是默认的per-tensor, MaxCalibrator, 后面要手动initiative来设置成Hstgram的量化校准器

      (conv): QuantConv2d(
        1024, 512, kernel_size=(1, 1), stride=(1, 1)
        (_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
        (_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
      )

1.5 手动initlize()

手动初始化插入节点可以更好的控制好整个流程, 这里的做了一件事，就是通过设置参数把校准器从Max换到Hisgram, 这里改的把conv2d, MaxPool2d, Linear, 的input设置成Hisgram, 这里也可以设置成weight的校准器成Histgram

    quant_desc_input = QuantDescriptor(calib_method="histogram")   # "max" 
    quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
    quant_nn.QuantConv2d.set_default_quant_desc_weight(quant_desc_input)

        (2): QuantConv2d(
          1024, 255, kernel_size=(1, 1), stride=(1, 1)
          (_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
          (_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
        )

这种就是设置weight和input的校准器都是直方图校准

        (2): QuantConv2d(
          1024, 255, kernel_size=(1, 1), stride=(1, 1)
          (_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=HistogramCalibrator scale=1.0 quant)
          (_weight_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=HistogramCalibrator scale=1.0 quant)
        )

下面是对比的手动插入QDQ节点的map计算和直接load模型的map, 从这里可以看出来因为没有训练，所以这边都是一样的

Scanning '/app/dataset/coco2017/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupted:
Origin Pth Model map: 
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100%|█| 625/625 [00
                 all        5000       36781       0.724       0.622       0.676       0.453
auto QDQ Model map: 
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100%|█| 625/625 [01
                 all        5000       36781       0.705       0.633       0.672       0.444

完整代码

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")
    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)
    
if __name__ == "__main__":
    weight = "./yolov7.pt"
    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)
    
    # 加载自动插入QDQ节点的模型
    model = prepare_model(weight=weight, device=device)
    replace_to_quantization_model(model)
    print(model)
    
    cocodir = "/app/dataset/coco2017"    #../dataset/coco2017
    
    dataloader = prepare_dataset(cocodir=cocodir, )
    
    print("Origin Pth Model map: ")
    ap = evaluate_coco(pth_model, dataloader)
    
    print("auto QDQ Model map: ")
    auto_ap = evaluate_coco(model, dataloader)