feat(keypoint-detection): enable ViTPose config/build/perf#905
Open
jeon185 wants to merge 1 commit into
Open
Conversation
ViTPose keypoint-detection models could not pass the wmk pipeline:
1. Task resolution: Optimum registers the ViTPose ONNX export config but has
no task-to-class entry for keypoint-detection, and transformers'
AutoModelForKeypointDetection only recognizes SuperPoint. Add
MODEL_CLASS_MAPPING[(vitpose, keypoint-detection)] = VitPoseForPoseEstimation
(models/hf/vitpose.py) so the resolver loads the correct class.
2. MoE export: the vitpose-plus checkpoints use a Mixture-of-Experts backbone
whose patcher injects a constant dataset_index at export time. Optimum's
patch_model_for_export defaults model_kwargs to None, so the patcher crashed
on init. Pass an explicit model_kwargs={} in _get_optimum_patcher. Also wrap
the Step 3 hierarchy trace in the same patcher context (it previously ran
the model forward without the injected dataset_index, failing before export).
Verified config -> build -> perf on all 6 acceptance models in #284
(base-simple, plus-{small,base,large,huge}, synthpose-vitpose-huge-hf).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Enables the 6 ViTPose keypoint-detection models from #284 to pass
wmk config -> build -> perf(CPU and OpenVINO). Eval isn't included here - I'll do the accuracy side in a follow-up since it needs a couple of design decisions first.Two things were blocking all 6 models:
wmk configfailed with "Task 'keypoint-detection' not supported by TasksManager". Optimum has the ViTPose ONNX export config but no task->class entry for keypoint-detection, andAutoModelForKeypointDetectiononly covers SuperPoint. Added the(vitpose, keypoint-detection) -> VitPoseForPoseEstimationmapping, same way we already do it for CLIP/SAM.The plus checkpoints (MoE backbone) crashed during export with "dataset_index must be provided when using multiple experts". Optimum's VitPoseModelPatcher injects a constant dataset_index, but
patch_model_for_exportdefaults model_kwargs to None so it crashed on init. Passing an explicitmodel_kwargs={}fixes that. The trace step (Step 3) was also running the model outside the patcher context, so I wrapped it the same way the export step already is.The exporter change isn't ViTPose-specific - it helps any MoE model whose patcher injects forward args.
Verified config/build/perf on all 6: vitpose-base-simple, vitpose-plus-small/base/large/huge, and synthpose-vitpose-huge-hf. Added unit tests for the mapping and the patcher model_kwargs handling.
One note: you still need to pass
--task keypoint-detectionexplicitly for now - the task isn't auto-detected from the config yet. I left auto-detection out of this PR to keep it small; can add it here or as a follow-up if you'd prefer.Refs #284.