环境说明：
tensorRT:8.2.4.2
CUDA:10.2
pytorch:1.7
显卡：NVIDIA 1650
Windows10
python 3.7
另一篇文章中写过C++版的trt推理。本篇文章是python版本的trt yolov5推理。
构建engine一般有两种方式。
方式1：torch模型->wts(序列化网络)->engine->推理
方式2：torch模型->onnx->engine->推理
第一种方式如果网络结构简单，在定义网络构建engine的时候还可以，但网络复杂的情况就麻烦了，写网络的时候还容易出错。
第二种方式也是很多人常用的方法，转onnx再转engine。转onnx就比较容易了，而转engine一般有两种方式，第一种是trt官方自带的方式，在你trt文件下的bin目录下有个trtexec.exe的文件，执行命令就可以将onnx转engine。而第二种python版trt自带工具，这也是本文要介绍的。
我这里的代码是用的v5 6.1代码，因为6.1以及之后版本的export.py中有engine格式的导出。
我们可以看一下官方提供的yolov5s.pt中都包含什么内容：
dict_keys(['epoch', 'best_fitness', 'model', 'ema', 'updates', 'optimizer', 'wandb_id', 'date'])
可以看到上述pt文件中包含了这些key值，其中的model就是我们要的，而且需要注意的是这个model不仅含有网络权重信息，还包含了整个网络结构【如果你想把其他网络转onnx，也需要主要必须torch保存的是整个网络】
导出onnx
执行下面的命令就可以得到我们的onnx模型。
python export.py --weights yolov5s.pt --include onnx
这里附上导出onnx的代码。
@try_export def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')): # YOLOv5 ONNX export check_requirements('onnx') import onnx LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') f = file.with_suffix('.onnx') output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0'] if dynamic: dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640) if isinstance(model, SegmentationModel): dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160) elif isinstance(model, DetectionModel): dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) torch.onnx.export( model.cpu() if dynamic else model, # --dynamic only compatible with cpu im.cpu() if dynamic else im, f, verbose=False, opset_version=opset, do_constant_folding=True, input_names=['images'], output_names=output_names, dynamic_axes=dynamic or None) # Checks model_onnx = onnx.load(f) # load onnx model onnx.checker.check_model(model_onnx) # check onnx model # Metadata d = {'stride': int(max(model.stride)), 'names': model.names} for k, v in d.items(): meta = model_onnx.metadata_props.add() meta.key, meta.value = k, str(v) onnx.save(model_onnx, f) # Simplify if simplify: try: cuda = torch.cuda.is_available() check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1')) import onnxsim LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') model_onnx, check = onnxsim.simplify(model_onnx) assert check, 'assert check failed' onnx.save(model_onnx, f) except Exception as e: LOGGER.info(f'{prefix} simplifier failure: {e}') return f, model_onnx
export_onnx函数中，model就是我们加载的torch网络，im是一个输入样例 ，file为yolov5s.pt的路径[我这里为F:/yolov5/yolov5s.pt]。opset就是版本这里是12，dynamic就说动态输入【我这里没开启】。
output_names是获取model的结点名字，由于这里是目标检测不是图像分割，因此仅有一个output，取名为output0。
这里需要注意一点的是，明明v5有三个head，为什么这里仅一个输出，如果你去看models/yolo.py中的Detect可以看到下面的代码，在export模型下会把三个输出拼接为1个。
# 如果export 为True，返回的输出是三个head合并为1个。 return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x) output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0']
下面这部分代码是onnx导出的核心代码，这里需要注意一下这里需要传入输入(input_names)输出结点(output_names)。【这里的结点名不要随意更改，因为后面还会用到】
torch.onnx.export( model.cpu() if dynamic else model, # --dynamic only compatible with cpu im.cpu() if dynamic else im, f, verbose=False, opset_version=opset, do_constant_folding=True, input_names=['images'], output_names=output_names, dynamic_axes=dynamic or None)
下面的图就onnx可视化，images就是我们前面定义好的输入结点。
下面这一部分就是输出部分，输出是三个头进行了整合，结点为output0，shape[1,25200,85]。这里的25200=80 * 80 *3 + 40 * 40 *3 + 20* 20 *3【3是anchors】，85就是80个类+(center_x,cente_y,w,h,conf)
在介绍导出engine过程需要先介绍一下会遇到的相关术语。
1.建立logger:日志记录器
2.建立Builder:网络元数据
用于搭建网络的入口，网络的TRT内部表示以及可执行程序引擎都是由该对象的成员方法生成
3.建立BuilderConfig:网络元数据的选项
负责配置模型的一些参数，比如是否开启FP16，int8模型等。
通常的语句为：config = builder.create_builder_config()
常用的成员函数：
config.max_workspace_size = 1<<30 # 指定构建期间可用显存(单位：Byte)
config.flag = .. # 设置标志位，如1<<int(trt.BuilderFlag.FP16)
4.创建Network:计算图内容
网络主体，使用api搭建网络过程中，将不断的向其中添加一些层，并标记网络的输入输出节点(这个可能大家在使用C++构建engine的时候遇到过，也就是wts->engine的过程)。不过这里还提供了其他的方法，可以采用解析器Parser加载来自onnx文件中的网络(推荐使用)，就不用一层一层手工添加。
语法：network = builder.create_network()
常用方法：
network.add_input('tensor',trt.float32,(3,4,5)) # 标记网络输入张量
convLayer = network.add_convolution_nd(XXX) # 添加各种网络层
network.mark_output(convLayer.get_output(0)) # 标记网络输出张量
常用获取网络信息的成员：
network.name/network.num_layers/network.num_inputs/network.num_outputs
network是计算图在TRT中的具体描述，由builder生成，在使用TRT原生api搭建网络的workflow中，我们需要不断地想network中添加一些层，并标记network的输入输出张量，而在使用parser导入onnx模型的workflow中，一旦模型解析完成，network的内容就会被自动填入。
5.生成SerializedNetwork:网络的TRT内部表示
模型网络在TRT中的内部表示，可用它生成可执行的推理引擎或者把它序列化保存为文件，方便以后读取和使用
导出engine
导出engine代码如下所示。
# engine TRT 必须在GPU上 @try_export def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')): # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt # 首先判断一下im是不是在GPU上 assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`' try: import tensorrt as trt except Exception: if platform.system() == 'Linux': # 判断操作系统 check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com') import tensorrt as trt # 判断trt版本 if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 grid = model.model[-1].anchor_grid model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid] export_onnx(model, im, file, 12, dynamic, simplify) # opset 12 model.model[-1].anchor_grid = grid else: # TensorRT >= 8 check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0 # 先转onnx export_onnx(model, im, file, 12, dynamic, simplify) # opset 12 onnx = file.with_suffix('.onnx') # 获取权重名 LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...') assert onnx.exists(), f'failed to export ONNX file: {onnx}' f = file.with_suffix('.engine') # TensorRT engine file # 记录trt转engine日志 logger = trt.Logger(trt.Logger.INFO) if verbose: logger.min_severity = trt.Logger.Severity.VERBOSE # 1.builder构造，记录日志 builder = trt.Builder(logger) # 2.builder.config建立 config = builder.create_builder_config() # 3.workspace workspace * 1 << 30 表示将workspace * 1 二进制左移30位后的10进制 config.max_workspace_size = workspace * 1 << 30 # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice # 4.定义Network并加载onnx解析器 flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) network = builder.create_network(flag) parser = trt.OnnxParser(network, logger) if not parser.parse_from_file(str(onnx)): raise RuntimeError(f'failed to load ONNX file: {onnx}') # 5.获得网络输入输出 inputs = [network.get_input(i) for i in range(network.num_inputs)] outputs = [network.get_output(i) for i in range(network.num_outputs)] # 下面的只是在log中打印input和output 的name和shape以及数据类型 for inp in inputs: LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}') for out in outputs: LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}') # 判断动态输入 if dynamic: if im.shape[0] <= 1: LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument") profile = builder.create_optimization_profile() for inp in inputs: profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape) config.add_optimization_profile(profile) LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}') # 判断是否支持FP16推理 if builder.platform_has_fast_fp16 and half: config.set_flag(trt.BuilderFlag.FP16) # build engine 文件的写入 这里的f是前面定义的engine文件 with builder.build_engine(network, config) as engine, open(f, 'wb') as t: # 序列化model t.write(engine.serialize()) return f, None
构建engine关键步骤如下：
1.builder构造。
其中的logger是记录trt转engine时的log信息。
这个步骤是构建引擎的核心部分。
# 记录trt转engine日志 logger = trt.Logger(trt.Logger.INFO) builder = trt.Builder(logger)
builder中的属性内容下 并输出下面的log内容，主要是一些内存上面的使用初始等。
[10/13/2022-17:40:17] [TRT] [I] [MemUsageChange] Init CUDA: CPU +285, GPU +0, now: CPU 7095, GPU 1776 (MiB)
[10/13/2022-17:40:17] [TRT] [I] [MemUsageSnapshot] Begin constructing builder kernel library: CPU 7129 MiB, GPU 1776 MiB
[10/13/2022-17:40:18] [TRT] [I] [MemUsageSnapshot] End constructing builder kernel library: CPU 7227 MiB, GPU 1810 MiB
2.builder.config构造
config = builder.create_builder_config()
3.workspace分配
# 3.workspace workspace * 1 << 30 表示将workspace * 1 二进制左移30位后的10进制 config.max_workspace_size = workspace * 1 << 30
4.网络定义并加载onnx解析器(写入网络)
网络的创建主要时调用builder中的create_network函数。
这一部分就是创建Network
flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) network = builder.create_network(flag) parser = trt.OnnxParser(network, logger)
此时的network还是一个创建的空网络，通过下面的各属性也能看出来，后续我们会把v5网络写入。
parser = trt.OnnxParser(network,logger)是加载onnx解析器。
加载一行此时的network变为下面这样，可以看到num_inputs和num_layers以及num_outputs均有改变：
5.获得网络的输入输出
inputs = [network.get_input(i) for i in range(network.num_inputs)] outputs = [network.get_output(i) for i in range(network.num_outputs)]
6.判断是否支持FP16推理
if builder.platform_has_fast_fp16 and half: config.set_flag(trt.BuilderFlag.FP16)
7.engine写入
写入engine文件需要调用前面定义的builder.build_engine函数，这里会传入两个参数，第一个就是我们定义好的网络，第二个就是针对网络的相关配置config【比如是否开发FP16等】。写入的网络也是序列化的。
实际就是生成网络TRT的内部表示。
# build engine 文件的写入 这里的f是前面定义的engine文件 with builder.build_engine(network, config) as engine, open(f, 'wb') as t: # 序列化model t.write(engine.serialize())
