87 lines
2.4 KiB
Python
87 lines
2.4 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import argparse
|
|
import os.path as osp
|
|
from collections import OrderedDict
|
|
from typing import Union
|
|
|
|
import mmengine
|
|
import torch
|
|
from mmengine.runner.checkpoint import _load_checkpoint
|
|
|
|
|
|
def convert_mmcls_to_timm(state_dict: Union[OrderedDict, dict]) -> OrderedDict:
|
|
"""Convert keys in MMClassification pretrained vit models to timm tyle.
|
|
|
|
Args:
|
|
state_dict (Union[OrderedDict, dict]): The state dict of
|
|
MMClassification pretrained vit models.
|
|
|
|
Returns:
|
|
OrderedDict: The converted state dict.
|
|
"""
|
|
# only keep the backbone weights and remove the backbone. prefix
|
|
state_dict = {
|
|
key.replace('backbone.', ''): value
|
|
for key, value in state_dict.items() if key.startswith('backbone.')
|
|
}
|
|
|
|
# replace projection with proj
|
|
state_dict = {
|
|
key.replace('projection', 'proj'): value
|
|
for key, value in state_dict.items()
|
|
}
|
|
|
|
# replace ffn.layers.0.0 with mlp.fc1
|
|
state_dict = {
|
|
key.replace('ffn.layers.0.0', 'mlp.fc1'): value
|
|
for key, value in state_dict.items()
|
|
}
|
|
|
|
# replace ffn.layers.1 with mlp.fc2
|
|
state_dict = {
|
|
key.replace('ffn.layers.1', 'mlp.fc2'): value
|
|
for key, value in state_dict.items()
|
|
}
|
|
|
|
# replace layers with blocks
|
|
state_dict = {
|
|
key.replace('layers', 'blocks'): value
|
|
for key, value in state_dict.items()
|
|
}
|
|
|
|
# replace ln with norm
|
|
state_dict = {
|
|
key.replace('ln', 'norm'): value
|
|
for key, value in state_dict.items()
|
|
}
|
|
|
|
# replace the last norm1 with norm
|
|
state_dict['norm.weight'] = state_dict.pop('norm1.weight')
|
|
state_dict['norm.bias'] = state_dict.pop('norm1.bias')
|
|
|
|
state_dict = OrderedDict({'model': state_dict})
|
|
return state_dict
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description='Convert keys in MMClassification '
|
|
'pretrained vit models to timm tyle')
|
|
parser.add_argument('src', help='src model path or url')
|
|
parser.add_argument('dst', help='save path')
|
|
args = parser.parse_args()
|
|
|
|
checkpoint = _load_checkpoint(args.src)
|
|
if 'state_dict' in checkpoint:
|
|
state_dict = checkpoint['state_dict']
|
|
else:
|
|
state_dict = checkpoint
|
|
|
|
state_dict = convert_mmcls_to_timm(state_dict)
|
|
mmengine.mkdir_or_exist(osp.dirname(args.dst))
|
|
torch.save(state_dict, args.dst)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|