Répertoire d'articles
Convertir le modèle PyTorch mmpose en TensorRT
1. code open source github
L'emplacement du code open source de l'inférence Yolov5 TensorRT est à https://github.com/linghu8812/tensorrt_inference/tree/master/mmpose , le code de PyTorch à onnx est le code de l'auteur original: pytorch2onnx.py , voir https: // github .com / open-mmlab / mmpose .
2. Convertissez le modèle PyTorch en modèle ONNX
Commencez par git clone [email protected]:open-mmlab/mmpose.git
cloner le code de mmpose via la commande , puis configurez l'environnement d'exécution mmpose selon les étapes de install.md . Après avoir terminé la configuration de l'environnement, suivez les tutoriels / 5_export_model.md pour exporter le fichier ONNX, l'opset est préférable de choisir 11 ou plus.
python3 tools/pytorch2onnx.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [--shape ${SHAPE}] \
[--verify] [--show] [--output-file ${OUTPUT_FILE}] [--opset-version ${VERSION}]
Exemple:
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
L'auteur original a terminé le développement du modèle PyTorch vers ONNX et peut être exporté en un seul clic.
3. Convertissez le modèle ONNX en modèle TensorRT
3.1 Aperçu
Le modèle TensorRT est le moteur de raisonnement de TensorRT et le code est implémenté en C ++. La configuration appropriée est écrite dans le fichier config.yaml , si engine_file
le chemin existe , il est lu engine_file
, sinon il est onnx_file
généré engine_file
, et le code de génération du moteur est hérité de la classe model.cpp .
Le fichier config.yaml peut définir la taille du lot, la taille de l'image et l'ancre du modèle.
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 Compiler
Compilez le projet avec la commande suivante pour générermmpose_trt
mkdir build && cd build
cmake ..
make -j
3.3 Fonctionnement
Exécutez le projet avec la commande suivante pour obtenir le résultat de l'inférence
./mmpose_trt../config.yaml ../samples
4. Résultats d'inférence
Le résultat de l'inférence est illustré dans la figure ci-dessous: