标签：模型 TensorRT hrnet PyTorch file mmpose
文章目录
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. 推理结果
推理结果如下图所示：
标签：模型,TensorRT,hrnet,PyTorch,file,mmpose 来源： https://blog.csdn.net/linghu8812/article/details/114411424
本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享； 2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关； 3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关； 4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除； 5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。