凑个热闹之美团 YOLOv6 ORT/MNN/TNN/NCNN C++推理部署

0. 前言

昨天美团开源了YOLOv6，又是一个YOLO系列的新作。此时距离YOLOX开源差不多刚好一年的时间。之前捏过很多YOLO系列的推理例子，比如YOLOv3、YOLOv4、YOLOv5、YOLOX、YOLOR和YOLOP等等。虽然最近已经没有在做detection方向了，但作为YOLO系列的老粉了，出来凑个热闹应该总是可以的。所以，这次也来凑凑YOLO系列的热闹，给出几个不同的推理引擎的例子，包括ONNXRuntime、MNN、NCNN和TNN，以及简单记录下模型的转换过程。总的来说，YOLOv6 的 C++推理，都是些重复性的工作，没什么太大的难度，刚好趁着周末，顺手捏一捏。这篇文章不会记录地很详细，只讲几个要点。

1. ONNX 和 TNN 模型转换

经过尝试，直接转换出来的ONNX和TNN模型文件在推理时，结果一切正常，不需要修改 YOLOv6 的 Detect 源码，使用官方提供的 deploy/ONNX/export_onnx.py 直接转换即可。但是 NCNN 和 MNN 都需要修改 Detect 的源码进行特殊处理才可正常推理。所以 ONNX 和 TNN 放在这一节讲，MNN 和 NCNN 的模型转换放在下一小节讲。

首先下载源码:

git clone --depth=1 https://github.com/meituan/YOLOv6.git

然后稍微修改下 export_onnx.py，源码是没有添加 onnxsim的，作为基操，我们把它添加上，修改后如下：

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import argparse
import time
import sys
import os
import torch
import torch.nn as nn
import onnx
import onnxsim
import onnxruntime as ort

ROOT = os.getcwd()
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

from yolov6.models.yolo import *
from yolov6.models.effidehead import Detect
from yolov6.layers.common import *
from yolov6.utils.events import LOGGER
from yolov6.utils.checkpoint import load_checkpoint

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='./yolov6s.pt', help='weights path')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
    parser.add_argument('--inplace', action='store_true', help='set Detect() inplace=True')
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0, 1, 2, 3 or cpu')
    args = parser.parse_args()
    args.img_size *= 2 if len(args.img_size) == 1 else 1  # expand
    print(args)
    t = time.time()
    apply_simplify = True  # 增加onnxsim

    # Check device
    cuda = args.device != 'cpu' and torch.cuda.is_available()
    device = torch.device('cuda:0' if cuda else 'cpu')
    assert not (device.type == 'cpu' and args.half), '--half only compatible with GPU export, i.e. use --device 0'
    # Load PyTorch model
    model = load_checkpoint(args.weights, map_location=device, inplace=True, fuse=True)  # load FP32 model
    for layer in model.modules():
        if isinstance(layer, RepVGGBlock):
            layer.switch_to_deploy()
    # Input
    img = torch.zeros(args.batch_size, 3, *args.img_size).to(device)  # image size(1,3,320,192) iDetection
    # Update model
    if args.half:
        img, model = img.half(), model.half()  # to FP16
    model.eval()
    for k, m in model.named_modules():
        if isinstance(m, Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.SiLU):
                m.act = SiLU()
        elif isinstance(m, Detect):
            m.inplace = args.inplace
    y = model(img)  # dry run
    # ONNX export
    h, w = args.img_size
    export_file = args.weights.replace('.pt', f'-{h}x{w}.onnx')  # filename 增加size标记
    try:
        LOGGER.info('\nStarting to export ONNX...')
        torch.onnx.export(model, img, export_file, verbose=False, opset_version=12,
                          training=torch.onnx.TrainingMode.EVAL,
                          do_constant_folding=True,
                          input_names=['image_arrays'],
                          output_names=["outputs"],
                          )
        # Checks
        onnx_model = onnx.load(export_file)  # load onnx model
        onnx.checker.check_model(onnx_model)  # check onnx model
        LOGGER.info(f'ONNX export success, saved as {export_file}')
    except Exception as e:
        LOGGER.info(f'ONNX export failure: {e}')
    if apply_simplify:  # 增加的onnxsim部分
        print(f'{export_file} simplifying with onnx-simplifier {onnxsim.__version__}...')
        try:
            onnx_model = onnx.load(export_file)  # load onnx model
            onnx_model, check = onnxsim.simplify(onnx_model, check_n=3)
            assert check, 'simplifying check failed'
            onnx.save(onnx_model, export_file)
        except Exception as e:
            print(f'{export_file} simplifier failure: {e}')
    # Running ORT check 增加ORT验证
    sess = ort.InferenceSession(export_file)
    print(f"ORT Loaded {export_file} !")
    for _ in sess.get_inputs(): print(f"Input: {_}")
    for _ in sess.get_outputs(): print(f"Output: {_}")
    print("ORT Check Done !")

    # Finish
    LOGGER.info('\nExport complete (%.2fs)' % (time.time() - t))

预训练好的pt模型文件可以从官方提供的链接下载，放在 YOLOv6 的根目录下，直接转换就行：

PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6n.pt --img 640 --batch 1
PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6n.pt --img 320 --batch 1
PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1
PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 320 --batch 1
PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6t.pt --img 640 --batch 1

这个过程比较顺利，暂时没发现什么坑。接下来转换为 TNN 模型文件，命令如下：

convert2tnn# python3 ./converter.py onnx2tnn ./tnn_models/yolov6/yolov6t-640x640.onnx -o ./tnn_models/yolov6/ -v v1.0 -optimize -align
convert2tnn# python3 ./converter.py onnx2tnn ./tnn_models/yolov6/yolov6s-640x640.onnx -o ./tnn_models/yolov6/ -v v1.0 -optimize -align
convert2tnn# python3 ./converter.py onnx2tnn ./tnn_models/yolov6/yolov6n-640x640.onnx -o ./tnn_models/yolov6/ -v v1.0 -optimize -align
convert2tnn# python3 ./converter.py onnx2tnn ./tnn_models/yolov6/yolov6s-320x320.onnx -o ./tnn_models/yolov6/ -v v1.0 -optimize -align
convert2tnn# python3 ./converter.py onnx2tnn ./tnn_models/yolov6/yolov6n-320x320.onnx -o ./tnn_models/yolov6/ -v v1.0 -optimize -align

TNN模型的转换需要用到 tnn-convert，如何使用tnn-convert就不展开了，感兴趣的同学可以看我之前写的一篇文章，传送门：

tnn-convert搭建简记-YOLOP转TNN5：https://zhuanlan.zhihu.com/p/431418709

2. MNN 和 NCNN 模型转换

NCNN 和 MNN 都需要修改 Detect 的源码进行特殊处理才可正常推理。MNN其实也可以直接转，但是转出来的模型文件虽然能推理，但是在decode完之后，结果很奇怪，所以我最后决定将MNN的模型文件转换采用和NCNN同样的处理方式。其实很多写的关于YOLO系列的部署文章都有提到过，该系列在部署时主要的一个问题就是如何处理 Detect Head 中关于 decode 部分的逻辑。这部分代码在 YOLOv6中长这样：

    def forward(self, x):
        z = []
        for i in range(self.nl):
            # ...
            if self.training:
                # ...
            else:
                y = torch.cat([reg_output, obj_output.sigmoid(), cls_output.sigmoid()], 1)
                bs, _, ny, nx = y.shape
                y = y.view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
                if self.grid[i].shape[2:4] != y.shape[2:4]:
                    d = self.stride.device
                    yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)])
                    self.grid[i] = torch.stack((xv, yv), 2).view(1, self.na, ny, nx, 2).float()
                if self.inplace:
                    y[..., 0:2] = (y[..., 0:2] + self.grid[i]) * self.stride[i]  # xy
                    y[..., 2:4] = torch.exp(y[..., 2:4]) * self.stride[i] # wh
                else:
                    xy = (y[..., 0:2] + self.grid[i]) * self.stride[i]  # xy
                    wh = torch.exp(y[..., 2:4]) * self.stride[i]  # wh
                    y = torch.cat((xy, wh, y[..., 4:]), -1)
                z.append(y.view(bs, -1, self.no))
        return x if self.training else torch.cat(z, 1)

这部分虽然有些框架可以直接支持导出，但会产生大量的胶水op，所以一个可选的做法是，只导出decode之前的Raw部分的内容，在C++侧做decode。另外，我们可以看到，decode部分存在一个5维度的操作：

 y = y.view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

这部分在NCNN中应该是不支持的（按照我对ncnn::Mat的理解，它有c,h,w三个维度，并假设b=1，所以可以处理<=4维的张量），也不能直接导出。所以，这个5维的处理，我们也要做相应的修改。至于MNN，其实可以直接转换这部分decode的逻辑，但是我在推理时，发现出来的结果不太对，于是决定采用NCNN同样的处理方式，就是只导出decode前的部分，把decode放在c++侧处理，后来验证了这样做可以得到正常的推理结果。

那么，这段 Detect Head 的逻辑到底要怎么改呢？长话短说，我直接放一个我修改后的代码吧：

models/effidehead.py修改之后

class Detect(nn.Module):
    # ...
    def forward(self, x):
        z = []
        for i in range(self.nl):
            # ...
            if self.training:
                # ...
            else:  # 修改之后长这样
                x[i] = torch.cat([reg_output, obj_output, cls_output], 1)
                bs, _, ny, nx = x[i].shape
                # x(bs,255,20,20) to x(bs,1,20,20,85=80+5) (bs,na,ny,nx,no=nc+5=4+1+nc)
                # x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
                x[i] = x[i].view(bs, self.na, 85, -1).permute(0, 1, 3, 2).contiguous()  # (b,self.na,20x20,85) for NCNN
                x[i] = x[i].view(bs, -1, 85)
        return torch.cat(x, dim=1)  # (b,?,85)

这里，我们把原来的5维操作修改成4维操作，因为导出时，实际的no(num_outputs=85)、na(num_anchors=1)的值都是可以事先计算出来的，通过看 YOLOv6 中的 configs 文件夹中的配置文件，我们也可以确定实际上 na一直为1，而 ny，nx是每个特征图的大小，对于固定的输入shape，各个特征图的ny，nx也是固定的。为了变成4维操作，我们可以把原来在最后的两个维度(ny,nx)直接拉平，按照行主序的线性内存来理解，这样是可行的。于是有：

x[i] = x[i].view(bs, self.na, 85, -1).permute(0, 1, 3, 2).contiguous()  # (b,self.na,20x20,85) for NCNN
x[i] = x[i].view(bs, -1, 85)

最后，在返回结果值之前，我们再做一个合并处理，这样就不需要在c++解码的时候单独对每个特征图都做一遍，合并后，只要对一个大的输出做decode就可以了。

return torch.cat(x, dim=1)  # (b,?,85)

这样修改之后，export_onnx.py的代码不用变，还是用原来的命令行直接导出就可以了（我在文件名增加了-for-ncnn作为后缀方便区分）。

还是先导出为 ONNX：

PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6n.pt --img 640 --batch 1
PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6n.pt --img 320 --batch 1
PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1
PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 320 --batch 1
PYTHONPATH=. python3 ./deploy/ONNX/export_onnx.py --weights yolov6t.pt --img 640 --batch 1

模型文件为：

yolov6 ls -lh | grep yolov6 | grep for-ncnn | grep onnx
-rw-r--r--  1 yanjunqiu  staff    16M Jun 25 12:52 yolov6n-320x320-for-ncnn.onnx
-rw-r--r--  1 yanjunqiu  staff    16M Jun 25 12:51 yolov6n-640x640-for-ncnn.onnx
-rw-r--r--  1 yanjunqiu  staff    66M Jun 25 12:53 yolov6s-320x320-for-ncnn.onnx
-rw-r--r--  1 yanjunqiu  staff    66M Jun 25 12:52 yolov6s-640x640-for-ncnn.onnx
-rw-r--r--  1 yanjunqiu  staff    57M Jun 25 12:53 yolov6t-640x640-for-ncnn.onnx

用netron打开来看，发现decode那部分已经没有了：

带decode的onnx，用netron打开是长这样的：

可以发现两者有很明显的区别，不带decode的图，在输出部分要简单很多。

接下来，就是按常规流程将ONNX转换成NCNN和MNN模型，命令行如下:

onnx2ncnn yolov6/yolov6s-320x320-for-ncnn.onnx yolov6/yolov6s-320x320-for-ncnn.param yolov6/yolov6s-320x320-for-ncnn.bin
ncnnoptimize yolov6/yolov6s-320x320-for-ncnn.param yolov6/yolov6s-320x320-for-ncnn.bin yolov6/yolov6s-320x320-for-ncnn.opt.param yolov6/yolov6s-320x320-for-ncnn.opt.bin 0

一切正常。转换为MNN模型的命令为：

YOLOv6 MNNConvert -f ONNX --modelFile yolov6n-640x640-for-ncnn.onnx --MNNModel yolov6n-640x640.mnn --bizCode MNN 
Start to Convert Other Model Format To MNN Model...
[15:19:09] /Users/yanjunqiu/Desktop/third_party/library/MNN/tools/converter/source/onnx/onnxConverter.cpp:30: ONNX Model ir version: 6
Start to Optimize the MNN Net...
inputTensors : [ image_arrays, ]
outputTensors: [ outputs, ]
Converted Success!

也是一切正常。

3. ONNX 和 TNN 模型 C++ 推理

ONNX和TNN的模型都是带decode的，因此在后处理时简单些，不用生成anchor了。模型推理直接输出的维度是(1,n,85)，这个n表示总共输出的anchors个数，85的含义是：

85=5+80=cxcy(2)+cwch(2)+obj_conf(1)+cls_conf(80)

由于输出的坐标就已经是归一化后的cx,cy和cw,ch，所以后处理就很简单了，直接转换成x1,y1,x2,y2格式就行。逻辑大概如下：

    float cx = offset_obj_cls_ptr[0];
    float cy = offset_obj_cls_ptr[1];
    float w = offset_obj_cls_ptr[2];
    float h = offset_obj_cls_ptr[3];
    float x1 = ((cx - w / 2.f) - (float) dw_) / r_;
    float y1 = ((cy - h / 2.f) - (float) dh_) / r_;
    float x2 = ((cx + w / 2.f) - (float) dw_) / r_;
    float y2 = ((cy + h / 2.f) - (float) dh_) / r_;

详细的推理代码就不展开了，会在文章最后放出。

4. NCNN 和 MNN 模型 C++ 推理

NCNN和MNN的模型文件没有导出decode部分，因此后处理复杂一点。后处理主要包括2部分，一是生成anchors，二是根据生成的anchors和输出的原始信息解码坐标。YOLOv6 的anchors生成逻辑其实和YOLOX基本是一致的，就是每个feature map上每个锚点生成一个anchor框，做过detection算法的同学应该很多理解这句话的意思，我也就不啰嗦了。直接放代码吧。

generate_anchors函数主要逻辑

void NCNNYOLOv6::generate_anchors(const int target_height,
                                  const int target_width,
                                  std::vector<int> &strides,
                                  std::vector<YOLOv6Anchor> &anchors)
{
  for (auto stride: strides)
  {
    int num_grid_w = target_width / stride;
    int num_grid_h = target_height / stride;
    for (int g1 = 0; g1 < num_grid_h; ++g1)
    {
      for (int g0 = 0; g0 < num_grid_w; ++g0)
      {
        YOLOv6Anchor anchor;
        anchor.grid0 = g0;
        anchor.grid1 = g1;
        anchor.stride = stride;
        anchors.push_back(anchor);
      }
    }
  }
}

坐标解码的主要逻辑

    const int grid0 = anchors.at(i).grid0;
    const int grid1 = anchors.at(i).grid1;
    const int stride = anchors.at(i).stride;

    float dx = offset_obj_cls_ptr[0];
    float dy = offset_obj_cls_ptr[1];
    float dw = offset_obj_cls_ptr[2];
    float dh = offset_obj_cls_ptr[3];

    float cx = (dx + (float) grid0) * (float) stride;
    float cy = (dy + (float) grid1) * (float) stride;
    float w = std::exp(dw) * (float) stride;
    float h = std::exp(dh) * (float) stride;
    float x1 = ((cx - w / 2.f) - (float) dw_) / r_;
    float y1 = ((cy - h / 2.f) - (float) dh_) / r_;
    float x2 = ((cx + w / 2.f) - (float) dw_) / r_;
    float y2 = ((cy + h / 2.f) - (float) dh_) / r_;

详细的推理代码就不展开了，会在文章最后放出。

5. YOLOv6 C++推理的使用例子

首先是放出YOLOv6 的4个推理引擎的C++ 源码，想必大家最关心就是能不能白嫖了。所有的代码都集成进了lite.ai.toolkit（https://github.com/DefTruth/lite.ai.toolkit） 工具箱中，零成本无压力白嫖。

• YOLOv6 ONNXRuntime C++ 源码：https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/ort/cv/yolov6.cpp
• YOLOv6 MNN C++ 源码：https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/mnn/cv/mnn_yolov6.cpp
• YOLOv6 NCNN C++ 源码：https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/ncnn/cv/ncnn_yolov6.cpp
• YOLOv6 TNN C++ 源码：https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/tnn/cv/tnn_yolov6.cpp

对源码感兴趣的同学，可自行选择关心的推理引擎版本进行阅读。接下来，在简单贴几个使用lite.ai.toolkit 工具箱一键调用的例子。

ONNXRuntime版本

#include "lite/lite.h"

static void test_default()
{
  std::string onnx_path = "../../../hub/onnx/cv/yolov6s-640x640.onnx";
  std::string test_img_path = "../../../examples/lite/resources/test_lite_yolov6_1.jpg";
  std::string save_img_path = "../../../logs/test_lite_yolov6_1.jpg";

  // 1. Test Default Engine ONNXRuntime
  auto *yolov6 = new lite::cv::detection::YOLOv6(onnx_path); // default

  std::vector<lite::types::Boxf> detected_boxes;
  cv::Mat img_bgr = cv::imread(test_img_path);
  yolov6->detect(img_bgr, detected_boxes);
  lite::utils::draw_boxes_inplace(img_bgr, detected_boxes);

  cv::imwrite(save_img_path, img_bgr);
  std::cout << "Default Version Detected Boxes Num: " << detected_boxes.size() << std::endl;
  delete yolov6;
}

MNN版本

static void test_mnn()
{
#ifdef ENABLE_MNN
  std::string mnn_path = "../../../hub/mnn/cv/yolov6s-640x640.mnn";
  std::string test_img_path = "../../../examples/lite/resources/test_lite_yolov6_2.jpg";
  std::string save_img_path = "../../../logs/test_lite_yolov6_mnn_2.jpg";

  // 3. Test Specific Engine MNN
  auto *yolov6 = new lite::mnn::cv::detection::YOLOv6(mnn_path);

  std::vector<lite::types::Boxf> detected_boxes;
  cv::Mat img_bgr = cv::imread(test_img_path);
  yolov6->detect(img_bgr, detected_boxes);

  lite::utils::draw_boxes_inplace(img_bgr, detected_boxes);
  cv::imwrite(save_img_path, img_bgr);
  std::cout << "MNN Version Detected Boxes Num: " << detected_boxes.size() << std::endl;
  delete yolov6;
#endif
}

NCNN版本

static void test_ncnn()
{
#ifdef ENABLE_NCNN
  std::string param_path = "../../../hub/ncnn/cv/yolov6s-640x640-for-ncnn.opt.param";
  std::string bin_path = "../../../hub/ncnn/cv/yolov6s-640x640-for-ncnn.opt.bin";
  std::string test_img_path = "../../../examples/lite/resources/test_lite_yolov6_2.jpg";
  std::string save_img_path = "../../../logs/test_lite_yolov6_ncnn_2.jpg";

  // 4. Test Specific Engine NCNN
  auto *yolov6 = new lite::ncnn::cv::detection::YOLOv6(param_path, bin_path);

  std::vector<lite::types::Boxf> detected_boxes;
  cv::Mat img_bgr = cv::imread(test_img_path);
  yolov6->detect(img_bgr, detected_boxes);

  lite::utils::draw_boxes_inplace(img_bgr, detected_boxes);
  cv::imwrite(save_img_path, img_bgr);
  std::cout << "NCNN Version Detected Boxes Num: " << detected_boxes.size() << std::endl;
  delete yolov6;
#endif
}

TNN版本

static void test_tnn()
{
#ifdef ENABLE_TNN
  std::string proto_path = "../../../hub/tnn/cv/yolov6s-640x640.opt.tnnproto";
  std::string model_path = "../../../hub/tnn/cv/yolov6s-640x640.opt.tnnmodel";
  std::string test_img_path = "../../../examples/lite/resources/test_lite_yolov6_2.jpg";
  std::string save_img_path = "../../../logs/test_lite_yolov6_tnn_2.jpg";

  // 5. Test Specific Engine TNN
  auto *yolov6 = new lite::tnn::cv::detection::YOLOv6(proto_path, model_path);

  std::vector<lite::types::Boxf> detected_boxes;
  cv::Mat img_bgr = cv::imread(test_img_path);
  yolov6->detect(img_bgr, detected_boxes);

  lite::utils::draw_boxes_inplace(img_bgr, detected_boxes);
  cv::imwrite(save_img_path, img_bgr);
  std::cout << "TNN Version Detected Boxes Num: " << detected_boxes.size() << std::endl;
  delete yolov6;
#endif
}

输出的结果如下：

6. 总结

本文主要讲解了美团6月24开源的YOLOv6在不同推理引擎的C++工程化过程，包含了ORT、MNN、NCNN和TNN下的推理处理，以及不同推理框架下模型转换需要注意的问题。并且，所有的代码都集成进了lite.ai.toolkit 工具箱中，零成本无压力白嫖。

• YOLOv6 ONNXRuntime C++ 源码
• YOLOv6 MNN C++ 源码
• YOLOv6 NCNN C++ 源码
• YOLOv6 TNN C++ 源码

趁着周末捏美团的这个YOLOv6模型，写完这篇文章的时候，差不多是傍晚6点，时值仲夏，我摸了摸饥肠辘辘的肚子，觉得是时候打开我手机上的美团外卖来下个单了~ 问题来了，那到底吃啥好呢 ...

公众号后台回复“项目实践”获取50+CV项目实践机会～