[Enhancement] Accelerate Associative Embedding inference

[QwQ2000]

Motivation

Improve the inference speed of bottom-up human pose estimators.

Modification

Re-implement AssociativeEmbedding.refine with torch instead of numpy. Accelerate tensor operations and avoid moving large tensors from GPU to CPU. AssociativeEmbedding.parse is also adjusted to fit the changes in AssociativeEmbedding.refine.

BC-breaking (Optional)

Use cases (Optional)

Checklist

Before PR:

• I have read and followed the workflow indicated in the CONTRIBUTING.md to create this PR.
• Pre-commit or linting tools indicated in CONTRIBUTING.md are used to fix the potential lint issues.
• Bug fixes are covered by unit tests, the case that causes the bug should be added in the unit tests.
• New functionalities are covered by complete unit tests. If not, please add more unit tests to ensure correctness.
• The documentation has been modified accordingly, including docstring or example tutorials.

After PR:

• CLA has been signed and all committers have signed the CLA in this PR.

[codecov]

Codecov Report

Merging #1099 (83083d6) into master (bf691c0) will decrease coverage by 0.00%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #1099      +/-   ##
==========================================
- Coverage   83.39%   83.38%   -0.01%     
==========================================
  Files         196      196              
  Lines       15142    15144       +2     
  Branches     2736     2736              
==========================================
+ Hits        12627    12628       +1     
  Misses       1823     1823              
- Partials      692      693       +1     

Flag	Coverage Δ
unittests	83.32% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
mmpose/core/post_processing/group.py	82.84% <100.00%> (+0.20%)	⬆️
mmpose/datasets/pipelines/shared_transform.py	85.36% <0.00%> (-0.49%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update bf691c0...83083d6. Read the comment docs.

[jin-s13]

Have you tested the model inference? Could it produce the same accuracy as before?

[QwQ2000]

I tested the accuracy on two different backbones. The results can be found in https://openmmlab.feishu.cn/docs/doccnZWiwgGNh3fuxcTQv49f8qg.

[jin-s13]

lint fails. Please re-run pre-commit.

[QwQ2000]

Change this pull request to branch dev-0.22