[Feature] Add RewriteCheckPointHook to rewrite key of state_dict in checkpoint

[HAOCHENYE]

Thanks for your contribution and we appreciate it a lot. The following instructions would make your pull request more healthy and more easily get feedback. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers.

Motivation

Add MigrateCheckPointHook to the rewriting key in the loaded checkpoint.

The state_dict(or any other keys, such as ema_state_dict) in checkpoint may not match the model strictly. MigrateCheckPointHook can rewrite these keys to the matched ones.

Modification

Please briefly describe what modification is made in this PR.

BC-breaking (Optional)

Does the modification introduce changes that break the backward-compatibility of the downstream repos?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.

Use cases (Optional)

import torch
import torch.nn as nn

from mmengine.hooks import MigrateCheckPointHook


class SubModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layer1 = nn.Linear(1, 1)
        self.layer2 = nn.Linear(1, 1)
class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layer1 = nn.Linear(1, 1)
        self.layer2 = nn.Linear(1, 1)
        self.submodule = SubModule()

# original `state_dict`.
model = Model()
model.state_dict().keys()
# ['layer1.weight', 'layer1.bias', 'layer2.weight', 'layer2.bias',
#  'submodule.layer1.weight', 'submodule.layer1.bias',
#  'submodule.layer2.weight', 'submodule.layer2.bias']

# remove `layer1` in `state_dict`.
checkpoint = dict(state_dict=model.state_dict())
hook = MigrateCheckPointHook(removed_prefix='layer1')
hook.after_load_checkpoint(None, checkpoint)
checkpoint['state_dict'].keys()
# ['layer2.weight', 'layer2.bias', 'submodule.layer1.weight',
#  'submodule.layer1.bias', 'submodule.layer2.weight',
#  'submodule.layer2.bias']

# remove key with prefix `submodule`.
checkpoint = dict(state_dict=model.state_dict())
hook = MigrateCheckPointHook(removed_prefix='submodule')
hook.after_load_checkpoint(None, checkpoint)
checkpoint['state_dict'].keys()
# ['layer1.weight', 'layer1.bias', 'layer2.weight', 'layer2.bias']

# remapping prefix `module` to `submodule`.
checkpoint = dict(state_dict=model.state_dict())
hook = MigrateCheckPointHook(prefix_mapping=[dict(src='submodule', dst='module')])  # noqa: E501
hook.after_load_checkpoint(None, checkpoint)
checkpoint['state_dict'].keys()
# ['layer1.weight', 'layer1.bias', 'layer2.weight', 'layer2.bias',
#  'module.layer1.weight', 'module.layer1.bias',
#  'module.layer2.weight', 'module.layer2.bias']

# remapping prefix `module` to `submodule`, `layer1` to `linear1`.
checkpoint = dict(state_dict=model.state_dict())
hook = MigrateCheckPointHook(
    prefix_mapping=[dict(src='submodule', dst='module'),
                dict(src='layer1', dst='linear1')])
hook.after_load_checkpoint(None, checkpoint)
checkpoint['state_dict'].keys()
# ['linear1.weight', 'linear1.bias', 'layer2.weight',
#  'layer2.bias', 'module.layer1.weight', 'module.layer1.bias',
#  'module.layer2.weight', 'module.layer2.bias']

# merge other `state_dict`.
checkpoint = dict(state_dict=model.state_dict())
merged_ckpt = dict(state_dict=nn.Conv2d(1, 1, 1).state_dict())
torch.save(merged_ckpt, 'docs_demo.pth')
hook = MigrateCheckPointHook(
    merged_state_dicts=['docs_demo.pth'])
hook.after_load_checkpoint(None, checkpoint)
checkpoint['state_dict'].keys()
# ['layer1.weight', 'layer1.bias', 'layer2.weight', 'layer2.bias',
#  'submodule.layer1.weight', 'submodule.layer1.bias',
#  'submodule.layer2.weight', 'submodule.layer2.bias',
#  'weight', 'bias'

Checklist

1. Pre-commit or other linting tools are used to fix the potential lint issues.
2. The modification is covered by complete unit tests. If not, please add more unit test to ensure the correctness.
3. If the modification has potential influence on downstream projects, this PR should be tested with downstream projects, like MMDet or MMCls.
4. The documentation has been modified accordingly, like docstring or example tutorials.

[gaotongxiao]

Don't forget to rename this file as well

[gaotongxiao]

Suggested change

	removed and remapped keys. Removed keys has the next highest
	priority, once the original key has been removed, it cannot be mapped
	removed and remapped keys. Removed keys have the second highest
	priority. Once the original key has been removed, it cannot be mapped

[gaotongxiao]

Suggested change

	removed_prefix (List[str]): Key starts with corresponding prefix will
	unused_prefix (str or List[str]): Keys starting with corresponding prefix(es) will

removed_prefix is ambiguous, as it sounds like the keys starting with these prefixes have already been removed.

[gaotongxiao]

Suggested change

	merged_state_dicts (List[str]): A list of checkpoint paths need to be
	merge_from (List[str]): A list of checkpoint paths needed to be

Might be more concise

[gaotongxiao]

type hint: Union[List[str], str]

[gaotongxiao]

Suggested change

	dict), ('prefix_mapping should be a list of dict a dict, but got '
	dict), ('prefix_mapping should be a list of dict or a dict, but got '

[gaotongxiao]

Suggested change

	applied_key (str): Target state dictionary saved in checkpoints, which
	state_key (str): Target state dictionary saved in checkpoints, which

Not sure if it's a better name, but applied_key sounds weird. What has been applied to the key?