TensorRT使用说明

一、TensorRT简介

TensorRT 是Nvidia 提出的深度学习推理平台，能够在GPU 上实现低延迟、高吞吐量的部属。基于TensorRT 的推论运行速度会比仅使用CPU 快40倍，提供精度INT8 和FP16 优化，支援TensorFlow、Caffe、Mxnet、Pytorch 等深度学习框架，其中Mxnet、Pytorch 需先转换为ONNX 格式。

trt-info.png

TensorRT 函式库以C++撰写，提供C++ 和Python API。

• TensorRT 的 API 具有 C++ 和 Python 的语言绑定，具有几乎相同的功能。Python API 促进了与 Python 数据处理工具包和库（如 NumPy 和 SciPy）的互操作性，C++ API 可以更高效。

1.TensorRT 的优化方法

1）层间融合或张量融合

TensorRT通过对层间的横向或纵向合并（合并后的结构称为CBR，意指 convolution, bias, and ReLU layers are fused to form a single layer），使得层的数量大大减少。合并之后的计算图的层次更少了，占用的CUDA核心数也少了，因此整个模型结构会更小，更快，更高效。

• 横向合并可以把卷积、偏置和激活层合并成一个CBR结构，只占用一个CUDA核心。
• 纵向合并可以把结构相同，但是权值不同的层合并成一个更宽的层，也只占用一个CUDA核心。

2）数据精度校准

大部分深度学习框架在训练神经网络时网络中的张量都是32位浮点数的精度（Full 32-bit precision，FP32），一旦网络训练完成，在部署推理的过程中由于不需要反向传播，完全可以适当降低数据精度，比如降为FP16或INT8的精度。更低的数据精度将会使得内存占用和延迟更低，模型体积更小。

3）Kernel Auto-Tuning

网络模型在推理计算时，是调用 GPU 的 CUDA 进行计算的，TensorRT 可以针对不同的算法、不同的模型结构、不同的 GPU 平台等，进行 CUDA 调整，以保证当前模型在特定平台上以最优的性能计算。

• 假设在 3090 和 T4 上要分别部署，需要分别在这两个平台上进行 TensorRT 的转换，然后在对应的平台上使用，而不能在相同同的平台上转换，在不同的平台上使用。

4）Dynamic Tensor Memory

在每个 tensor 使用期间，TensorRT 会为其指定显存，避免显存重复申请，减少内存占用和提高重复使用效率。

2.相关软件

• NVIDIA Triton™推理服务器是一个更高级别的库，可提供跨 CPU 和 GPU 的优化推理。它提供了启动和管理多个模型的功能，以及用于服务推理的 REST 和 gRPC 端点。

• NVIDIA DALI ®为预处理图像、音频和视频数据提供高性能原语。TensorRT 推理可以作为自定义算子集成到 DALI 管道中。可以在此处找到作为 DALI 的一部分集成的 TensorRT 推理的工作示例。

• TensorFlow-TensorRT (TF-TRT)是将 TensorRT 直接集成到 TensorFlow 中。它选择 TensorFlow 图的子图由 TensorRT 加速，同时让图的其余部分由 TensorFlow 本地执行。结果仍然是您可以照常执行的 TensorFlow 图。有关 TF-TRT 示例，请参阅TensorFlow 中的 TensorRT 示例。

• PyTorch 量化工具包提供了以降低精度训练模型的工具，然后可以将其导出以在 TensorRT 中进行优化 。

• PyTorch Automatic SParsity (ASP)工具提供了用于训练具有结构化稀疏性的模型的工具，然后可以将其导出并允许 TensorRT 在 NVIDIA Ampere GPU 上利用更快的稀疏策略。

• TensorRT 与 NVIDIA 的分析工具、NVIDIA Nsight™ Systems和NVIDIA® Deep Learning Profiler (DLProf)集成。

二、安装TensorRT及其依赖

参考

1.安装TensorRT

1）查看tensorrt能力支持

根据tensorrt compatibility确定对CUDA、cuDNN、Pytorch和ONNX版本的支持。例如：

[图片上传失败...(image-7469eb-1650335485672)]

TensorRT 支持的环境和Python 版本关系

2）执行安装

TensorRT 有四种安装方式: 使用Debian, RPM, Tar, Zip 档案，其中Zip 档案只支持Windows。

通过TensorRT官网https://developer.nvidia.com/nvidia-tensorrt-download下载所需的TensorRT版本，以Ubuntu 18.04.5 LTS，CUDA 10.2, cuDNN 7.6.5, python 3.7.9 环境为例：

直接安装到系统

# 要记得换成刚下载的deb 安装包
sudo dpkg -i nv-tensorrt-repo-ubuntu1804-cuda10.2-trt7.1.3.4-ga-20200617_1–1_amd64.deb
sudo apt-key add /var/nv-tensorrt-repo-cuda10.2-trt7.1.3.4-ga-20200617/7fa2af80.pub
sudo apt-get update
sudo apt-get install tensorrt 
    # 如果所需依赖项因为版本匹配问题导致无法安装，需要指定版本，如：
    apt-get install libnvinfer-dev=7.1.3-1+cuda10.2 -y
    apt-get install libnvinfer-plugin-dev=7.1.3-1+cuda10.2
    apt-get install libnvparsers-dev=7.1.3-1+cuda10.2 -y
    apt-get install libnvonnxparsers-dev=7.1.3-1+cuda10.2 -y
    apt-get install libnvinfer-samples=7.1.3-1+cuda10.2 -y


# If using Python 2.7:
sudo apt-get install python-libnvinfer-dev
# If using Python 3.x:
sudo apt-get install python3-libnvinfer-dev

# 验证安装
$ dpkg -l | grep TensorRT

通过tar包安装tensorrt

# 下载TensorRT-7.1.3.4.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz后，进行解压
tar -xzvf TensorRT-7.1.3.4.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz
cd TensorRT-7.1.3.4/python
# 安装跟python 版本一样的whl
pip install tensorrt-7.1.3.4-cp37-none-linux_x86_64.whl

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/share/TensorRT-7.1.3.4/lib

# 用于编译C++ tensorrt代码
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/share/TensorRT-7.1.3.4/targets/x86_64-linux-gnu/lib/
export C_INCLUDE_PATH=/home/share/TensorRT-7.1.3.4/include
export CPLUS_INCLUDE_PATH=/home/share/TensorRT-7.1.3.4/include

2.安装pytorch（及cuda）

基于Anaconda，根据所需的cuda版本安装pytorch。

# CUDA 10.2
conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch
# CUDA 11.3
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
# CPU
conda install pytorch torchvision torchaudio cpuonly -c pytorch

# 查看pytorch版本
$ python
import torch
torch.__version__
    '1.8.1+cu102'
torch.version.cuda
10.2
torch.backends.cudnn.version()  # cudnn版本
7605
    
# 查看cuda版本
$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Wed_Oct_23_19:24:38_PDT_2019
Cuda compilation tools, release 10.2, V10.2.89

# 检验驱动是否安装成功
$ nvidia-smi -L
GPU 0: GeForce RTX 2080 Ti (UUID: GPU-09b6f910-63d7-e56d-0885-0862fea0bc0c)
GPU 1: GeForce RTX 2080 Ti (UUID: GPU-38d16434-c477-cea8-150a-b88cce8aaacb)
GPU 2: GeForce RTX 2080 Ti (UUID: GPU-94eb0194-e443-7214-b6a5-d430989dd703)

3.安装cudnn

点击cuDNN Archive下载链接，按照要求进行下载。

# 安装cudnn， 根据实际版本修改
tar -xzvf cudnn-11.2-linux-x64-v8.1.1.33.tgz
sudo cp cuda/include/* /usr/local/cuda/include/
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/include/cudnn.h
sudo chmod a+r /usr/local/cuda/lib64/libcudnn*

# 查看cudnn版本
cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2
# 或
$ python
import torch
torch.backends.cudnn.version()  

三、TensorRT 样例

TensorRT 官方范例

1.验证demo

https://blog.csdn.net/irving512/article/details/107165757

进入TensorRT的解压目录

# 数据准备
$ cd TensorRT-7.1.3.4/data/mnist
$ python download_pgms.py   
# 下载了一些 *.pgm
$ ls *.pgm
0.pgm  1.pgm  2.pgm  3.pgm  4.pgm  5.pgm  6.pgm  7.pgm  8.pgm  9.pgm

# 代码编译与运行
$ cd TensorRT-7.1.3.4/samples/sampleMNIST
$ make
    # 可执行文件生成在 ./bin/ 目录下
$ cd TensorRT-7.1.3.4/
$ ./bin/sample_mnist
[04/12/2022-11:48:05] [I] Building and running a GPU inference engine for MNIST
[04/12/2022-11:48:11] [I] [TRT] Detected 1 inputs and 1 output network tensors.
[04/12/2022-11:48:11] [I] Input:
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@#- -#@@@@@@@@
@@@@@@@@@@@@@@#     @@@@@@@@
@@@@@@@@@@@@@#.     #@@@@@@@
@@@@@@@@@@@@#.   :*  +@@@@@@
@@@@@@@@@@@-      *: -@@@@@@
@@@@@@@@@@#   :+ .%* -@@@@@@
@@@@@@@@@#   :@*+@@@  #@@@@@
@@@@@@@@%-  .*@@@@@@  -@@@@@
@@@@@@@@:  #@%@@@@@@  :@@@@@
@@@@@@@#  #@@@@@@@@@  :@@@@@
@@@@@@@: :@@@@@@@@@@  :@@@@@
@@@@@@*  +@@@@@@@@@@  =@@@@@
@@@@@@*  %@@@@@@@@@= :@@@@@@
@@@@@@* .@@@@@@@@@= .#@@@@@@
@@@@@@* =@@@@@@@#- -@@@@@@@@
@@@@@@* .@@@@@@+  -@@@@@@@@@
@@@@@@*  =#%*:. .-#@@@@@@@@@
@@@@@@*   ..   :=@@@@@@@@@@@
@@@@@@%:      =@@@@@@@@@@@@@
@@@@@@@%=   =%@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@

源码的基本流程

• 第一步：将训练好的神经网络模型转换为TensorRT的形式，并用TensorRT Optimizer进行优化。
• 第二步：将在TensorRT Engine中运行优化好的TensorRT网络结构。

2.python demo

$ cd TensorRT-7.1.3.4/samples/python/yolov3_onnx
# 安装依赖
$ pip install -r requirements.txt
# 下载yolov3 权重并转为onnx
$ python yolov3_to_onnx.py
# build TensorRT engine from the generated ONNX file and run inference on a sample image
$ python onnx_to_tensorrt.py
Loading ONNX file from path yolov3.onnx...
Beginning ONNX file parsing
Completed parsing of ONNX file
Building an engine from file yolov3.onnx; this may take a while...
Completed creating Engine
Running inference on image dog.jpg...
[[135.14841345 219.59879372 184.30209058 324.0265199 ]
 [ 98.30804939 135.72612864 499.71263113 299.2558099 ]
 [478.00607009  81.25702312 210.57786012  86.91503109]] [0.99854713 0.99880403 0.93829261] [16  1  7]
 $ python onnx_to_tensorrt.py
Reading engine from file yolov3.trt
Running inference on image dog.jpg...
[[135.14841345 219.59879372 184.30209058 324.0265199 ]
 [ 98.30804939 135.72612864 499.71263113 299.2558099 ]
 [478.00607009  81.25702312 210.57786012  86.91503109]] [0.99854713 0.99880403 0.93829261] [16  1  7]
Saved image with bounding boxes of detected objects to dog_bboxes.png.

四、pytorch转TensorRT部署

1.ONNX简介

ONNX 是 Open Neural Network Exchange 的简称，也叫开放神经网络交换，是一个用于表示深度学习模型的标准，可使模型在不同框架直接转换。

在深度学习模型落地的过程中，会面临将模型部署到边端设备的问题，模型训练使用不同的框架，则推理的时候也需要使用相同的框架，但不同类型的平台，调优和实现起来非常困难，因为每个平台都有不同的功能和特性。使用ONNX可以通过将不同框架训练的模型转换成通用的 ONNX 模型，再进而转换成各个平台支持的格式，就可以实现简化部署。

ONNX 目前支持的框架有：Caffe2、PyTorch、TensorFlow、MXNet、TensorRT、CNTK 等。

[图片上传失败...(image-7442c4-1650335485672)]

常用部署方案

• cpu: pytorch/tensorflow/Caffe->onnx->onnxruntime

• gpu: pytorch/tensorflow/Caffe->onnx->onnx2trt->tensorRT

• arm: pytorch/tensorflow/Caffe->onnx->ncnn/mace/mnn等

2.PyTorch to ONNX

只要模型中所有OP均被ONNX支持，即可利用Pytorch中的ONN库进行转换。

需要提供的有：加载好的Pytorch模型、一个输入样例。

• 其中模型需要按照自己的方式导入并加载模型；
• 输入样例的格式为BCHW，B为batch_size，CHW为通道、高、宽，CHW的值需要与你自己的模型相匹配，否则后面转换成功后输出结果也不对。

#--*-- coding:utf-8 --*--
import onnx 
import torch
import torchvision 
import netron

net = torchvision.models.resnet18(pretrained=True).cuda()   # pretrained=True：使用模型自带预训练模型
# net.eval()

export_onnx_file = "./resnet18.onnx"
x=torch.onnx.export(net,  # 待转换的网络模型和参数
                torch.randn(1, 3, 224, 224, device='cuda'), # 虚拟的输入，用于确定输入尺寸和推理计算图每个节点的尺寸
                export_onnx_file,  # 输出文件的名称
                verbose=False,      # 是否以字符串的形式显示计算图
                input_names=["input"]+ ["params_%d"%i for i in range(120)],  # 输入节点的名称，这里也可以给一个list，list中名称分别对应每一层可学习的参数，便于后续查询
                output_names=["output"], # 输出节点的名称
                opset_version=10,   # onnx 支持采用的operator set, 应该和pytorch版本相关，目前我这里最高支持10
                do_constant_folding=True, # 是否压缩常量
                dynamic_axes={"input":{0: "batch_size", 2: "h"}, "output":{0: "batch_size"},} #设置动态维度，此处指明input节点的第0维度可变，命名为batch_size
                )

# import onnx  # 注意这里导入onnx时必须在torch导入之前，否则会出现segmentation fault
net = onnx.load("./resnet18.onnx")  # 加载onnx 计算图
onnx.checker.check_model(net)  # 检查文件模型是否正确
onnx.helper.printable_graph(net.graph)  # 输出onnx的计算图


# 使用onnxruntime进行cpu推理onnx
import onnxruntime
import numpy as np

# netron.start("./resnet18.onnx")

session = onnxruntime.InferenceSession("./resnet18.onnx") # 创建一个运行session，类似于tensorflow
out_r = session.run(None, {"input": np.random.rand(16, 3, 256, 224).astype('float32')})  # 模型运行，注意这里的输入必须是numpy类型
print(len(out_r))   # 输出是list类型
print(out_r[0].shape)

• 如果出现“RuntimeError: ONNX export failed: Couldn't export Python operator XXXX”错误提示，说明模型中有ONNX不支持的OP，可以尝试升级Pytorch版本，或者编写自定义op。
• 注意：dynamic_axes可能会导致转tenorrt时失败，谨慎使用。

3.onnxruntime验证onnx模型【可选】

onnxruntime基于cpu推理，以试验转换的onnx模型是否正常。

import cv2
import numpy as np
import onnxruntime
 
def image_process(image_path):
    mean = np.array([[[0.485, 0.456, 0.406]]])      # 训练的时候用来mean和std
    std = np.array([[[0.229, 0.224, 0.225]]])
 
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (96, 96))                 # (96, 96, 3)
 
    image = img.astype(np.float32)/255.0
    image = (image - mean)/ std
 
    image = image.transpose((2, 0, 1))              # (3, 96, 96)
    image = image[np.newaxis,:,:,:]                 # (1, 3, 96, 96)
 
    image = np.array(image, dtype=np.float32)
    
    return image
 
def onnx_runtime():
    imgdata = image_process('test.jpg')
    
    sess = onnxruntime.InferenceSession('test.onnx')
    input_name = sess.get_inputs()[0].name  
    output_name = sess.get_outputs()[0].name
    pred_onnx = sess.run([output_name], {input_name: imgdata})  # 定义一路输入，一路输出
    
    # pred_onnx = session.run(None, input_feed={input_name: imgdata}) # 定义一路输入，可多路输出
 
    print("outputs:")
    print(np.array(pred_onnx))
 
onnx_runtime()

4.ONNX to TensorRT

https://github.com/qq995431104/Pytorch2TensorRT

根据生成的ONNX文件，将其转换为TensorRT engine，过程如下：

• 先以trt的Logger为参数，使用TensorRT创建一个builder，然后用builder创建一个network。

   with trt.Builder(G_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, \
              trt.OnnxParser(network, G_LOGGER) as parser:
  

• 指定builder的参数设置，如：max_batch_size、max_workspace_size；

  • 如需转为特定格式，如fp16或int8，需指定相应参数：fp16_mode或int8_mode设为True；

• 利用对应的Parser（OnnxParser、CaffPaser、UffParser）加载ONNX文件，解析得到网络架构，并填充计算图。

  • 由于onnx文件包含模型网络信息，因此并不需要重构网络。
  • 当然也可以使用tensorrt的API手动构建网络（c程序wts转engine需要这种方式）。

  with open(args.onnx_file_path, 'rb') as model:
              print('Beginning ONNX file parsing')
              parser.parse(model.read())
  

• builder以network为参数，创建engine。

  engine = builder.build_cuda_engine(network)
  

• 将engine序列化为字符串，写入到engine文件。

  with open(args.engine_file_path, "wb") as f:
              f.write(engine.serialize())
  

• 可以读取engine文件并反序列化为engine。【可选】

  G_LOGGER = trt.Logger(trt.Logger.WARNING)
      # 反序列化引擎
      with open(filepath, "rb") as f, trt.Runtime(G_LOGGER) as runtime:
          engine = runtime.deserialize_cuda_engine(f.read())
          return engine
  

1）通过python代码转换

import tensorrt as trt

def ONNX2TRT(args, calib=None):
    ''' convert onnx to tensorrt engine, use mode of ['fp32', 'fp16', 'int8']
    :return: trt engine
    '''

    assert args.mode.lower() in ['fp32', 'fp16', 'int8'], "mode should be in ['fp32', 'fp16', 'int8']"

    G_LOGGER = trt.Logger(trt.Logger.WARNING)
    # TRT7中的onnx解析器的network，需要指定EXPLICIT_BATCH
    EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    with trt.Builder(G_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, \
            trt.OnnxParser(network, G_LOGGER) as parser:

        builder.max_batch_size = args.batch_size
        builder.max_workspace_size = 1 << 30
        if args.mode.lower() == 'int8':
            assert (builder.platform_has_fast_int8 == True), "not support int8"
            builder.int8_mode = True
            builder.int8_calibrator = calib
        elif args.mode.lower() == 'fp16':
            assert (builder.platform_has_fast_fp16 == True), "not support fp16"
            builder.fp16_mode = True

        print('Loading ONNX file from path {}...'.format(args.onnx_file_path))
        with open(args.onnx_file_path, 'rb') as model:
            print('Beginning ONNX file parsing')
            if not parser.parse(model.read()):
                for e in range(parser.num_errors):
                    print(parser.get_error(e))
                raise TypeError("Parser parse failed.")

        print('Completed parsing of ONNX file')

        print('Building an engine from file {}; this may take a while...'.format(args.onnx_file_path))
        engine = builder.build_cuda_engine(network)
        print("Created engine success! ")

        # 保存计划文件
        print('Saving TRT engine file to path {}...'.format(args.engine_file_path))
        with open(args.engine_file_path, "wb") as f:
            f.write(engine.serialize())
        print('Engine file has already saved to {}!'.format(args.engine_file_path))
        return engine

• 对于Int8格式，需要：
  • 准备一个校准集，用于在转换过程中寻找使得转换后的激活值分布与原来的FP32类型的激活值分布差异最小的阈值;
  • 并写一个校准器类，该类需继承trt.IInt8EntropyCalibrator2父类，并重写get_batch_size, get_batch, read_calibration_cache, write_calibration_cache这几个方法。具体做法参考脚本myCalibrator.py.
  • 使用时，需额外指定cache_file，该参数是校准集cache文件的路径，会在校准过程中生成，方便下一次校准时快速提取。

2）通过trtexec工具转换

trtexec的参数使用说明

cd TensorRT-7.1.3.4/bin

#生成静态batchsize的engine
./trtexec   --onnx= \                        #指定onnx模型文件
            ----maxBatch=1                          # max batch size，默认为1
            --explicitBatch \                           #在构建引擎时使用显式批大小(默认=隐式)显示批处理
            --saveEngine= \       #输出engine
            --workspace= \           #设置工作空间大小单位是MB(默认为16MB)
            --fp16                                      #除了fp32之外，还启用fp16精度(默认=禁用)
        
#生成动态batchsize的engine
./trtexec   --onnx= \                        #指定onnx模型文件
            --minShapes=input: \    #最小的batchsize x 通道数 x 输入尺寸x x 输入尺寸y
            --optShapes=input: \    #最佳输入维度，跟maxShapes一样就好
            --maxShapes=input: \    #最大输入维度
            --workspace= \           #设置工作空间大小单位是MB(默认为16MB)
            --saveEngine= \                #输出engine
            --fp16   

# 示例：
./trtexec  --onnx=yolov4_-1_3_416_416_dynamic.onnx \
                                        --minShapes=input:1x3x416x416 \
                                        --optShapes=input:8x3x416x416 \
                                        --maxShapes=input:8x3x416x416 \
                                        --workspace=4096 \
                                        --saveEngine=yolov4_-1_3_416_416_dynamic_b8_fp16.engine \
                                        --fp16

./trtexec --onnx=onnx/mobilenet_v2.onnx --saveEngine=engine/mobilenet_v2_b1_fp16.engine --workspace=4096 --fp16
----------------------------------------------------------------
Input filename:   onnx/mobilenet_v2.onnx
ONNX IR version:  0.0.6
Opset version:    9
Producer name:    pytorch
Producer version: 1.8
Domain:           
Model version:    0
Doc string:       
----------------------------------------------------------------
[04/16/2022-08:40:21] [I] Host Latency
[04/16/2022-08:40:21] [I] min: 0.490479 ms (end to end 0.496094 ms)
[04/16/2022-08:40:21] [I] max: 4.66577 ms (end to end 4.79541 ms)
[04/16/2022-08:40:21] [I] mean: 0.50234 ms (end to end 0.748152 ms)
[04/16/2022-08:40:21] [I] median: 0.497772 ms (end to end 0.793213 ms)
[04/16/2022-08:40:21] [I] percentile: 0.538818 ms at 99% (end to end 0.828125 ms at 99%)
[04/16/2022-08:40:21] [I] throughput: 2295.34 qps
[04/16/2022-08:40:21] [I] walltime: 3.00129 s
[04/16/2022-08:40:21] [I] Enqueue Time
[04/16/2022-08:40:21] [I] min: 0.217041 ms
[04/16/2022-08:40:21] [I] max: 4.73901 ms
[04/16/2022-08:40:21] [I] median: 0.243286 ms
[04/16/2022-08:40:21] [I] GPU Compute
[04/16/2022-08:40:21] [I] min: 0.414551 ms
[04/16/2022-08:40:21] [I] max: 4.58044 ms
[04/16/2022-08:40:21] [I] mean: 0.425963 ms
[04/16/2022-08:40:21] [I] median: 0.422241 ms
[04/16/2022-08:40:21] [I] percentile: 0.454712 ms at 99%
[04/16/2022-08:40:21] [I] total compute time: 2.93446 s

• 注意：保存为.trt和.engine文件没有区别

5.TensorRT Python推理

推理过程完全独立于原先模型所依赖的框架，基本过程如下：

• 按照原模型的输入输出格式，准备数据，如：输入的shape、均值、方差，输出的shape等；

• 根据得到的引擎文件，利用TensorRT Runtime反序列化为引擎engine；

• 创建上下文环境；

  context = engine.create_execution_context()
  

• 使用Pycuda的mem_alloc对输入输出分配cuda内存；

  d_input = cuda.mem_alloc(1 * input.size * input.dtype.itemsize)
  d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
  bindings = [int(d_input), int(d_output)]
  

• 创建Stream，即pycuda操作缓冲区；

  stream = cuda.Stream()
  

• 使用memcpy_htod_async将IO数据放入device（一般为GPU）；

  cuda.memcpy_htod_async(d_input, input, stream)
  

• 使用context.execute_async执行推理（异步）；

  context.execute_async(batch_size, bindings, stream.handle, None)
  

• 使用memcpy_dtoh_async取出结果：从cuda从缓冲区取出结果并复制到cpu；

  cuda.memcpy_dtoh_async(output, d_output, stream)
  

• 对输出结果进行后处理，因模型而异。

1）直接编写推理程序

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from torchvision import transforms
import numpy as np
from PIL import Image
import time
import argparse

def loadEngine2TensorRT(filepath):
    G_LOGGER = trt.Logger(trt.Logger.WARNING)
    # 反序列化引擎
    with open(filepath, "rb") as f, trt.Runtime(G_LOGGER) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())
        return engine

def do_inference(engine, batch_size, input, output_shape):

    # 创建上下文
    context = engine.create_execution_context()
    output = np.empty(output_shape, dtype=np.float32)

    # 分配内存
    d_input = cuda.mem_alloc(1 * input.size * input.dtype.itemsize)
    d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
    bindings = [int(d_input), int(d_output)]

    # pycuda操作缓冲区
    stream = cuda.Stream()
    # 将输入数据放入device
    cuda.memcpy_htod_async(d_input, input, stream)

    start = time.time()
    # 执行模型
    context.execute_async(batch_size, bindings, stream.handle, None)
    # 将预测结果从从缓冲区取出
    cuda.memcpy_dtoh_async(output, d_output, stream)
    end = time.time()

    # 线程同步
    stream.synchronize()

    print("\nTensorRT {} test:".format(engine_path.split('/')[-1].split('.')[0]))
    print("output:", output)
    print("time cost:", end - start)

def get_shape(engine):
    for binding in engine:
        if engine.binding_is_input(binding):
            input_shape = engine.get_binding_shape(binding)
        else:
            output_shape = engine.get_binding_shape(binding)
    return input_shape, output_shape

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description = "TensorRT do inference")
    parser.add_argument("--batch_size", type=int, default=1, help='batch_size')
    parser.add_argument("--img_path", type=str, default='test_image/1.jpg', help='cache_file')
    parser.add_argument("--engine_file_path", type=str, default='my_files/test.engine', help='engine_file_path')
    args = parser.parse_args()

    engine_path = args.engine_file_path
    engine = loadEngine2TensorRT(engine_path)
    img = Image.open(args.img_path)
    input_shape, output_shape = get_shape(engine)
    transform = transforms.Compose([
        transforms.Resize([input_shape[1], input_shape[2]]),  # [h,w]
        transforms.ToTensor()
        ])
    img = transform(img).unsqueeze(0)
    img = img.numpy()

    do_inference(engine, args.batch_size, img, output_shape)

2）使用官方库torch2trt进行推理

https://github.com/NVIDIA-AI-IOT/torch2trt

def torch2trt(module,
              inputs,
              input_names=None,
              output_names=None,
              log_level=trt.Logger.ERROR,
              max_batch_size=1,
              fp16_mode=False,
              max_workspace_size=1<<25,
              strict_type_constraints=False,
              keep_network=True,
              int8_mode=False,
              int8_calib_dataset=None,
              int8_calib_algorithm=DEFAULT_CALIBRATION_ALGORITHM,
              int8_calib_batch_size=1,
              use_onnx=False,
              **kwargs):

import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()
# create example data
x = torch.ones((1, 3, 224, 224)).cuda()
# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])

优点

• 使用简单，对于Mobilenet、Unet、Resnet等可以直接使用

缺点

• 不支持多输出
• 自定义的tensor或者list都没有_trt这个属性

6.TensorRT C++ 推理

把模型部署到内存有限的嵌入式板的过程：

• 电脑上安装的有anaconda, pytorch等，但是在电脑上转的不能直接在板子上用。
• 板子的内存有限，不能安装anaconda, pytorch这些，但是需要部署模型上去。这时就可以现在电脑上把pth转成wts，再把wts传到板子上，在板子上转成tensorrt。

1）pth转wts

import torch
import torch.nn as nn
from torchvision import models
import struct
from torchsummary import summary

def get_model():
    net = getattr(models.quantization, 'mobilenet_v2')(pretrained=False, num_classes=2, quantize=False)
    net.load_state_dict(torch.load('weights/xxx.pth'))
    net = net.eval().cuda()
    return net

def pth_to_wts(model, wts_name):
    f = open(wts_name, 'w')
    f.write('{}\n'.format(len(model.state_dict().keys())))
    for k, v in model.state_dict().items():
        vr = v.reshape(-1).cpu().numpy()
        f.write('{} {} '.format(k, len(vr)))
        for vv in vr:
            f.write(' ')
            f.write(struct.pack('>f',float(vv)).hex())
        f.write('\n')

if __name__ == '__main__':
    model = get_model()
    summary(model, (3, 256, 256))   # 可选，将模型视觉化，了解模型每一层输入输出
    pth_to_wts(model, "wts/mobilenet_v2.wts")

2）wts转tensorrt

关键头文件

• #include "cuda_runtime_api.h"

  • 路径：/usr/local/cuda-10.2/targets/x86_64-linux/include/cuda_runtime_api.h

• #include "NvInfer.h"

  • 本地路径：/usr/include/x86_64-linux-gnu/NvInfer.h，在线API地址
  • 提供IRuntime、IBuilder、IHostMemory、IExecutionContext、ICudaEngine和英伟达gpu 算子等调用接口。

常用的模型转换可参考Tensorrt C API。

wts转tensorrt的原理

• 从wts文件把weight给load出来，存到一个map里，key是网络每层的名称，value就是对应的权重
• 利用tensorrt的API把网络重建出来，同时导入key对应的value，也就是weightMap的形式
• 定义网络的输出，设置内存空间
• build engine输出一个engine文件

wts转tensorrt示例程序

const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";

// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] 
std::map loadWeights(const std::string file)
{
    std::cout << "Loading weights: " << file << std::endl;
    std::map weightMap;

    // Open weights file
    std::ifstream input(file);
    assert(input.is_open() && "Unable to load weight file.");

    // Read number of weight blobs
    int32_t count;
    input >> count;
    assert(count > 0 && "Invalid weight map file.");

    while (count--)
    {
        Weights wt{DataType::kFLOAT, nullptr, 0};
        uint32_t size;

        // Read name and type of blob
        std::string name;
        input >> name >> std::dec >> size;
        wt.type = DataType::kFLOAT;

        // Load blob
        uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size));
        for (uint32_t x = 0, y = size; x < y; ++x)
        {
            input >> std::hex >> val[x];
        }
        wt.values = val;
        
        wt.count = size;
        weightMap[name] = wt;
    }

    return weightMap;
}

// Creat the engine using only the API and not any parser.  重建网络
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt)
{
    INetworkDefinition* network = builder->createNetworkV2(0U);

    // Create input tensor of shape { 3, INPUT_H, INPUT_W } with name INPUT_BLOB_NAME
    ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W}); // 定义输入层
    assert(data);

    std::map weightMap = loadWeights("../mobilenet.wts");
    Weights emptywts{DataType::kFLOAT, nullptr, 0};

    // 输入传入各卷积层
    auto ew1 = convBnRelu(network, weightMap, *data, 32, 3, 2, 1, "features.0.");
    ILayer* ir1 = invertedRes(network, weightMap, *ew1->getOutput(0), "features.1.", 32, 16, 1, 1);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.2.", 16, 24, 2, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.3.", 24, 24, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.4.", 24, 32, 2, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.5.", 32, 32, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.6.", 32, 32, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.7.", 32, 64, 2, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.8.", 64, 64, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.9.", 64, 64, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.10.", 64, 64, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.11.", 64, 96, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.12.", 96, 96, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.13.", 96, 96, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.14.", 96, 160, 2, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.15.", 160, 160, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.16.", 160, 160, 1, 6);
    ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.17.", 160, 320, 1, 6);
    IElementWiseLayer* ew2 = convBnRelu(network, weightMap, *ir1->getOutput(0), 1280, 1, 1, 1, "features.18.");

    // 池化
    IPoolingLayer* pool1 = network->addPoolingNd(*ew2->getOutput(0), PoolingType::kAVERAGE, DimsHW{7, 7});
    assert(pool1);

    // 全连接
    IFullyConnectedLayer* fc1 = network->addFullyConnected(*pool1->getOutput(0), 1000, weightMap["classifier.1.weight"], weightMap["classifier.1.bias"]);
    assert(fc1);

    fc1->getOutput(0)->setName(OUTPUT_BLOB_NAME);   // 定义输出
    std::cout << "set name out" << std::endl;
    network->markOutput(*fc1->getOutput(0));    // 指定网络输出内存空间

    // Build engine
    builder->setMaxBatchSize(maxBatchSize);     // 设置一个engine内存空间
    config->setMaxWorkspaceSize(1 << 20);
    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    std::cout << "build out" << std::endl;

    // Don't need the network any more
    network->destroy();

    // Release host memory
    for (auto& mem : weightMap)
    {
        free((void*) (mem.second.values));
    }

    return engine;
}

void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream)
{
    // Create builder
    IBuilder* builder = createInferBuilder(gLogger);
    IBuilderConfig* config = builder->createBuilderConfig();

    // Create model to populate the network, then set the outputs and create an engine
    ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
    assert(engine != nullptr);

    // Serialize the engine
    (*modelStream) = engine->serialize();

    // Close everything down
    engine->destroy();
    builder->destroy();
    config->destroy();
}

int main(int argc, char** argv)
{
    if (std::string(argv[1]) == "-s") {
        IHostMemory* modelStream{nullptr};
        APIToModel(1, &modelStream);    // 加载wts文件，利用tensorrt的API把网络重建，并转为engine
        assert(modelStream != nullptr);

        std::ofstream p("mobilenet.engine", std::ios::binary);
        if (!p)
        {
            std::cerr << "could not open plan output file" << std::endl;
            return -1;
        }
        p.write(reinterpret_cast(modelStream->data()), modelStream->size());   // 写入engine数据到文件
        modelStream->destroy();
        return 1;
    }
}

3）engine推理

void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
{
    const ICudaEngine& engine = context.getEngine();

    // Pointers to input and output device buffers to pass to engine.
    // Engine requires exactly IEngine::getNbBindings() number of buffers.
    assert(engine.getNbBindings() == 2);
    void* buffers[2];

    // In order to bind the buffers, we need to know the names of the input and output tensors.
    // Note that indices are guaranteed to be less than IEngine::getNbBindings()
    const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
    const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);

    // Create GPU buffers on device
    CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
    CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));

    // Create stream
    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));

    // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
    CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
    context.enqueue(batchSize, buffers, stream, nullptr);   // 上下文将数据入列，执行推理
    CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));  // 将预测结果从从缓冲区取出
    cudaStreamSynchronize(stream);      // 数据同步

    // Release stream and buffers
    cudaStreamDestroy(stream);
    CHECK(cudaFree(buffers[inputIndex]));
    CHECK(cudaFree(buffers[outputIndex]));
}

int main(int argc, char** argv)
{
    char *trtModelStream{nullptr};
    if (std::string(argv[1]) == "-d") {
        std::ifstream file("mobilenet.engine", std::ios::binary);
        if (file.good()) {
            file.seekg(0, file.end);
            size = file.tellg();
            file.seekg(0, file.beg);
            trtModelStream = new char[size];
            assert(trtModelStream);
            file.read(trtModelStream, size);    // 读取engine文件到trtModelStream
            file.close();
        }
    }
    
    // trtModelStream转ICudaEngine，并创建上下文
    IRuntime* runtime = createInferRuntime(gLogger);
    assert(runtime != nullptr);
    ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
    assert(engine != nullptr);
    IExecutionContext* context = engine->createExecutionContext();
    assert(context != nullptr);
    delete[] trtModelStream;
    
    
    // Subtract mean from image
    static float data[3 * INPUT_H * INPUT_W];
    for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
        data[i] = 1.0;
    // Run inference
    static float prob[OUTPUT_SIZE];
    doInference(*context, data, prob, 1);   // 构建图像进行推理
}

4）tensorrtx

tensorrtx提供多种模型的C版本的的wts转engine和gpu推理程序， github地址：https://github.com/wang-xinyu/tensorrtx。

tensorrtx工程目录

$ git clone https://github.com/wang-xinyu/tensorrtx.git
$ tree -L 1 tensorrtx
├── alexnet
├── arcface
├── crnn
├── dbnet
├── Dockerfile
├── googlenet
├── hrnet
├── inceptionv3
├── lenet
├── LICENSE
├── mnasnet
├── mobilenetv2
├── mobilenetv3
├── psenet
├── README.md
├── resnet
├── retinaface
├── retinafaceAntiCov
├── senet
├── shufflenetv2
├── squeezenet
├── tutorials
├── ufld
├── vgg
├── yolov3
├── yolov3-spp
├── yolov3-tiny
├── yolov4
└── yolov5

参考

• tensorrt developer-guide 【官方】
• TensorRT详细入门指北【推荐】
• yolov5 pytorch转tensorrt注意
• 详细步骤：pytorch pth转wts转tensorrt