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.gitcloner 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_filele chemin existe , il est lu engine_file, sinon il est onnx_filegé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:
Insérez la description de l'image ici