Optional batched pose processing.

2024-12-18 16:22:08 +01:00
parent 7b8d209601
commit 07426fac2f
8 changed files with 151 additions and 75 deletions
--- a/extras/mmdeploy/README.md
+++ b/extras/mmdeploy/README.md
@ -11,39 +11,53 @@ docker build --progress=plain -f extras/mmdeploy/dockerfile -t rpt_mmdeploy .
 ## ONNX

 ```bash
-export withFP16="_fp16"
-cp /RapidPoseTriangulation/extras/mmdeploy/configs/detection_onnxruntime_static-320x320$withFP16.py configs/mmdet/detection/
-
 cd /mmdeploy/
+export withFP16="_fp16"
+cp /RapidPoseTriangulation/extras/mmdeploy/configs/detection_onnxruntime_static-320x320"$withFP16".py configs/mmdet/detection/
+
 python3 ./tools/deploy.py \
-    configs/mmdet/detection/detection_onnxruntime_static-320x320$withFP16.py \
+    configs/mmdet/detection/detection_onnxruntime_static-320x320"$withFP16".py \
    /mmpose/projects/rtmpose/rtmdet/person/rtmdet_nano_320-8xb32_coco-person.py \
    https://download.openmmlab.com/mmpose/v1/projects/rtmpose/rtmdet_nano_8xb32-100e_coco-obj365-person-05d8511e.pth \
    /mmpose/projects/rtmpose/examples/onnxruntime/human-pose.jpeg \
    --work-dir work_dir \
    --show

-mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320$withFP16.onnx
+mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_1x3x320x320"$withFP16".onnx
 ```

 ```bash
-export withFP16="_fp16"
-cp /RapidPoseTriangulation/extras/mmdeploy/configs/pose-detection_simcc_onnxruntime_static-384x288$withFP16.py configs/mmpose/
-
 cd /mmdeploy/
+export withFP16="_fp16"
+cp /RapidPoseTriangulation/extras/mmdeploy/configs/pose-detection_simcc_onnxruntime_static-384x288"$withFP16".py configs/mmpose/
+cp /RapidPoseTriangulation/extras/mmdeploy/configs/pose-detection_simcc_onnxruntime_dynamic-384x288"$withFP16".py configs/mmpose/
+
 python3 ./tools/deploy.py \
-    configs/mmpose/pose-detection_simcc_onnxruntime_static-384x288$withFP16.py \
+    configs/mmpose/pose-detection_simcc_onnxruntime_static-384x288"$withFP16".py \
    /mmpose/projects/rtmpose/rtmpose/body_2d_keypoint/rtmpose-m_8xb256-420e_coco-384x288.py \
    https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/rtmpose-m_simcc-body7_pt-body7_420e-384x288-65e718c4_20230504.pth \
    /mmpose/projects/rtmpose/examples/onnxruntime/human-pose.jpeg \
    --work-dir work_dir \
    --show

-mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288$withFP16.onnx
+mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_1x3x384x288"$withFP16".onnx
+
+python3 ./tools/deploy.py \
+    configs/mmpose/pose-detection_simcc_onnxruntime_dynamic-384x288"$withFP16".py \
+    /mmpose/projects/rtmpose/rtmpose/body_2d_keypoint/rtmpose-m_8xb256-420e_coco-384x288.py \
+    https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/rtmpose-m_simcc-body7_pt-body7_420e-384x288-65e718c4_20230504.pth \
+    /mmpose/projects/rtmpose/examples/onnxruntime/human-pose.jpeg \
+    --work-dir work_dir \
+    --show
+
+mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_Bx3x384x288"$withFP16".onnx
 ```

 ```bash
 python3 /RapidPoseTriangulation/extras/mmdeploy/make_extra_graphs.py
+```
+
+```bash
 python3 /RapidPoseTriangulation/extras/mmdeploy/add_extra_steps.py
 ```

@ -57,14 +71,17 @@ Run this directly in the inference container (the TensorRT versions need to be t
 export withFP16="_fp16"

 trtexec --fp16 \
-  --onnx=/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320"$withFP16"_extra-steps.onnx \
+  --onnx=/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_1x320x320x3"$withFP16"_extra-steps.onnx \
  --saveEngine=end2end.engine

 mv ./end2end.engine /RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_1x320x320x3"$withFP16"_extra-steps.engine

 trtexec --fp16 \
-  --onnx=/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288"$withFP16"_extra-steps.onnx \
-  --saveEngine=end2end.engine
+  --onnx=/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_Bx384x288x3"$withFP16"_extra-steps.onnx \
+  --saveEngine=end2end.engine \
+  --minShapes=image_input:1x384x288x3 \
+  --optShapes=image_input:1x384x288x3 \
+  --maxShapes=image_input:1x384x288x3

 mv ./end2end.engine /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_1x384x288x3"$withFP16"_extra-steps.engine
 ```
@ -74,14 +91,14 @@ mv ./end2end.engine /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_1x
 ## Benchmark

 ```bash
+cd /mmdeploy/
 export withFP16="_fp16"

-cd /mmdeploy/
 python3 ./tools/profiler.py \
-    configs/mmpose/pose-detection_simcc_onnxruntime_static-384x288$withFP16.py \
+    configs/mmpose/pose-detection_simcc_onnxruntime_static-384x288"$withFP16".py \
    /mmpose/projects/rtmpose/rtmpose/body_2d_keypoint/rtmpose-m_8xb256-420e_coco-384x288.py \
    /RapidPoseTriangulation/extras/mmdeploy/testimages/ \
-    --model /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288$withFP16.onnx \
+    --model /RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_1x3x384x288"$withFP16".onnx \
    --shape 384x288 \
    --device cuda \
    --warmup 50 \
--- a/extras/mmdeploy/add_extra_steps.py
+++ b/extras/mmdeploy/add_extra_steps.py
@ -1,12 +1,15 @@
+import re
+
 import numpy as np
 import onnx
-from onnx import TensorProto, compose, helper, numpy_helper
+from onnx import TensorProto, helper, numpy_helper

 # ==================================================================================================

 base_path = "/RapidPoseTriangulation/extras/mmdeploy/exports/"
-pose_model_path = base_path + "rtmpose-m_384x288.onnx"
-det_model_path = base_path + "rtmdet-nano_320x320.onnx"
+det_model_path = base_path + "rtmdet-nano_1x3x320x320.onnx"
+pose_model_path1 = base_path + "rtmpose-m_Bx3x384x288.onnx"
+pose_model_path2 = base_path + "rtmpose-m_1x3x384x288.onnx"

 norm_mean = -1 * (np.array([0.485, 0.456, 0.406]) * 255)
 norm_std = 1.0 / (np.array([0.229, 0.224, 0.225]) * 255)
@ -97,6 +100,11 @@ def add_steps_to_onnx(model_path):
    for i, j in enumerate([0, 3, 1, 2]):
        input_shape[j].dim_value = dims[i]

+    # Set the batch size to a defined string
+    input_shape = graph.input[0].type.tensor_type.shape.dim
+    if input_shape[0].dim_value == 0:
+        input_shape[0].dim_param = "batch_size"
+
    # Rename the input tensor
    main_input_image_name = model.graph.input[0].name
    for node in model.graph.node:
@ -108,7 +116,8 @@ def add_steps_to_onnx(model_path):
    # Set input image type to int8
    model.graph.input[0].type.tensor_type.elem_type = TensorProto.UINT8

-    path = model_path.replace(".onnx", "_extra-steps.onnx")
+    path = re.sub(r"(x)(\d+)x(\d+)x(\d+)", r"\1\3x\4x\2", model_path)
+    path = path.replace(".onnx", "_extra-steps.onnx")
    onnx.save(model, path)


@ -116,10 +125,12 @@ def add_steps_to_onnx(model_path):


 def main():
-    add_steps_to_onnx(pose_model_path)
    add_steps_to_onnx(det_model_path)
+    add_steps_to_onnx(pose_model_path1)
+    add_steps_to_onnx(pose_model_path2)
    add_steps_to_onnx(det_model_path.replace(".onnx", "_fp16.onnx"))
-    add_steps_to_onnx(pose_model_path.replace(".onnx", "_fp16.onnx"))
+    add_steps_to_onnx(pose_model_path1.replace(".onnx", "_fp16.onnx"))
+    add_steps_to_onnx(pose_model_path2.replace(".onnx", "_fp16.onnx"))


 # ==================================================================================================
--- a/extras/mmdeploy/configs/detection_onnxruntime_static-320x320.py
+++ b/extras/mmdeploy/configs/detection_onnxruntime_static-320x320.py
@ -5,7 +5,7 @@ onnx_config = dict(
 )

 codebase_config = dict(
-    # For later TensorRT inference, the number of output boxes needs to be as stable as possible, 
+    # For later TensorRT inference, the number of output boxes needs to be as stable as possible,
    # because a drop in the box count leads to a re-optimization which takes a lot of time,
    # therefore reduce the maximum number of output boxes to the smallest usable value and sort out
    # low confidence boxes outside the model.
--- a/extras/mmdeploy/configs/detection_onnxruntime_static-320x320_fp16.py
+++ b/extras/mmdeploy/configs/detection_onnxruntime_static-320x320_fp16.py
@ -5,7 +5,7 @@ onnx_config = dict(
 )

 codebase_config = dict(
-    # For later TensorRT inference, the number of output boxes needs to be as stable as possible, 
+    # For later TensorRT inference, the number of output boxes needs to be as stable as possible,
    # because a drop in the box count leads to a re-optimization which takes a lot of time,
    # therefore reduce the maximum number of output boxes to the smallest usable value and sort out
    # low confidence boxes outside the model.
--- a/extras/mmdeploy/configs/pose-detection_simcc_onnxruntime_dynamic-384x288.py
+++ b/extras/mmdeploy/configs/pose-detection_simcc_onnxruntime_dynamic-384x288.py
@ -0,0 +1,19 @@
+_base_ = ["./pose-detection_static.py", "../_base_/backends/onnxruntime.py"]
+
+onnx_config = dict(
+    input_shape=[288, 384],
+    output_names=["kpts", "scores"],
+    dynamic_axes={
+        "input": {
+            0: "batch",
+        },
+        "kpts": {
+            0: "batch",
+        },
+        "scores": {
+            0: "batch",
+        },
+    },
+)
+
+codebase_config = dict(export_postprocess=True)  # export get_simcc_maximum
--- a/extras/mmdeploy/configs/pose-detection_simcc_onnxruntime_dynamic-384x288_fp16.py
+++ b/extras/mmdeploy/configs/pose-detection_simcc_onnxruntime_dynamic-384x288_fp16.py
@ -0,0 +1,19 @@
+_base_ = ["./pose-detection_static.py", "../_base_/backends/onnxruntime-fp16.py"]
+
+onnx_config = dict(
+    input_shape=[288, 384],
+    output_names=["kpts", "scores"],
+    dynamic_axes={
+        "input": {
+            0: "batch",
+        },
+        "kpts": {
+            0: "batch",
+        },
+        "scores": {
+            0: "batch",
+        },
+    },
+)
+
+codebase_config = dict(export_postprocess=True)  # export get_simcc_maximum