Pytorch->ONNX->TensorRT

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

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

Pytorch2ONNX

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

ONNX2TensorRT

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

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

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

相关实现

相关阅读