[Enhance] Enhance ArcFaceClsHead

[Ezra-Yu]

Motivation

Enhance ArcFaceClsHead.

related paper : subcenterArcFace.
related issue #1169

Modification

0. refactor the Arcface Code
1. support subcenterArcFace
2. support AdvMargins
3. fix arcface when using amp

Use cases (Optional)

use vanilla ArcFace

mode = dict(
    backbone = xxx,
    neck = xxxx,
    head=dict(
        type='ArcFaceClsHead',
        num_classes=5000,
        in_channels=1024,
        loss = dict(type='CrossEntropyLoss', loss_weight=1.0),
        init_cfg=None),
)

use SubCenterArcFace with 3 sub-centers

mode = dict(
    backbone = xxx,
    neck = xxxx,
    head=dict(
        type='ArcFaceClsHead',
        num_classes=5000,
        in_channels=1024,
        num_subcenters=3,
        loss = dict(type='CrossEntropyLoss', loss_weight=1.0),
        init_cfg=None),
)

use ArcFaceWithAdvMargin

mode = dict(
    backbone = xxx,
    neck = xxxx,
    head=dict(
        type='ArcFaceClsHead',
        num_classes=5000,
        in_channels=1024,
        num_subcenters=3,
        loss = dict(type='CrossEntropyLoss', loss_weight=1.0),
        init_cfg=None),
)

custom_hooks = [dict(type='SetCountPowAdvMarginsHook')]

Checklist

Before PR:

• Pre-commit or other linting tools are used to fix the potential lint issues.
• Bug fixes are fully covered by unit tests, the case that causes the bug should be added in the unit tests.
• The modification is covered by complete unit tests. If not, please add more unit test to ensure the correctness.
• The documentation has been modified accordingly, like docstring or example tutorials.

After PR:

• If the modification has potential influence on downstream or other related projects, this PR should be tested with those projects, like MMDet or MMSeg.
• CLA has been signed and all committers have signed the CLA in this PR.

[codecov]

Codecov Report

Base: 0.02% // Head: 89.05% // Increases project coverage by +89.03% 🎉

Coverage data is based on head (dd0979a) compared to base (b8b31e9).
Patch has no changes to coverable lines.

Additional details and impacted files

@@             Coverage Diff              @@
##           dev-1.x    #1181       +/-   ##
============================================
+ Coverage     0.02%   89.05%   +89.03%     
============================================
  Files          121      147       +26     
  Lines         8217    11205     +2988     
  Branches      1368     1777      +409     
============================================
+ Hits             2     9979     +9977     
+ Misses        8215      975     -7240     
- Partials         0      251      +251     

Flag	Coverage Δ
unittests	89.05% <ø> (+89.03%)	⬆️

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

Impacted Files	Coverage Δ
mmcls/apis/inference.py	0.00% <0.00%> (ø)
mmcls/datasets/transforms/compose.py
mmcls/models/backbones/davit.py	84.94% <0.00%> (ø)
mmcls/models/classifiers/timm.py	25.67% <0.00%> (ø)
mmcls/models/heads/margin_head.py	90.00% <0.00%> (ø)
mmcls/models/utils/layer_scale.py	86.66% <0.00%> (ø)
mmcls/models/retrievers/image2image.py	92.38% <0.00%> (ø)
mmcls/models/retrievers/__init__.py	100.00% <0.00%> (ø)
mmcls/models/backbones/swin_transformer_v2.py	89.63% <0.00%> (ø)
mmcls/models/backbones/mobileone.py	94.47% <0.00%> (ø)
... and 137 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[okotaku]

@Ezra-Yu The implementation is exactly as I intended in my proposal👍
It is a little complicated, but I can't think of any other simple plan.

[okotaku]

LGTM