ViT Pose Pipeline Examlpe

[dbogunowicz]

Stage 0 exploration for the ViT Pose model

Source: https://github.com/ViTAE-Transformer/ViTPose

Installation

Follow the instructions in the readme file. Note:

• installing one of the dependencies, mmcv takes a lot of time and may look often like it is stuck. Be patient, it will eventually terminate successfully.
• after the setup completes, it is also advisable to downgrade the default torch version from 1.3 to 1.2 to avoid CUDA errors (at least on my server, which does not support Cuda setup for 1.3)

Export

Exporting the sample onnx model is quite easy. Before running the onnx install, one needs to manually install timm, onnx and onnxruntime. Then, launch the export script:

python tools/deployment/pytorch2onnx.py /home/ubuntu/damian/ViTPose/configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/coco/ViTPose_base_coco_256x192.py /home/ubuntu/damian/ViTPose/vitpose-b.pth 

The first argument is a config file (for ViTpose B) the second argument is the .pth checkpoint (weights). Both can be found on the main site of the repository:

The resulting model is about 400mb.

Benchmarking in DeepSparse:

Naive benchmarking shows that for the dense model, the engine is roughly x2 faster then ORT:

Postprocessing

ViT-Pose might be our first candidate for a "composed" deepsparse pipeline. It is a top-down pose estimation approach.
We take an image and run object detection (e.g. look for all the humans in the image)
We pass the cropped bounding boxes to ViT to get an array (batch, no_keypoints, h, w) array. To decode this array, according to the original paper, we need some simple composition of transposed convolutions.

What I do naively in this PR: I "squash" the array to (h,w) and then overlay it on the original image. We can see that the heatmap roughly coincides with the joints of the model.

[bfineran]

lgtm through engine forward

• any idea on complexity of postprocessing?
• for stage zero pipelines let's get them added in first into an examples directory potentially

[corey-nm]

nice! clean integration

[corey-nm]

Perhaps include some of the contents of the PR description as a README in the example folder?

[bfineran]

LGTM pending comment on removing changes to task.py

+1 to a simple README

[corey-nm]

🔥