环境说明：
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格式的导出。
git clone https://github.com/ultralytics/yolov5
我们可以看一下官方提供的yolov5s.pt中都包含什么内容：
dict_keys([‘epoch’, ‘best_fitness’, ‘model’, ’ema’, ‘updates’, ‘optimizer’, ‘wandb_id’, ‘date’])
可以看到上述pt文件中包含了这些key值，其中的model就是我们要的，而且需要注意的是这个model不仅含有网络权重信息，还包含了整个网络结构【如果你想把其他网络转onnx，也需要主要必须torch保存的是整个网络】
目录
导出onnx
导出engine
导出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_onnxexport_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内部表示以及可执行程序引擎都是由该对象的成员方法生成
常用成员函数：
builder.max_batch_size = 256 # 指定最大batch_size(static shape模型下使用)
builder.max_workspace_size = 1 <<30# 指定最大可用显存
builder.fp16_model = True/Flase# 开启或者关闭FP16模式
builder.int8_model = True/Flase# # 开启或者关闭int8模式注意：(上面的成员变量即将弃用，上面这些设置将被放置在Config中，而builder仅仅是作为一个构建引擎的入口)
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 << 304.网络定义并加载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())
注：本文会在后续学习中不定时更新。
文章出处登录后可见！
