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

这篇具有很好参考价值的文章主要介绍了利用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就是我们前面定义好的输入结点。 

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

下面这一部分就是输出部分,输出是三个头进行了整合,结点为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) 

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


在介绍导出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内容,主要是一些内存上面的使用初始等。

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

[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网络写入。 

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

 parser = trt.OnnxParser(network,logger)是加载onnx解析器。

加载一行此时的network变为下面这样,可以看到num_inputs和num_layers以及num_outputs均有改变:

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

 

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

 

 注:本文会在后续学习中不定时更新。文章来源地址https://www.toymoban.com/news/detail-410227.html

到了这里,关于利用python版tensorRT导出engine【以yolov5为例】的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 使用TensorRT对yolov5进行部署(基于python,超详细)

    哈喽大家好 ! 我是唐宋宋宋,很荣幸与您相见!!! 我的环境: cuda:11.1 cudnn:8.7.0 TensorRT:8.4.1.5 首先需要下载TensorRT,官网链接附下: NVIDIA TensorRT 8.x Download | NVIDIA Developer 注:下载TensorRT版本需要和你电脑上的cuda版本对应 yolov5的代码需要大家上github自己扒 链接已经提供

    2024年02月07日
    浏览(42)
  • 【TensorRT】TensorRT 部署Yolov5模型(C++)

      该项目代码在本人GitHub代码仓库开源,本人GitHub主页为:GitHub   项目代码:   NVIDIA TensorRT™ 是用于高性能深度学习推理的 SDK,可为深度学习推理应用提供低延迟和高吞吐量。详细安装方式参考以下博客: NVIDIA TensorRT 安装 (Windows C++)   经典的一个TensorRT部署模型步骤为

    2023年04月26日
    浏览(48)
  • 【YOLOv5】LabVIEW+TensorRT的yolov5部署实战(含源码)

    ‍‍🏡博客主页: virobotics的CSDN博客:LabVIEW深度学习、人工智能博主 🎄所属专栏:『LabVIEW深度学习实战』 🍻上期文章: 手把手教你使用LabVIEW TensorRT实现图像分类实战(含源码) 📰如觉得博主文章写的不错或对你有所帮助的话,还望大家多多支持呀! 欢迎大家✌关注、

    2024年02月14日
    浏览(31)
  • jetson nx目标检测环境配置遇到的一万个坑,安装v1.12.0版本的pytorch和v0.13.0版本的vision torchvision,以及使用TensorRT部署YOLOv5.

    本文参考了许多官网和博客,肯定是存在抄袭的,请各位大哥不要喷我啊。 自己工作找到的是医学信号方向的算法工程师,所以以后和CV可能无缘了,将自己一个多星期的心血历程发表出来,希望大家接起我的CV火炬,接着前行,各位加油!(后面也学习了yolov5-6.0 yolov7的模型

    2024年02月05日
    浏览(46)
  • 【TensorRT】基于C#调用TensorRT 部署Yolov5模型 - 上篇:构建TensorRTSharp

      NVIDIA TensorRT™ 是用于高性能深度学习推理的 SDK,可为深度学习推理应用提供低延迟和高吞吐量。详细安装方式参考以下博客: NVIDIA TensorRT 安装 (Windows C++)   前文中已经介绍了在C++中利用TensorRT 部署Yolov5模型,但在实际应用中,经常会出现在C#中部署模型的需求,目前T

    2023年04月24日
    浏览(50)
  • 【目标检测】YOLOv5推理加速实验:TensorRT加速

    前两篇博文分别讨论了YOLOv5检测算法的两种加速思路:采用多进程或批量检测,不过效果均收效甚微。本问将讨论使用TensorRT加速以及半精度推理/模型量化等优化策略对检测加速的实际影响。 测试图片分辨率:13400x9528 GPU:RTX4090 Cuda:11.7 YOLOv5版本:最新版(v7.0+) 检测策略:将

    2024年02月05日
    浏览(53)
  • 改进的yolov5目标检测-yolov5替换骨干网络-yolo剪枝(TensorRT及NCNN部署)

    2022.10.30 复现TPH-YOLOv5 2022.10.31 完成替换backbone为Ghostnet 2022.11.02 完成替换backbone为Shufflenetv2 2022.11.05 完成替换backbone为Mobilenetv3Small 2022.11.10 完成EagleEye对YOLOv5系列剪枝支持 2022.11.14 完成MQBench对YOLOv5系列量化支持 2022.11.16 完成替换backbone为EfficientNetLite-0 2022.11.26 完成替换backbone为

    2024年01月17日
    浏览(68)
  • 【YOLOv5】1.搭建Pycharm+Python+yolov5环境

    目录 一、安装Python 二、安装PyCharm 三、创建项目和虚拟环境 四、下载YOLOv5和依赖库 五、配置Pytorch 六、检验YOLOv5环境 1.Python官方下载网址:Download Python | Python.org 2.安装python3.10即可,笔者使用的pytorch版本在python3.10下不可用,所以多装了一个python3.9。 3.心得:如果安装了多个

    2024年02月13日
    浏览(39)
  • Jetson AGX Xavier实现TensorRT加速YOLOv5进行实时检测

    link 上一篇:Jetson AGX Xavier安装torch、torchvision且成功运行yolov5算法 下一篇:Jetson AGX Xavier测试YOLOv4         由于YOLOv5在Xavier上对实时画面的检测速度较慢,需要采用TensorRT对其进行推理加速。接下来记录一下我的实现过程。  如果还没有搭建YOLOv5的python环境,按照下文步骤

    2024年02月10日
    浏览(40)
  • YOLOV5使用(一): docker跑通,详解TensorRT下plugin的onnx

    使用ubuntu为案例 使用端口映射功能也就是说打开jupyter lab的指令是 当然,个人建议直接去vscode端口点击就打开jupyter lab和tensorboard比较方便 YOLO格式的标签文件是一个纯文本文件,每个文件名对应一张图像,每个标签文件中包含了该图像中所有检测到的目标的信息。 YOLOv5的标

    2024年02月04日
    浏览(43)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包