美团 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 ortROOT = 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_checkpointif __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, widthparser.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  # expandprint(args)t = time.time()apply_simplify = True  # 增加onnxsim# Check devicecuda = 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 modelmodel = load_checkpoint(args.weights, map_location=device, inplace=True, fuse=True)  # load FP32 modelfor layer in model.modules():if isinstance(layer, RepVGGBlock):layer.switch_to_deploy()# Inputimg = torch.zeros(args.batch_size, 3, *args.img_size).to(device)  # image size(1,3,320,192) iDetection# Update modelif args.half:img, model = img.half(), model.half()  # to FP16model.eval()for k, m in model.named_modules():if isinstance(m, Conv):  # assign export-friendly activationsif isinstance(m.act, nn.SiLU):m.act = SiLU()elif isinstance(m, Detect):m.inplace = args.inplacey = model(img)  # dry run# ONNX exporth, w = args.img_sizeexport_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"],)# Checksonnx_model = onnx.load(export_file)  # load onnx modelonnx.checker.check_model(onnx_model)  # check onnx modelLOGGER.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 modelonnx_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 !")# FinishLOGGER.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.shapey = 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.deviceyv, 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]  # xyy[..., 2:4] = torch.exp(y[..., 2:4]) * self.stride[i] # whelse:xy = (y[..., 0:2] + self.grid[i]) * self.stride[i]  # xywh = torch.exp(y[..., 2:4]) * self.stride[i]  # why = 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 NCNNx[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 ONNXRuntimeauto *yolov6 = new lite::cv::detection::YOLOv6(onnx_path); // defaultstd::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_MNNstd::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 MNNauto *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_NCNNstd::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 NCNNauto *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_TNNstd::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 TNNauto *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
}

输出的结果如下：

延庆川北小区45孙老师东屯收卖废品垃圾破烂炒股废品孙一起人工智能把

whaosoft aiot http://143ai.com 没有篇 c++ yolo用不了模型转换有用哦虽然yolo7 也出来了还么i找到哦

6. 总结

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

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