Python API构建TensorRT加速模型的步骤详解

ytusdc

已于 2024-07-08 09:48:58 修改

阅读量1.2k

点赞数 18

文章标签： python 开发语言

于 2024-07-04 18:07:00 首次发布

本文链接：https://blog.csdn.net/ytusdc/article/details/140185771

版权

写了一半不想写了，看这个 https://github.com/ytusdc/TensorRT-NMS-YOLO， utils/export.py 吧

先来一段摘抄自网上的TensorRT介绍：

TensorRT是英伟达针对自家平台做的加速包，TensorRT主要做了这么两件事情，来提升模型的运行速度。

TensorRT支持INT8和FP16的计算。深度学习网络在训练时，通常使用 32 位或 16 位数据。TensorRT则在网络的推理时选用不这么高的精度，达到加速推断的目的。
TensorRT对于网络结构进行了重构，把一些能够合并的运算合并在了一起，针对GPU的特性做了优化。现在大多数深度学习框架是没有针对GPU做过性能优化的，而英伟达，GPU的生产者和搬运工，自然就推出了针对自己GPU的加速工具TensorRT。一个深度学习模型，在没有优化的情况下，比如一个卷积层、一个偏置层和一个reload层，这三层是需要调用三次cuDNN对应的API，但实际上这三层的实现完全是可以合并到一起的，TensorRT会对一些可以合并网络进行合并。我们通过一个典型的inception block来看一看这样的合并运算。

TensorRT用来做模型的推理优化，也是有Python接口的，实际使用测试下来，python接口的模型推理速度C++基本差不多的。这里较为详细的记录TensorRT python接口从环境的配置到模型的转换，再到推理过程，还有模型的INT8量化，有时间的话也一并总结记录了，笔者使用的版本是TensorRT7.0版本，此版本支持模型动态尺寸的前向推理，下面也会分为静态推理和动态推理来介绍。

3.3.1 builder介绍

builder功能之一是搜索cuda内核目录，找到最快的cuda以求获得最快的实现，因此有必要使用相同的GPU进行构建（相同的操作，算子进行融合，减少IO操作），engine就是在此基础上运行的，builder还可以控制网络以什么精度运行（FP32,FP16,INT8），还有两个特别重要的属性是最大批处理大小和最大工作空间大小。

安装完 tensorRT后， pip安装pycuda

pip install pycuda

TensorRT 模型转换

简单来说就是用tensorrt的语言来构建模型，如果自己构建的话，主要是灵活但是工作量so large，一般还是用tensorrt parser来构建。自己将wts放入自己构建的模型中，工作量so large，但是很灵活。博主用的ONNXparser来构建engine的，下面就介绍以下ONNX构建engine

1、创建logger：日志记录器

要创建构建器，您需要首先创建一个记录器，为tensorrt实现日志报错接口方便报错。 TensorRT中包含一个简单的日志记录器与Python绑定。包括一个简单的记录器实现，它将高于特定严重性的所有消息记录到stdout 。在下面的代码我们只允许警告和错误消息才打印

logger = trt.Logger(trt.Logger.WARNING)

2、构建一个序列构建器

build a serialized network from scratch

builder = trt.Builder(logger)

3、构建一个空的网络计算图

Creating a Network Definition in Python

创建构建器后，优化模型的第一步是创建网络定义，

network_flags = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
self.network = builder.create_network(network_flags)

设置网络读取的 flag， EXPLICIT_BATCH 相较于 IMPLICIT_BATCH 模式，会显示的将 batch 的维度包含在张量维度当中，有了 batch大小的，我们就可以进行一些必须包含 batch 大小的操作了，如 Layer Normalization。不然在推理阶段，应当指定推理的 batch 的大小。目前主流的使用的 EXPLICIT_BATCH 模式

为了使用 ONNX 解析器导入模型，需要EXPLICIT_BATCH标志。有关详细信息，请参阅显式与隐式批处理部分。

4、创建一个构建配置

create BuidlerConfig to set meta data of the network

指定 TensorRT 应该如何优化模型：

config = builder.create_builder_config()

这个接口有很多属性，你可以设置这些属性来控制 TensorRT 如何优化网络。一个重要的属性是最大工作空间大小。层实现通常需要一个临时工作空间，并且此参数限制了网络中任何层可以使用的最大大小。如果提供的工作空间不足，TensorRT 可能无法找到层的实现：

set workspace for the optimization process (default value is total GPU memory)

变换：

2^30 = 1<<30 = 1G

1<<20：1M

config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

4.1.2. Importing a Model using the ONNX Parser

现在，需要从 ONNX 表示中填充网络定义。您可以创建一个 ONNX 解析器来填充网络，如下所示：

将空的网络计算图和相应的 logger 设置装载进一个解析器里面

parser = trt.OnnxParser(network, logger)

with open(onnx_path, "rb") as f_r:
    parser_flag = parser.parse(f_r.read())
    if not parser_flag:
        print(f"ERROR: Failed to parse the ONNX file: {onnx_path}")
        for error in range(parser.num_errors):
            # 出错了，将相关错误的地方打印出来，进行可视化处理
            print(parser.num_errors)
            print(parser.get_error(error))
        # sys.exit(1)
        return None

指定配置后，可以使用以下命令构建和序列化引擎：

serialized_engine = builder.build_serialized_network(network, config)

将引擎保存到文件以供将来使用可能很有用。你可以这样做：

with open(“sample.engine”, “wb”) as f:
    f.write(serialized_engine)

https://blog.51cto.com/u_16099272/10992508

https://zhuanlan.zhihu.com/p/645010134

6. TensorRT 进阶用法 - NVIDIA 技术博客

TensorRT_Plugin:手写Plugin详细步骤教程_tensorrt plugin-CSDN博客

TensorRT教程15：使用Python API 添加自定义插件层_python trt插件如何使用-CSDN博客

从0实现-＞训练pytorch模型-＞转onnx-＞tensorrt模型序列化-＞Tensorrt推理_torch .th转 onnx-CSDN博客

TensorRT的Python接口解析_tensorrt10.0.0中executev2-CSDN博客

yolov5 C++部署学习笔记_nvinfer1::iscalelayer* addbatchnorm2d(nvinfer1::in-CSDN博客

较为详细的记录总结TensorRT的python接口的使用，环境配置，模型转换和静态动态模型推理_pycuda 调用tensorrt-CSDN博客