一、PTH转换为ONNX
之前介绍过ONNX，ONNX是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。可以使得在Pytorch、Tensorflow或其他框架下训练好的模型采用相同的格式来存储模型数据并进行交互。
整体转换代码：
import argparse
import time
import sys
import os
import torch
import torch.nn as nn
import onnx
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
from io import BytesIO
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./yolov6s.pt', help='weights path')# 获取Pytorch下的模型权重
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # 获取图像的size：宽高
parser.add_argument('--batch-size', type=int, default=1, help='batch size')# 每次推理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')# 表示要decode的地方，是否直接在tensor上进行解码
parser.add_argument('--simplify', action='store_true', help='simplify onnx model')# 是否需要简化onnx模型
parser.add_argument('--end2end', action='store_true', help='export end2end onnx')# 指的是是否需要把nms后处理操作也加在里面。 仅当onnxruntime 和 TensorRT >= 8.0.0才支持
parser.add_argument('--max-wh', type=int, default=None, help='None for trt int for ort')# 主要是为ort的后端设置int数据类型
parser.add_argument('--topk-all', type=int, default=100, help='topk objects for every images')# 获取前top-k个的目标
parser.add_argument('--iou-thres', type=float, default=0.45, help='iou threshold for NMS')# iou的阈值大小
parser.add_argument('--conf-thres', type=float, default=0.25, help='conf threshold for NMS')# nms的阈值大小
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')# 所使用的GPU设备
args = parser.parse_args()
args.img_size *= 2 if len(args.img_size) == 1 else 1 # expand
print(args)
t = time.time()
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'
# 导入PT模型
model = load_checkpoint(args.weights, map_location=device, inplace=True, fuse=True) # 加载fp32格式的模型
for layer in model.modules():
if isinstance(layer, RepVGGBlock):
layer.switch_to_deploy()# 如果遇到了RepVGG的block，都转为部署模式，用于加速
# 定义一个为零的输入，观察模型运行时是否会出错
img = torch.zeros(args.batch_size, 3, *args.img_size).to(device)
# 判断是否模型使用fp16的模式
if args.half:
img, model = img.half(), model.half()
model.eval()
for k, m in model.named_modules():
if isinstance(m, Conv):
if isinstance(m.act, nn.SiLU):
m.act = SiLU()
elif isinstance(m, Detect):
m.inplace = args.inplace
if args.end2end:# 判断是否导出为端到端的模型
from yolov6.models.end2end import End2End
model = End2End(model, max_obj=args.topk_all, iou_thres=args.iou_thres,
score_thres=args.conf_thres, max_wh=args.max_wh, device=device)
y = model(img) # 运行一张没有内容的空图片，让模型warmup
# ONNX模型导出
try:
LOGGER.info('\nStarting to export ONNX...')
export_file = args.weights.replace('.pt', '.onnx') # filename
with BytesIO() as f:
torch.onnx.export(model, img, f, verbose=False, opset_version=12,
training=torch.onnx.TrainingMode.EVAL,
do_constant_folding=True,
input_names=['image_arrays'],
output_names=['num_dets', 'det_boxes', 'det_scores', 'det_classes']
if args.end2end and args.max_wh is None else ['outputs'],)
f.seek(0)
# 导入onnx，并检查onnx是否有错误
onnx_model = onnx.load(f)
onnx.checker.check_model(onnx_model)
# 在这里会固定输出的大小
if args.end2end and args.max_wh is None:
shapes = [args.batch_size, 1, args.batch_size, args.topk_all, 4,
args.batch_size, args.topk_all, args.batch_size, args.topk_all]
for i in onnx_model.graph.output:
for j in i.type.tensor_type.shape.dim:
j.dim_param = str(shapes.pop(0))
# 判断是否需要使用onnxsim.simplify简化onnx模型，即去除无用的op
if args.simplify:
try:
import onnxsim
LOGGER.info('\nStarting to simplify ONNX...')
onnx_model, check = onnxsim.simplify(onnx_model)
assert check, 'assert check failed'
except Exception as e:
LOGGER.info(f'Simplifier failure: {
e}')
onnx.save(onnx_model, export_file)
LOGGER.info(f'ONNX export success, saved as {
export_file}')
except Exception as e:
LOGGER.info(f'ONNX export failure: {
e}')
# 至此，模型为onnx已经结束，下面是打印的日志信息
LOGGER.info('\nExport complete (%.2fs)' % (time.time() - t))
if args.end2end:
if args.max_wh is None:
LOGGER.info('\nYou can export tensorrt engine use trtexec tools.\nCommand is:')
LOGGER.info(f'trtexec --onnx={
export_file} --saveEngine={
export_file.replace(".onnx",".engine")}')
二、RepVGG结构
YOLOv6引入了RepVGG结构，RepVGG在训练时具有多分支，而在实际部署时可以等效地融合为单个3x3卷积的一种可重参数化的结构。通过融合成的3x3卷积，可以有效利用计算密集型的硬件的计算能力，比如GPU设备，同时也可以获得 GPU/CPU上已经优化过的NVIDIA cuDNN和Intel MKL编译框架的正向作用。
通过实验表明，使用RepVGG，YOLOv6减少了在硬件上的延时，并显著提升了算法的精度，让检测网络更快更强。以nano尺寸模型为例，对比YOLOv5-nano采用的网络结构，使用了这个方法后，YOLOv6在速度上提升了21%，同时精度也提升了3.6%AP。
三、转换RepVGG结构时的方法
以下就是遇到RepVGG时，转换方法switch_to_deploy的详细实现。主要逻辑是：如果是推理时，则直接构建一个卷积层，执行卷积运算即可。get_equivalent_kernel_bias()方法是融合RepVGGBlock中的参数，并将融合后的参数赋值给新创建的卷积对象。
def switch_to_deploy(self):
if hasattr(self, 'rbr_reparam'):
return
kernel, bias = self.get_equivalent_kernel_bias()
self.rbr_reparam = nn.Conv2d(in_channels=self.rbr_dense.conv.in_channels, out_channels=self.rbr_dense.conv.out_channels,
kernel_size=self.rbr_dense.conv.kernel_size, stride=self.rbr_dense.conv.stride,
padding=self.rbr_dense.conv.padding, dilation=self.rbr_dense.conv.dilation, groups=self.rbr_dense.conv.groups, bias=True)
self.rbr_reparam.weight.data = kernel
self.rbr_reparam.bias.data = bias
for para in self.parameters():
para.detach_()
self.__delattr__('rbr_dense')
self.__delattr__('rbr_1x1')
if hasattr(self, 'rbr_identity'):
self.__delattr__('rbr_identity')
if hasattr(self, 'id_tensor'):
self.__delattr__('id_tensor')
self.deploy = True
四、转换ONNX的技巧
-
对于任何用到shape、size返回值的参数时，例如: tensor.view(tensor.size(0),-1)这类操作，最好不要直接使用tensor.size的返回值，建议加上int转换，tensor.view(int(tensor.size(0)),-1)。
-
对于nn.Upsample或nn.functional.interpolate这些算子，使用scale_factor指定倍率，而不是使用size参数来指定最终的大小。
-
使用reshape、view这些操作时，-1的指定请放到batch维度。其他维度能通过计算得出就可以了，需要注意的是，batch维度禁止指定为大于-1的明确数字。
-
torch.onnx.export指定dynamic_axes参数，并且只指定batch维度，而不指定其他维度。因为这里只需要动态的batch。
五、ONNX转换为TensorRT
当成功地将PTH模型转换为ONNX的格式后，就可以使用ONNX来转换TensorRT的Engine了，ONNX转换为TensorRT的流程如下所示：
1、转换流程
一般来说，构建为引擎的过程都比较固定：
- 首先会构建Builder，以及通过实例化一个Logger接口来捕获异常，如果构建过程有什么问题，就可以通过查找日志信息来定位。
- 然后就是使用模型结构实例化Builder，这一步也可以叫定义网络结构。
- 定义好网络以后，就会创建ONNX解析器来对填充网络。
- 接着在构建Engine时会创建优化的配置。在这里，需要特别注意的是workspace-size，它表示的是构建TensorRT引擎所需的GPU工作空间大小，需要你自己指定大小。如果没指定，那么就默认为1GB。在TensorRT中，层的实现通常需要一个临时工作空间，并且这个参数限制了网络中任何层可以使用的最大的大小。如果提供的工作空间不足，TensorRT 就可能无法找到层的实现。
- 最后就可以按照配置来构建引擎，也就是序列化模型，通过这样，可以把Engine转换为可存储的格式方便后面推理使用。那么当推理时，直接对Engine进行反序列化操作就可以了。因为从头构建一个Engine是很耗费时间的，序列化保存后可以避免每次推理时都重新创建Engine，从而节省时间。
- 转换结束后可以在上下文中进行推理，来验证在这个过程中是否完成了正确的转换。
整体转换代码：
import tensorrt as trt
import sys
import argparse
"""
takes in onnx model
converts to tensorrt
"""
def onnx_export_trt():
desc = 'compile Onnx model to TensorRT'
parser = argparse.ArgumentParser(description=desc)
parser.add_argument('-m', '--model', help='onnx file location inside ./lib/models')# 导入onnx模型
parser.add_argument('-fp', '--floatingpoint', type=int, default=32, help='floating point precision. 16 or 32')# 判断使用fp32还是fp16模式
parser.add_argument('-o', '--output', help='name of trt output file')# 指定所要导出的tensorRT的路径
args = parser.parse_args()
model = args.model
fp = args.floatingpoint
if fp != 16 and fp != 32:
print('floating point precision must be 16 or 32')
sys.exit()
output = args.output
return {
'model': model,
'fp': fp,
'output': output
}
if __name__ == '__main__':
args = onnx_export_trt()
model = '{}'.format(args['model'])
output = '{}'.format(args['output'])
logger = trt.Logger(trt.Logger.VERBOSE)
EXPLICIT_BATCH = []
print('trt version', trt.__version__)
if trt.__version__[0] >= '7':
EXPLICIT_BATCH.append(
1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))# 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间，和执行时最大可以使用的batchsize大小相关
with trt.Builder(logger) as builder, builder.create_network(*EXPLICIT_BATCH) as network, trt.OnnxParser(network, logger) as parser:
# 使用onnx的解析器来绑定计算图，之后将通过所解析的来填充后面的计算图
config = builder.create_builder_config()
config.max_workspace_size = 1 << 28
builder.max_batch_size = 1
if args['fp'] == '16':
builder.fp16_mode = True
with open(model, 'rb') as f:
# 解析onnx文件，用于填充后面的计算图
if not parser.parse(f.read()):
for error in range(parser.num_errors):
print(parser.get_error(error))
shape = list(network.get_input(0).shape)
engine = builder.build_serialized_network(network,config)# 根据模型来构建引擎，用于后面的推理
# 将引擎进行持久化的序列化存储
with open(output, 'wb') as f:
f.write(engine)
2、trtexec导出engine引擎
当获得ONNX模型后，也可以直接使用trtexec来序列化引擎，并保存下来，只是这个构建过程的耗时会比较长。
trtexec --onnx=yolov6s.onnx --saveEngine=yolov6s.engine --fp16
可以查看trtexec可以设置哪一些输入参数，支持哪一些可以转换的模式。
六、测试结果
以下是将PTH模型转换为ONNX后再转换为TensorRT引擎的推理结果：
由下面的表格可以看到，转换后的模型在A100上表现出的FPS还是很优秀的，并且FP16比FP32还会快一些，从表中可以得出，使用TensorRT部署YOLOv6可以获得一个比较可观的实时性。
| model | input | FP16/FP32 | FPS | Device |
|---|---|---|---|---|
| yolov6s | 640*640 | FP16 | 360FPS | A100 |
| yolov6s | 640*640 | FP32 | 350FPS | A100 |
参考链接：
https://tech.meituan.com/2022/09/22/yolov6-quantization-in-meituan.html
https://arxiv.org/pdf/2209.02976.pdf
https://github.com/meituan/YOLOv6
https://mp.weixin.qq.com/s?__biz=MzkyMDE2OTA3Mw==&mid=2247502379&idx=2&sn=4d09f25a675758f6cd8090a26e764a28&chksm=c1947535f6e3fc23320e3aaf58fd43d5dec395476c67a9ea7fa5528d925d373b0dcd347f1753&cur_album_id=2405123118882357250&scene=189#wechat_redirect