当前位置：首页 > news >正文

使用Tensorrt的一般步骤

news 来源：原创 2025/5/10 22:25:40

使用Tensorrt的一般步骤

TensorRT的使用包括两个阶段：build and deployment。

build：该阶段主要完成模型转换（从caffe或TensorFlow到TensorRT），如下图所示，在模型转换时会完成前述优化过程中的层间融合，精度校准。这一步的输出是一个针对特定GPU平台和网络模型的优化过的TensorRT模型，这个TensorRT模型可以序列化存储到磁盘或内存中。存储到磁盘中的文件称之为 planfile。

Deploy：该阶段主要完成推理过程，如下图所示。将上一个步骤中的plan文件首先反序列化，并创建一个 runtime engine，然后就可以输入数据（比如测试集或数据集之外的图片），然后输出分类向量结果或检测结果。

以onnx模型为例检测介绍，主要分为3步，如下图所示，第一步是导入模型，这包括从磁盘上保存的文件加载模型，并将其从原始框架转换为TensorRT网络。ONNX是表示深度学习模型的标准，使它们能够在框架之间传输（Caffe2、Chainer、CNTK、paddle、PyTorch和MXNet都支持ONNX格式）。接下来，基于输入模型、目标GPU平台和指定的其他配置参数，构建一个优化的TensorRT引擎。最后一步是向TensorRT引擎提供输入数据以执行推理。

需要用的tensorrt的组件如下：

ONNX解析器：以ONNX格式的经过训练的模型作为输入，并用TensorRT填充网络对象
Builder：在TensorRT中获取一个网络并生成一个为目标平台优化的引擎
Engine：获取输入数据，执行推理并发出推理输出
Logger：与生成器和引擎关联的对象，用于在生成和推断阶段捕获错误、警告和其他信息

大家可以根据以下coding的例子进行相关实验：

>> git clone https://github.com/parallel-forall/code-samples.git
>> cd code-samples/posts/TensorRT-introduction
>> wget https://s3.amazonaws.com/onnx-model-zoo/resnet/resnet50v2/resnet50v2.tar.gz // Get ONNX model and test data
>> tar xvf resnet50v2.tar.gz # unpack model data into resnet50v2 folder
>> apt-get update
>> apt install libprotobuf-dev protobuf-compiler        # install protobuf to read the input data which is in .pb format
>> git clone --recursive https://github.com/onnx/onnx.git # pull onnx repository from github
>> cd onnx
>> cmake . # compile and install onnx
>> make install -j12
>> cd ..
>> make # compile the TensorRT C++ sample code