利用python版tensorRT导出engine【以yolov5为例】

环境说明：

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_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内部表示以及可执行程序引擎都是由该对象的成员方法生成   

常用成员函数：

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<

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())

---

 

 注：本文会在后续学习中不定时更新。