File size: 2,142 Bytes
186701e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
from collections import OrderedDict
import torch
def convert(src, dst):
"""Convert keys in pretrained RTMDet models to MMYOLO style."""
blobs = torch.load(src)['state_dict']
state_dict = OrderedDict()
for key, weight in blobs.items():
if 'neck.reduce_layers.0' in key:
new_key = key.replace('.0', '.2')
state_dict[new_key] = weight
elif 'neck.reduce_layers.1' in key:
new_key = key.replace('reduce_layers.1', 'top_down_layers.0.1')
state_dict[new_key] = weight
elif 'neck.top_down_blocks.0' in key:
new_key = key.replace('down_blocks', 'down_layers.0')
state_dict[new_key] = weight
elif 'neck.top_down_blocks.1' in key:
new_key = key.replace('down_blocks', 'down_layers')
state_dict[new_key] = weight
elif 'downsamples' in key:
new_key = key.replace('downsamples', 'downsample_layers')
state_dict[new_key] = weight
elif 'bottom_up_blocks' in key:
new_key = key.replace('bottom_up_blocks', 'bottom_up_layers')
state_dict[new_key] = weight
elif 'out_convs' in key:
new_key = key.replace('out_convs', 'out_layers')
state_dict[new_key] = weight
elif 'bbox_head' in key:
new_key = key.replace('bbox_head', 'bbox_head.head_module')
state_dict[new_key] = weight
elif 'data_preprocessor' in key:
continue
else:
new_key = key
state_dict[new_key] = weight
print(f'Convert {key} to {new_key}')
# save checkpoint
checkpoint = dict()
checkpoint['state_dict'] = state_dict
checkpoint['meta'] = blobs.get('meta')
torch.save(checkpoint, dst)
def main():
parser = argparse.ArgumentParser(description='Convert model keys')
parser.add_argument('src', help='src rtm model path')
parser.add_argument('dst', help='save path')
args = parser.parse_args()
convert(args.src, args.dst)
if __name__ == '__main__':
main()
|