史上最快的实例分割SparseInst Int8量化实录

近期，YOLOv7里面借鉴（复 制 粘 贴）了一个新的模型，SparseInst，我借助YOLOv7的基建能力，将其导出到了ONNX, 获得了一个非常不错的可以直接用OnnxRuntime， 或者TensorRT跑的实例分割 (后续也可能把link加到官方的repo当中)。索性就一不作二不休，把int8也给他加上。于是就有了这个踩坑记录博客。

本文将带你从0开始量化一个复杂网络，这个SparseInst也是基于Transformer的，网络结构够复杂。最终实现Int8的量化推理。

我们会cover的知识点包括：

• 是如何将一个transformer导出到onnx的，中间会遇到哪些问题？
• 如果重写后处理得到一个可用的onnx模型？
• 如何对coco的模型实现普适化的量化（无需重复写乱七八糟的dataloader)
• 如何评估量化模型与float32的误差？
• 如何对量化后的模型进行推理，并通用的迁移到X86, ARM等架构,或者TensorRT， OpenVINO, NCNN等前推框架

先来看看SparseInst的精度。这或许是今年，速度最快，精度最高的实例分割模型：

可以看到，在mAP达到 37.9 的情况下，可以跑到 40 FPS. 

这是我在CPU下，用int8的跑到结果，请注意，是CPU, 大概是 30ms. 小目标检测的效果依旧非常不错。

这还是还没有用TensorRT,也没有用TensorRT的int8加速的情况下，如果TensorRT int8,可以达到4-6ms. 输入尺寸是640x640,不大也不小的图。TensorRT我将在下一篇继续测试。

这个模型，我认为如果你感兴趣的话，是可以在数据集上得到一个更高的精度的。这里面的attention machanism 是可以被改进的，潜力依旧很大。但毫无疑问，这是目前看来，第一个不是特别慢的实例分割的范式，也不是无脑的堆叠encoder layer寻求高精度的方法。能做到非常好的精度与推理速度的权衡。

那我们就来把他量化到int8吧。

0x1 导出onnx

首先导出到onnx ， 这一步可以直接一不到 yolov7仓库，参考对应的导出流程。

https://github.com/jinfagang/yolov7

如果你遇到了什么问题，请尽情的在github提issue.

0x2 量化

这部分坑，实在是太深了，一言难尽。一入量化深似海。所以接下来这部分，我建议你慎重观摩，可能引起你的心理落差。

因为如果你不经过一些非常深的套路，结局可能是这样的：

不管那么多了，早晚是要quant的。那不如我们先来吧calibrator高一波：

import sys
from torchvision import transforms
import torchvision
import torch
from atomquant.onnx.dataloader import get_calib_dataloader_coco
import os
import cv2
import numpy as np
import onnxruntime as ort
from torchvision.datasets.coco import CocoDetection
from alfred.dl.torch.common import device


def preprocess_func(img, target):
    w = 640
    h = 640
    a = cv2.resize(img, (w, h))
    a_t = np.array(a).astype(np.float32)
    boxes = []
    for t in target:
        boxes.append(t["bbox"])
    target = np.array(boxes)
    a_t = torch.as_tensor(a_t)
    target = torch.as_tensor(target)
    return a_t, target


def collate_fn(batch):
    images, targets = zip(*batch)
    if isinstance(images[0], torch.Tensor):
        images = torch.stack(images)
        targets = torch.stack(targets)
    else:
        images = np.array(images)
    return images


if __name__ == "__main__":
    ONNX_PATH = sys.argv[1]

    coco_root = os.path.expanduser("~/data/coco/images/val2017")
    anno_f = os.path.expanduser(
        "~/data/coco/annotations/instances_val2017_val_val_train.json"
    )

    # coco_ds = CocoDetection(coco_root, anno_f, )

    session = ort.InferenceSession(ONNX_PATH)
    input_name = session.get_inputs()[0].name

    calib_dataloader = get_calib_dataloader_coco(
        coco_root,
        anno_f,
        preprocess_func=preprocess_func,
        input_names=input_name,
        bs=1,
        max_step=50,
        collate_fn=collate_fn
    )

这个地方我比较建议你使用pqq来量化。上面的atomquant是我的一个还未开源的包，由于太菜以至于不敢开源。但是我就暂且用这个里面提供的一些calibrator构造函数来丢一波coco的数据来量化。

接下来丢进来我们的onnx模型，就可以开始量化了：

REQUIRE_ANALYSE = False
BATCHSIZE = 1
# INPUT_SHAPE = [3, 224, 224]
INPUT_SHAPE = [640, 640, 3]
DEVICE = "cuda"  # only cuda is fully tested :(  For other executing device there might be bugs.
# PLATFORM = TargetPlatform.PPL_CUDA_INT8  # identify a target platform for your network.
PLATFORM = (
TargetPlatform.ORT_OOS_INT8
# TargetPlatform.PPL_CUDA_INT8
)  # identify a target platform for your network.
# PLATFORM = TargetPlatform.ONNXRUNTIME  # identify a target platform for your network.

EXECUTING_DEVICE = "cpu"  # 'cuda' or 'cpu'.

# create a setting for quantizing your network with PPL CUDA.
# quant_setting = QuantizationSettingFactory.pplcuda_setting()
quant_setting = QuantizationSettingFactory.default_setting()
quant_setting.equalization = True  # use layerwise equalization algorithm.
quant_setting.dispatcher = (
"conservative"  # dispatch this network in conservertive way.
)


# quantize your model.
quantized = quantize_onnx_model(
onnx_import_file=ONNX_PATH,
calib_dataloader=calib_dataloader.dataloader_holder,
calib_steps=88,
input_shape=[BATCHSIZE] + INPUT_SHAPE,
setting=quant_setting,
# collate_fn=collate_fn,
platform=PLATFORM,
device=DEVICE,
verbose=0,
)

记得import你需要的东西，这里默认我们进行x86的计算模拟，因为我们想跑在CPU上，至于GPU,那是下一篇的事情。

然后我们就可以得到一个quantize的模型：

在这里面，我们展示的是一个实例分割模型，这里面包含了非常多的复杂操作，例如各种shape的组合，以及各种concat,各种interpolate, 其中很多算子是没有办法去量化的，至少很多前推引擎并不支持。

但是我们不管那么多，一顿梭哈，无脑梭哈。

然后我们就可以得到这么一个int8的模型：

模型体积从 140M -> 50M.

好了，接下来进入下一段。

0x3 量化模型推理

虽然我上面写起来很简单，这当中忽略了很多坑，多得数不过来，更别说记录了。如果你也遇到了坑，欢迎前往YOLOv7留言。

接着就是量化模型推理了，：

def load_test_image(f, h, w):
    a = cv2.imread(f)
    a = cv2.resize(a, (w, h))
    a_t = np.expand_dims(np.array(a).astype(np.float32), axis=0)
    return a_t, a


def preprocess_image(img, h, w):
    a = cv2.resize(img, (w, h))
    a_t = np.expand_dims(np.array(a).astype(np.float32), axis=0)
    return a_t, img


if __name__ == "__main__":
    args = make_parser().parse_args()
    input_shape = tuple(map(int, args.input_shape.split(",")))
    session = onnxruntime.InferenceSession(args.model)

    iter = ImageSourceIter(args.image_path)
    while True:
        im = next(iter)
        if isinstance(im, str):
            im = cv2.imread(im)

        inp, ori_img = preprocess_image(im, h=input_shape[0], w=input_shape[1])

        ort_inputs = {session.get_inputs()[0].name: inp}
        output = session.run(None, ort_inputs)

        if "sparse" in args.model:
            masks, scores, labels = None, None, None
            for o in output:
                if o.dtype == np.float32:
                    scores = o
                if o.dtype == np.int32 or o.dtype == np.int64:
                    labels = o
                if o.dtype == bool:
                    masks = o
            masks = masks[0]
            print(masks.shape)
            if len(masks.shape) > 3:
                masks = np.squeeze(masks, axis=1)
            scores = scores[0]
            labels = labels[0]
            # keep = scores > 0.15
            keep = scores > (0.15 if args.int8 else 0.32)
            scores = scores[keep]
            labels = labels[keep]
            masks = masks[keep]
            print(scores)
            print(labels)
            print(masks.shape)
            img = vis_res_fast(im, None, masks, scores, labels)
        else:
            predictions = demo_postprocess(output[0], input_shape, p6=args.with_p6)[0]
            boxes = predictions[:, :4]
            scores = predictions[:, 4:5] * predictions[:, 5:]

            boxes_xyxy = np.ones_like(boxes)
            boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
            boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
            boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
            boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
            # boxes_xyxy /= ratio
            dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.65, score_thr=0.1)
            final_boxes, final_scores, final_cls_inds = (
                dets[:, :4],
                dets[:, 4],
                dets[:, 5],
            )
            img = visualize_det_cv2_part(
                ori_img, final_scores, final_cls_inds, final_boxes
            )
            cv2.imshow("aa", img)
            cv2.waitKey(0)

        cv2.imshow("YOLOv7 SparseInst CPU int8", img)
        if iter.video_mode:
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
        else:
            cv2.waitKey(0)

这部分代码在YOLOv7 deploy当中。

最后我们传入命令：

python ort_infer.py -m ../weights/sparse_inst.onnx -i ../datasets/public/images --int8

就 可以看到int8的推理结果：

0x4 总结

本文实现了一个较为复杂的transformer的实例分割的ONNX导出，同时实现了精度一定范围内得到了保证的int8量化。但其实这还只是一个粗狂的尝试，未来我们会进一步的精细化int8的量化误差、让量化误差进一步可控。

下一篇预告：使用ncnn前端推理量化模型。让int8在CPU下跑得更快。

同时，也会吧量化的结果，迁移到TensorRT,用int8来跑tensorrt.

如果你对量化感兴趣，可以扫码加入我们的 “高端” 量化交流群，群内大佬云集。

本文所有代码，将会在下面链接公布：

https://manaai.cn/