文章目录
mmpose PyTorch模型转TensorRT
1. github开源代码
yolov5 TensorRT推理的开源代码位置在https://github.com/linghu8812/tensorrt_inference/tree/master/mmpose,PyTorch转onnx的代码是原作者的代码:pytorch2onnx.py,原作者仓库见https://github.com/open-mmlab/mmpose。
2. PyTorch模型转ONNX模型
首先通过命令git clone git@github.com:open-mmlab/mmpose.git
clone mmpose的代码,然后按照install.md的步骤配置mmpose运行环境。完成环境配置后,按照tutorials/5_export_model.md导出ONNX文件,opset最好选择11以上。
python3 tools/pytorch2onnx.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [--shape ${SHAPE}] \
[--verify] [--show] [--output-file ${OUTPUT_FILE}] [--opset-version ${VERSION}]
样例:
python3 tools/pytorch2onnx.py configs/top_down/hrnet/coco/hrnet_w48_coco_256x192.py https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth --output-file hrnet_w48_coco_256x192.onnx
PyTorch转ONNX模型的步骤,原作者已完成了开发,可以一键导出。
3. ONNX模型转TensorRT模型
3.1 概述
TensorRT模型即TensorRT的推理引擎,代码中通过C++实现。相关配置写在config.yaml文件中,如果存在engine_file
的路径,则读取engine_file
,否则从onnx_file
生成engine_file
,生成engine的代码从model.cpp类中继承。
config.yaml文件可以设置batch size,图像的size及模型的anchor等。
mmpose:
onnx_file: "../hrnet_w48_coco_256x192.onnx"
engine_file: "../hrnet_w48_coco_256x192.trt"
BATCH_SIZE: 1
INPUT_CHANNEL: 3
IMAGE_WIDTH: 192
IMAGE_HEIGHT: 256
img_mean: [0.485, 0.456, 0.406]
img_std: [0.229, 0.224, 0.225]
num_key_points: 17
skeleton: [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11],
[6, 12], [5, 6], [5, 7], [6, 8], [7, 9], [8, 10], [1, 2],
[0, 1], [0, 2], [1, 3], [2, 4], [3, 5], [4, 6]]
point_thresh: 0.5
3.2 编译
通过以下命令对项目进行编译,生成mmpose_trt
mkdir build && cd build
cmake ..
make -j
3.3 运行
通过以下命令运行项目,得到推理结果
./mmpose_trt../config.yaml ../samples
4. 推理结果
推理结果如下图所示: