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
        )