Pytorch->ONNX->TensorRT

发表于 更新于 分类于
本文字数： 2.2k 阅读时长 ≈ 4 分钟

Pytorch分类/检测/分割模型转换到ONNX格式，最后转换成TensorRT Engine文件进行推理（Python）。

Pytorch2ONNX

Pytorch提供了转换函数torch.onnx.export进行模型格式转换，关键点在于是否指定动态批量大小。注意，低版本的TensorRT版本不一定适配于高等级的opset_version。

ONNX2TensorRT

对于ONNX模型转换成TensorRT模型，需要设置好输入大小（固定还是动态）和精度（FP32/FP16）。除了调用API来手动转换外，TensorRT提供了命令行转换工具trtexec：

注意：TensorRT自带了trtexec，位于/path/to/tensorrt/bin/trtexec，比如/usr/bin/tensorrt/bin/trtexec

1
2
3
4
# 固定批量大小
trtexec --onnx=resnet18_pytorch.onnx --saveEngine=resnet_engine_pytorch.trt  --explicitBatch
# 半精度转换
trtexec --onnx=resnet18_pytorch.onnx --saveEngine=resnet_engine_pytorch_fp16.trt  --explicitBatch --inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --fp16
  • 模型加载相关参数
    • --onnx=<path>：指定输入模型为 ONNX 格式
    • --explicitBatch：启用显式批量模式（Explicit Batch Mode）
  • 引擎生成与保存相关参数
    • --saveEngine=<path>：将生成的 TensorRT 引擎保存到指定路径
  • 输入数据相关参数
    • --shapes=<input_name>:<shape>：
      • 定义输入张量的形状（适用于动态形状模型）
      • 示例：--shapes=image:32x3x640x640,scale_factor:32x2
        • image:32x3x640x640：表示输入图像张量的形状为 [batch_size, channels, height, width]，即[32, 3, 640, 640]
        • scale_factor:32x2：表示另一个输入张量 scale_factor 的形状为[32, 2]
    • --minShapes=<input_name>:<shape>：
      • 定义动态形状模型的最小输入形状
      • 示例：--minShapes=input:1x3x224x224
    • --maxShapes=<input_name>:<shape>：
      • 定义动态形状模型的最大输入形状
      • 示例：--maxShapes=input:16x3x224x224
    • --workspace=<size_in_MB>
      • 设置 GPU 内存的工作空间大小（以 MB 为单位）
      • 示例：--workspace=1024
    • --inputIOFormats=fp16:chw
      • 定义输入张量的格式为 FP16 和通道优先（Channel-Height-Width, CHW）
        • FP16：表示输入数据将以半精度浮点数（FP16）格式传递
        • CHW：表示输入张量的维度顺序为 [channel, height, width]
    • --outputIOFormats=fp16:chw
      • 定义输出张量的格式为 FP16 和通道优先（CHW）
  • 性能相关参数
    • --batch=<size>
      • 设置推理的批量大小（仅在隐式批量模式下有效）
      • 示例：--batch=32
    • --avgRuns=<num>
      • 计算平均性能时使用的运行次数
      • 示例：--avgRuns=5
  • 构建选项相关参数
    • --fp16：启用 FP16（半精度浮点数）推理以提高性能

TensorRT

TensorRT不同版本的变化非常大，往往不能够兼容，所以转换和推理实现均在同一个TensorRT环境下实现。

1
2
# python -c "import tensorrt as trt; print(trt.__version__)"
7.2.2.1

Docker

Nvidia提供了配置好依赖环境的Docker镜像，可以在容器中进行模型转换和推理：

1
2
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:20.12-py3
docker run -it --runtime nvidia --gpus=all --shm-size=16g -v /etc/localtime:/etc/localtime -v $(pwd):/workdir --workdir=/workdir --name tensorrt-v8.x ultralytics/yolov5:v7.0

相关实现

相关阅读