手册   游戏   WEB   字典   单词   在线工具  

TensorRT 开始

TensorRT 是 NVIDIA 自家的高性能推理库, 其 Getting Started 列出了各资料入口, 如下:

本文基于当前的 TensorRT 8.2 版本, 将一步步介绍从安装, 直到加速推理自己的 ONNX 模型.

安装

进 TensorRT 下载页 选择版本下载, 需注册登录.

本文选择了 TensorRT-8.2.2.1.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz, 可以注意到与 CUDA cuDNN 要匹配好版本. 也可以准备 NVIDIA Docker 拉取对应版本的 nvidia/cuda 镜像, 再 ADD TensorRT 即可.

  1. # 解压进 $HOME (以免 sudo 编译样例, 为当前用户)
  2. tar -xzvf TensorRT-*.tar.gz -C $HOME/
  3. # 软链到 /usr/local/TensorRT (以固定一个路径)
  4. sudo ln -s $HOME/TensorRT-8.2.2.1 /usr/local/TensorRT

之后, 编译运行样例, 保证 TensorRT 安装正确.

编译样例

样例在 TensorRT/samples, 说明见 Sample Support Guide 或各样例目录里的 README.md.

  1. cd /usr/local/TensorRT/samples/
  2. # 设定环境变量, 可见 Makefile.config
  3. export CUDA_INSTALL_DIR=/usr/local/cuda
  4. export CUDNN_INSTALL_DIR=/usr/local/cuda
  5. export ENABLE_DLA=
  6. export TRT_LIB_DIR=../lib
  7. export PROTOBUF_INSTALL_DIR=
  8. # 编译
  9. make -j`nproc`
  10. # 运行
  11. export LD_LIBRARY_PATH=/usr/local/TensorRT/lib:$LD_LIBRARY_PATH
  12. cd /usr/local/TensorRT/
  13. ./bin/trtexec -h
  14. ./bin/sample_mnist -d data/mnist/ --fp16

运行结果参考:

  1. $ ./bin/sample_mnist -d data/mnist/ --fp16
  2. &&&& RUNNING TensorRT.sample_mnist [TensorRT v8202] # ./bin/sample_mnist -d data/mnist/ --fp16
  3. [12/23/2021-20:20:16] [I] Building and running a GPU inference engine for MNIST
  4. [12/23/2021-20:20:16] [I] [TRT] [MemUsageChange] Init CUDA: CPU +322, GPU +0, now: CPU 333, GPU 600 (MiB)
  5. [12/23/2021-20:20:16] [I] [TRT] [MemUsageSnapshot] Begin constructing builder kernel library: CPU 333 MiB, GPU 600 MiB
  6. [12/23/2021-20:20:16] [I] [TRT] [MemUsageSnapshot] End constructing builder kernel library: CPU 468 MiB, GPU 634 MiB
  7. [12/23/2021-20:20:17] [I] [TRT] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +518, GPU +224, now: CPU 988, GPU 858 (MiB)
  8. [12/23/2021-20:20:17] [I] [TRT] [MemUsageChange] Init cuDNN: CPU +114, GPU +52, now: CPU 1102, GPU 910 (MiB)
  9. [12/23/2021-20:20:17] [I] [TRT] Local timing cache in use. Profiling results in this builder pass will not be stored.
  10. [12/23/2021-20:20:33] [I] [TRT] Some tactics do not have sufficient workspace memory to run. Increasing workspace size may increase performance, please check verbose output.
  11. [12/23/2021-20:20:34] [I] [TRT] Detected 1 inputs and 1 output network tensors.
  12. [12/23/2021-20:20:34] [I] [TRT] Total Host Persistent Memory: 8448
  13. [12/23/2021-20:20:34] [I] [TRT] Total Device Persistent Memory: 1626624
  14. [12/23/2021-20:20:34] [I] [TRT] Total Scratch Memory: 0
  15. [12/23/2021-20:20:34] [I] [TRT] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 2 MiB, GPU 13 MiB
  16. [12/23/2021-20:20:34] [I] [TRT] [BlockAssignment] Algorithm ShiftNTopDown took 0.01595ms to assign 3 blocks to 8 nodes requiring 57857 bytes.
  17. [12/23/2021-20:20:34] [I] [TRT] Total Activation Memory: 57857
  18. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +0, GPU +8, now: CPU 1621, GPU 1116 (MiB)
  19. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] Init cuDNN: CPU +0, GPU +8, now: CPU 1621, GPU 1124 (MiB)
  20. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +4, now: CPU 0, GPU 4 (MiB)
  21. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] Init CUDA: CPU +0, GPU +0, now: CPU 1622, GPU 1086 (MiB)
  22. [12/23/2021-20:20:34] [I] [TRT] Loaded engine size: 1 MiB
  23. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +0, GPU +8, now: CPU 1622, GPU 1096 (MiB)
  24. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] Init cuDNN: CPU +1, GPU +8, now: CPU 1623, GPU 1104 (MiB)
  25. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] TensorRT-managed allocation in engine deserialization: CPU +0, GPU +1, now: CPU 0, GPU 1 (MiB)
  26. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +0, GPU +8, now: CPU 1485, GPU 1080 (MiB)
  27. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] Init cuDNN: CPU +0, GPU +8, now: CPU 1485, GPU 1088 (MiB)
  28. [12/23/2021-20:20:34] [I] [TRT] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +2, now: CPU 0, GPU 3 (MiB)
  29. [12/23/2021-20:20:34] [I] Input:
  30. @@@@@@@@@@@@@@@@@@@@@@@@@@@@
  31. @@@@@@@@@@@@@@@@@@@@@@@@@@@@
  32. @@@@@@@@@@@@@@@@@@@@@@@@@@@@
  33. @@@@@@@@@@@@@@@@@@@@@@@@@@@@
  34. @@@@@@@@@@@@@@@@@@@@@@@@@@@@
  35. @@@@@@@@@@@@@@@@@@@@@@@@@@@@
  36. @@@@@@@@@% -:  =@@@@@@@@@@@@
  37. @@@@@@@%=      -@@@**@@@@@@@
  38. @@@@@@@   :%#@-#@@@. #@@@@@@
  39. @@@@@@*  +@@@@:*@@@  *@@@@@@
  40. @@@@@@#  +@@@@ @@@%  @@@@@@@
  41. @@@@@@@.  :%@@.@@@. *@@@@@@@
  42. @@@@@@@@-   =@@@@. -@@@@@@@@
  43. @@@@@@@@@%:   +@- :@@@@@@@@@
  44. @@@@@@@@@@@%.  : -@@@@@@@@@@
  45. @@@@@@@@@@@@@+   #@@@@@@@@@@
  46. @@@@@@@@@@@@@@+  :@@@@@@@@@@
  47. @@@@@@@@@@@@@@+   *@@@@@@@@@
  48. @@@@@@@@@@@@@@: =  @@@@@@@@@
  49. @@@@@@@@@@@@@@ :@  @@@@@@@@@
  50. @@@@@@@@@@@@@@ -@  @@@@@@@@@
  51. @@@@@@@@@@@@@# +@  @@@@@@@@@
  52. @@@@@@@@@@@@@* ++  @@@@@@@@@
  53. @@@@@@@@@@@@@*    *@@@@@@@@@
  54. @@@@@@@@@@@@@#   =@@@@@@@@@@
  55. @@@@@@@@@@@@@@. +@@@@@@@@@@@
  56. @@@@@@@@@@@@@@@@@@@@@@@@@@@@
  57. @@@@@@@@@@@@@@@@@@@@@@@@@@@@
  58. [12/23/2021-20:20:34] [I] Output:
  59. 0:
  60. 1:
  61. 2:
  62. 3:
  63. 4:
  64. 5:
  65. 6:
  66. 7:
  67. 8: **********
  68. 9:
  69. &&&& PASSED TensorRT.sample_mnist [TensorRT v8202] # ./bin/sample_mnist -d data/mnist/ --fp16

快速开始

Quick Start Guide / Using The TensorRT Runtime API

准备教程代码, 编译:

  1. Git clone --depth 1 https://github.com/NVIDIA/TensorRT.git
  2. export CUDA_INSTALL_DIR=/usr/local/cuda
  3. export CUDNN_INSTALL_DIR=/usr/local/cuda
  4. export TRT_LIB_DIR=/usr/local/TensorRT/lib
  5. # 编译 quickstart
  6. cd TensorRT/quickstart
  7. # Makefile.config
  8. #  INCPATHS += -I"/usr/local/TensorRT/include"
  9. # common/logging.h
  10. #  void log(Severity severity, const char* msg) noexcept override
  11. make
  12. # 运行环境
  13. export PATH=/usr/local/TensorRT/bin:$PATH
  14. export LD_LIBRARY_PATH=/usr/local/TensorRT/lib:$LD_LIBRARY_PATH
  15. cd SemanticSegmentation

获取预训练 FCN-ResNet-101 模型, 转成 ONNX:

  1. # 创建本地环境
  2. #  conda create -n torch python=3.9 -y
  3. #  conda activate torch
  4. #  conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch -y
  5. # 不然, 容器环境
  6. #  docker run --rm -it --gpus all -p 8888:8888 -v `pwd`:/workspace/SemanticSegmentation -w /workspace nvcr.io/nvidia/pytorch:20.12-py3 bash
  7. $ python export.py
  8. Exporting ppm image input.ppm
  9. Downloading: "https://github.com/pytorch/vision/archive/v0.6.0.zip" to /home/john/.cache/torch/hub/v0.6.0.zip
  10. Downloading: "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth" to /home/john/.cache/torch/hub/checkpoints/resnet101-5d3b4d8f.pth

100%|████████████████████████████████████████| 170M/170M [00:27<00:00, 6.57MB/s]

Downloading: "https://download.pytorch.org/models/fcn_resnet101_coco-7ecb50ca.pth" to /home/john/.cache/torch/hub/checkpoints/fcn_resnet101_coco-7ecb50ca.pth

100%|████████████████████████████████████████| 208M/208M [02:26<00:00, 1.49MB/s]

Exporting ONNX model fcn-resnet101.onnx

再用 trtexec 将 ONNX 转成 TensorRT engine:

  1. $ trtexec --onnx=fcn-resnet101.onnx --fp16 --workspace=64 --minShapes=input:1x3x256x256 --optShapes=input:1x3x1026x1282 --maxShapes=input:1x3x1440x2560 --buildOnly --saveEngine=fcn-resnet101.engine
  2. ...
  3. [01/07/2022-20:20:00] [I] Engine built in 406.011 sec.
  4. &&&& PASSED TensorRT.trtexec [TensorRT v8202] ...

随机输入, 测试 engine:

  1. $ trtexec --shapes=input:1x3x1026x1282 --loadEngine=fcn-resnet101.engine
  2. ...
  3. [01/07/2022-20:20:00] [I] === Performance summary ===
  4. [01/07/2022-20:20:00] [I] Throughput: 12.4749 qps
  5. [01/07/2022-20:20:00] [I] Latency: min = 76.9746 ms, max = 98.8354 ms, mean = 79.5844 ms, median = 78.0542 ms, percentile(99%) = 98.8354 ms
  6. [01/07/2022-20:20:00] [I] End-to-End Host Latency: min = 150.942 ms, max = 188.431 ms, mean = 155.834 ms, median = 152.444 ms, percentile(99%) = 188.431 ms
  7. [01/07/2022-20:20:00] [I] Enqueue Time: min = 0.390625 ms, max = 1.61279 ms, mean = 1.41182 ms, median = 1.46136 ms, percentile(99%) = 1.61279 ms
  8. [01/07/2022-20:20:00] [I] H2D Latency: min = 1.25977 ms, max = 1.53467 ms, mean = 1.27415 ms, median = 1.26514 ms, percentile(99%) = 1.53467 ms
  9. [01/07/2022-20:20:00] [I] GPU Compute Time: min = 75.2869 ms, max = 97.1318 ms, mean = 77.8847 ms, median = 76.3599 ms, percentile(99%) = 97.1318 ms
  10. [01/07/2022-20:20:00] [I] D2H Latency: min = 0.408447 ms, max = 0.454346 ms, mean = 0.425577 ms, median = 0.423004 ms, percentile(99%) = 0.454346 ms
  11. [01/07/2022-20:20:00] [I] Total Host Walltime: 3.2866 s
  12. [01/07/2022-20:20:00] [I] Total GPU Compute Time: 3.19327 s
  13. [01/07/2022-20:20:00] [I] Explanations of the performance metrics are printed in the verbose logs.
  14. [01/07/2022-20:20:00] [I]
  15. &&&& PASSED TensorRT.trtexec [TensorRT v8202] ...

运行教程, 使用 engine:

  1. $ ./bin/segmentation_tutorial
  2. [01/07/2022-20:20:34] [I] [TRT] [MemUsageChange] Init CUDA: CPU +322, GPU +0, now: CPU 463, GPU 707 (MiB)
  3. [01/07/2022-20:20:34] [I] [TRT] Loaded engine size: 132 MiB
  4. [01/07/2022-20:20:35] [I] [TRT] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +520, GPU +224, now: CPU 984, GPU 1065 (MiB)
  5. [01/07/2022-20:20:35] [I] [TRT] [MemUsageChange] Init cuDNN: CPU +115, GPU +52, now: CPU 1099, GPU 1117 (MiB)
  6. [01/07/2022-20:20:35] [I] [TRT] [MemUsageChange] TensorRT-managed allocation in engine deserialization: CPU +0, GPU +131, now: CPU 0, GPU 131 (MiB)
  7. [01/07/2022-20:20:35] [I] Running TensorRT inference for FCN-ResNet101
  8. [01/07/2022-20:20:35] [I] [TRT] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +0, GPU +10, now: CPU 966, GPU 1109 (MiB)
  9. [01/07/2022-20:20:35] [I] [TRT] [MemUsageChange] Init cuDNN: CPU +0, GPU +8, now: CPU 966, GPU 1117 (MiB)
  10. [01/07/2022-20:20:35] [I] [TRT] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +722, now: CPU 0, GPU 853 (MiB)

实践

以上给到了官方样例与教程的编译使用. 这里, 另外找了个 RVM 的模型, 从头开始试一试.

准备模型

Robust Video Matting (RVM) 稳定视频抠像, 可在任意视频上做实时高清抠像. 有 Webcam Demo 可以网页上体验.

准备 ONNX 模型 rvm_mobilenetv3_fp32.onnx, 其 推断文档 给出了模型输入输出:

输入: [src, r1i, r2i, r3i, r4i, downsample_ratio]

src: 输入帧, RGB 通道, 形状为 [B, C, H, W], 范围为 0~1

rXi: 记忆输入, 初始值是是形状为 [1, 1, 1, 1] 的零张量

downsample_ratio 下采样比, 张量形状为 [1]

只有 downsample_ratio 必须是 FP32, 其他输入必须和加载的模型使用一样的 dtype

输出: [fgr, pha, r1o, r2o, r3o, r4o]

fgr, pha: 前景和透明度通道输出, 范围为 0~1

rXo: 记忆输出

准备输入图像 input.jpg . 不用视频, 保持代码简单些.

准备环境

  1. Anaconda
  2. PyTorch
  3. conda create -n torch python=3.9 -y
  4. conda activate torch
  5. conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch -y
  6. # Requirements
  7. #  https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements
  8. pip install onnx onnxruntime-gpu==1.10

运行 ONNX 模型

  1. rvm_onnx_infer.py:
  2. import onnxruntime as ort
  3. import numpy as np
  4. from PIL import Image
  5. # 读取图像
  6. with Image.open('input.jpg') as img:
  7.     img.load()
  8. #  HWC [0,255]> BCHW [0,1]
  9. src = np.array(img)
  10. src = np.moveaxis(src, -1, 0) .astype(np.float32)
  11. src = src[np.newaxis, :] / 255.
  12. # 载入模型
  13. sess = ort.InferenceSession('rvm_mobilenetv3_fp32.onnx', providers=['CUDAExecutionProvider'])
  14. # 创建 io binding
  15. io = sess.io_binding()
  16. # 在 CUDA 上创建张量
  17. rec = [ ort.OrtValue.ortvalue_from_numpy(np.zeros([1, 1, 1, 1], dtype=np.float32), 'cuda') ] * 4
  18. downsample_ratio = ort.OrtValue.ortvalue_from_numpy(np.asarray([0.25], dtype=np.float32), 'cuda')
  19. # 设置输出项
  20. for name in ['fgr', 'pha', 'r1o', 'r2o', 'r3o', 'r4o']:
  21.     io.bind_output(name, 'cuda')
  22. # 推断
  23. io.bind_cpu_input('src', src)
  24. io.bind_ortvalue_input('r1i', rec[0])
  25. io.bind_ortvalue_input('r2i', rec[1])
  26. io.bind_ortvalue_input('r3i', rec[2])
  27. io.bind_ortvalue_input('r4i', rec[3])
  28. io.bind_ortvalue_input('downsample_ratio', downsample_ratio)
  29. sess.run_with_iobinding(io)
  30. fgr, pha, *rec = io.get_outputs()
  31. # 只将 `fgr` 和 `pha` 回传到 CPU
  32. fgr = fgr.numpy()
  33. pha = pha.numpy()
  34. # 合成 RGBA
  35. com = np.where(pha> 0, fgr, pha)
  36. com = np.concatenate([com, pha], axis=1) # + alpha
  37. #  BCHW [0,1]> HWC [0,255]
  38. com = np.squeeze(com, axis=0)
  39. com = np.moveaxis(com, 0, -1) * 255
  40. img = Image.fromarray(com.astype(np.uint8))
  41. img.show()

运行:

python rvm_onnx_infer.py --model "rvm_mobilenetv3_fp32.onnx" --input-image "input.jpg" --precision float32 --show

结果 (背景透明):

ONNX 转成 TRT 模型

trtexec 将 ONNX 转成 TensorRT engine:

  1. export PATH=/usr/local/TensorRT/bin:$PATH
  2. export LD_LIBRARY_PATH=/usr/local/TensorRT/lib:$LD_LIBRARY_PATH
  3. trtexec --onnx=rvm_mobilenetv3_fp32.onnx --workspace=64 --saveEngine=rvm_mobilenetv3_fp32.engine --verbose

发生问题:

  1. [01/08/2022-20:20:36] [E] [TRT] ModelImporter.cpp:773: While parsing node number 3 [Resize -> "389"]:
  2. [01/08/2022-20:20:36] [E] [TRT] ModelImporter.cpp:774: --- Begin node ---
  3. [01/08/2022-20:20:36] [E] [TRT] ModelImporter.cpp:775: input: "src"
  4. input: "386"
  5. input: "388"
  6. output: "389"
  7. name: "Resize_3"
  8. op_type: "Resize"
  9. attribute {
  10.   name: "coordinate_transformation_mode"
  11.   s: "pytorch_half_pixel"
  12.   type: STRING
  13. }
  14. attribute {
  15.   name: "cubic_coeff_a"
  16.   f: -0.75
  17.   type: FLOAT
  18. }
  19. attribute {
  20.   name: "mode"
  21.   s: "linear"
  22.   type: STRING
  23. }
  24. attribute {
  25.   name: "nearest_mode"
  26.   s: "floor"
  27.   type: STRING
  28. }
  29. [01/08/2022-20:20:36] [E] [TRT] ModelImporter.cpp:776: --- End node ---
  30. [01/08/2022-20:20:36] [E] [TRT] ModelImporter.cpp:779: ERROR: builtin_op_importers.cpp:3608 In function importResize:
  31. [8] Assertion failed: scales.is_weights() && "Resize scales must be an initializer!"

这时, 需要动手改动模型了.

首先, 安装必要工具:

  1. snap install netron
  2. pip install onnx-simplifier
  3. pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com

之后, Netron 查看模型 Resize_3 节点:

发现其 scales 输入是依据 downsample_ratio 得到的, 即 [1,1,downsample_ratio,downsample_ratio], 可用 ONNX GraphSurgeon 修改成常量.

最后, 模型改动步骤如下:

  1. # ONNX 模型简化, 并改为静态输入尺寸
  2. python -m onnxsim rvm_mobilenetv3_fp32.onnx rvm_mobilenetv3_fp32_sim.onnx \
  3. --input-shape src:1,3,1080,1920 r1i:1,1,1,1 r2i:1,1,1,1 r3i:1,1,1,1 r4i:1,1,1,1
  4. # ONNX GraphSurgeon 修改模型
  5. python rvm_onnx_modify.py -i rvm_mobilenetv3_fp32_sim.onnx --input-size 1920 1280
  6. # trtexec 将 ONNX 转成 TensorRT engine
  7. trtexec --onnx=rvm_mobilenetv3_fp32_sim_modified.onnx --workspace=64 --saveEngine=rvm_mobilenetv3_fp32_sim_modified.engine
  8. rvm_onnx_modify.py:
  9. def modify(input: str, output: str, downsample_ratio: float = 0.25) -> None:
  10.     print(f'\nonnx load: {input}')
  11.     graph = gs.import_onnx(onnx.load(input))
  12.     _print_graph(graph)
  13.     # update node Resize_3: scales
  14.     resize_3 = [n for n in graph.nodes if n.name == 'Resize_3'][0]
  15.     print()
  16.     print(resize_3)
  17.     scales = gs.Constant('388',
  18.         np.asarray([1, 1, downsample_ratio, downsample_ratio], dtype=np.float32))
  19.     resize_3.inputs = [i if i.name != '388' else scales for i in resize_3.inputs]
  20.     print()
  21.     print(resize_3)
  22.     # remove input downsample_ratio
  23.     graph.inputs = [i for i in graph.inputs if i.name != 'downsample_ratio']
  24.     # remove node Concat_2
  25.     concat_2 = [n for n in graph.nodes if n.name == 'Concat_2'][0]
  26.     concat_2.outputs.clear()
  27.     # remove unused nodes/tensors
  28.     graph.cleanup()
  29.     onnx.save(gs.export_onnx(graph), output)

ONNX 与 TRT 模型输出差异

可用 Polygraphy 查看 ONNX 与 TRT 模型的输出差异.

首先, 安装

  1. # 安装 TensorRT Python API
  2. cd /usr/local/TensorRT/python/
  3. pip install tensorrt-8.2.2.1-cp39-none-linux_x86_64.whl
  4. export LD_LIBRARY_PATH=/usr/local/TensorRT/lib:$LD_LIBRARY_PATH
  5. python -c "import tensorrt; print(tensorrt.__version__)"
  6. # 安装 Polygraphy, 或者通过 TensorRT/tools/Polygraphy 源码安装
  7. python -m pip install colored polygraphy --extra-index-url https://pypi.ngc.nvidia.com

运行 ONNX 与 TRT 模型, 对比输出误差:

  1. # 运行 ONNX 模型, 保存输入输出
  2. polygraphy run rvm_mobilenetv3_fp32_sim_modified.onnx --onnxrt --val-range [0,1] --save-inputs onnx_inputs.JSON --save-outputs onnx_outputs.JSON
  3. # 运行 TRT 模型, 载入 ONNX 输入输出, 对比输出的相对误差与绝对误差
  4. polygraphy run rvm_mobilenetv3_fp32_sim_modified.engine --model-type engine --trt --load-inputs onnx_inputs.JSON --load-outputs onnx_outputs.JSON --rtol 1e-3 --atol 1e-3

可见 fp32 精度误差在 1e-3 以内, PASSED:

  1. [I]     PASSED | All outputs matched | Outputs: ['r4o', 'r3o', 'r2o', 'r1o', 'fgr', 'pha']
  2. [I] PASSED | Command: /home/john/anaconda3/envs/torch/bin/polygraphy run rvm_mobilenetv3_fp32_sim_modified.engine --model-type engine --trt --load-inputs onnx_inputs.JSON --load-outputs onnx_outputs.JSON --rtol 1e-3 --atol 1e-3

也试了 fp16, 其精度损失就比较大, FAILED:

  1. [E]     FAILED | Mismatched outputs: ['r4o', 'r3o', 'r2o', 'r1o', 'fgr', 'pha']
  2. [!] FAILED | Command: /home/john/anaconda3/envs/torch/bin/polygraphy run rvm_mobilenetv3_fp16_sim_modified.engine --model-type engine --trt --load-inputs onnx_inputs.JSON --load-outputs onnx_outputs.JSON --rtol 1e-3 --atol 1e-3

运行 TRT 模型

这里以 TensorRT C++ runtime APIs 为例, 将转出的 RVM TRT 模型运行起来. 完整代码见 rvm_infer.cc.

1. 载入模型: 创建 runtime, 反序列化 TRT 模型文件的数据

  1. static Logger logger{
  2. 	Logger::Severity::kINFO	
  3. };
  4. auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));
  5. auto engine = runtime->deserializeCudaEngine(engine_data.data(), fsize, nullptr);

遍历全部输入输出 bindings:

  1. auto nb = engine->getNbBindings();
  2. for (int32_t i = 0; i <nb; i++) {
  3.   auto is_input = engine->bindingIsInput(i);
  4.   auto name = engine->getBindingName(i);
  5.   auto dims = engine->getBindingDimensions(i);
  6.   auto datatype = engine->getBindingDataType(i);
  7.   // ...
  8. }
  9. Engine
  10.  Name=Unnamed Network 0
  11.  DeviceMemorySize=148 MiB
  12.  MaxBatchSize=1
  13. Bindings
  14.  Input[0] name=src dims=[1,3,1080,1920] datatype=FLOAT
  15.  Input[1] name=r1i dims=[1,1,1,1] datatype=FLOAT
  16.  Input[2] name=r2i dims=[1,1,1,1] datatype=FLOAT
  17.  Input[3] name=r3i dims=[1,1,1,1] datatype=FLOAT
  18.  Input[4] name=r4i dims=[1,1,1,1] datatype=FLOAT
  19.  Output[5] name=r4o dims=[1,64,18,32] datatype=FLOAT
  20.  Output[6] name=r3o dims=[1,40,36,64] datatype=FLOAT
  21.  Output[7] name=r2o dims=[1,20,72,128] datatype=FLOAT
  22.  Output[8] name=r1o dims=[1,16,144,256] datatype=FLOAT
  23.  Output[9] name=fgr dims=[1,3,1080,1920] datatype=FLOAT
  24.  Output[10] name=pha dims=[1,1,1080,1920] datatype=FLOAT

之后, 分配好所有 bindings 的 device 内存:

  1. auto nb = engine->getNbBindings();
  2. std::vector<void *> bindings(nb, nullptr);
  3. std::vector<int32_t> bindings_size(nb, 0);
  4. for (int32_t i = 0; i <nb; i++) {
  5.   auto dims = engine->getBindingDimensions(i);
  6.   auto size = GetMemorySize(dims, sizeof(float));
  7.   if (cudaMalloc(&bindings[i], size) != cudaSuccess) {
  8.     std::cerr <<"ERROR: cuda memory allocation failed, size =" << size
  9.         << "bytes" << std::endl;
  10.     return false;
  11.   }
  12.   bindings_size[i] = size;
  13. }

到此, 准备工作就好了.

2. 前处理: 输入数据处理成输入格式, 存进输入 bindings

用 OpenCV 读取图像, 缩放成 src 的输入尺寸. 再把数据从 BGR [0,255] 处理成 RGB [0,1]. 因 batch=1, 所以处理时可忽略.

  1. // img: HWC BGR [0,255] u8
  2. auto img = cv::imread(input_filename, cv::IMREAD_COLOR);
  3. if (src_h != img.rows || src_w != img.cols) {
  4.   cv::resize(img, img, cv::Size(src_w, src_h));
  5. }
  6. // src: BCHW RGB [0,1] fp32
  7. auto src = cv::Mat(img.rows, img.cols, CV_32FC3);
  8. {
  9.   auto src_data = (float*)(src.data);
  10.   for (int y = 0; y < src_h; ++y) {
  11.     for (int x = 0; x < src_w; ++x) {
  12.       auto &&bgr = img.at<cv::Vec3b>(y, x);
  13.       /*r*/*(src_data + y*src_w + x) = bgr[2] / 255.;
  14.       /*g*/*(src_data + src_n + y*src_w + x) = bgr[1] / 255.;
  15.       /*b*/*(src_data + src_n*2 + y*src_w + x) = bgr[0] / 255.;
  16.     }
  17.   }
  18. }
  19. if (cudaMemcpyAsync(bindings[0], src.data, bindings_size[0],
  20.     cudaMemcpyHostToDevice, stream) != cudaSuccess) {
  21.   std::cerr <<"ERROR: CUDA memory copy of src failed, size ="
  22.       << bindings_size[0] << "bytes" << std::endl;
  23.   return false;
  24. }

3. 推理: 将 bindings 给到 engine 执行上下文进行推理

  1. auto context = std::unique_ptr<nvinfer1::IExecutionContext>(
  2.     engine->createExecutionContext());
  3. if (!context) {
  4.   return false;
  5. }
  6. bool status = context->enqueueV2(bindings.data(), stream, nullptr);
  7. if (!status) {
  8.   std::cout <<"ERROR: TensorRT inference failed" << std::endl;
  9.   return false;
  10. }

4. 后处理: 从输出 bindings 取出数据, 根据输出格式处理数据

用 cv::Mat 接收输出的前景 fgr 和透明通道 pha:

  1. auto fgr = cv::Mat(src_h, src_w, CV_32FC3);  // BCHW RGB [0,1] fp32
  2. if (cudaMemcpyAsync(fgr.data, bindings[9], bindings_size[9],
  3.     cudaMemcpyDeviceToHost, stream) != cudaSuccess) {
  4.   std::cerr << "ERROR: CUDA memory copy of output failed, size ="
  5.       << bindings_size[9] << "bytes" << std::endl;
  6.   return false;
  7. }
  8. auto pha = cv::Mat(src_h, src_w, CV_32FC1);  // BCHW A [0,1] fp32
  9. if (cudaMemcpyAsync(pha.data, bindings[10], bindings_size[10],
  10.     cudaMemcpyDeviceToHost, stream) != cudaSuccess) {
  11.   std::cerr << "ERROR: CUDA memory copy of output failed, size ="
  12.       << bindings_size[10] << "bytes" << std::endl;
  13.   return false;
  14. }
  15. cudaStreamSynchronize(stream);

再将 fgr pha 合成 RGBA 数据, 并复原成原尺寸:

  1. // Compose `fgr` and `pha`
  2. auto com = cv::Mat(src_h, src_w, CV_8UC4);  // HWC BGRA [0,255] u8
  3. {
  4.   auto fgr_data = (float*)(fgr.data);
  5.   auto pha_data = (float*)(pha.data);
  6.   for (int y = 0; y < com.rows; ++y) {
  7.     for (int x = 0; x < com.cols; ++x) {
  8.       auto &&elem = com.at<cv::Vec4b>(y, x);
  9.       auto alpha = *(pha_data + y*src_w + x);
  10.       if (alpha> 0) {
  11.         /*r*/ elem[2] = *(fgr_data + y*src_w + x) * 255;
  12.         /*g*/ elem[1] = *(fgr_data + src_n + y*src_w + x) * 255;
  13.         /*b*/ elem[0] = *(fgr_data + src_n*2 + y*src_w + x) * 255;
  14.       } else {
  15.         /*r*/ elem[2] = 0;
  16.         /*g*/ elem[1] = 0;
  17.         /*b*/ elem[0] = 0;
  18.       }
  19.       /*a*/ elem[3] = alpha * 255;
  20.     }
  21.   }
  22. }
  23. if (dst_h != com.rows || dst_w != com.cols) {
  24.   cv::resize(com, com, cv::Size(dst_w, dst_h));
  25. }

5. 运行得到的抠像结果 (背景透明):

最后

想入门 TensorRT 的, 动手实践一下吧!

GoCoding 个人实践的经验分享, 可关注公众号!

来源: https://segmentfault.com/a/1190000041273519


与本文相关文章

  1. 我的 Python3.x 的开始 - 2018.8.4
  2. 从 0 开始的 Python 学习 016 异常
  3. python 之开始
  4. 如何开始学习 Go 语言
  5. 从头开始：用Python实现随机森林算法
  6. 开始写文章了
  7. 开始学习 python 的第五天
  8. 从头开始：用Python实现决策树算法
暂无,快来抢沙发吧！
更多





黑ICP备17005318号-1