File size: 4,511 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
from collections import OrderedDict
import torch
convert_dict_p5 = {
'model.0': 'backbone.stem',
'model.1': 'backbone.stage1.0',
'model.2': 'backbone.stage1.1',
'model.3': 'backbone.stage2.0',
'model.4': 'backbone.stage2.1',
'model.5': 'backbone.stage3.0',
'model.6': 'backbone.stage3.1',
'model.7': 'backbone.stage4.0',
'model.8': 'backbone.stage4.1',
'model.9.cv1': 'backbone.stage4.2.conv1',
'model.9.cv2': 'backbone.stage4.2.conv2',
'model.10': 'neck.reduce_layers.2',
'model.13': 'neck.top_down_layers.0.0',
'model.14': 'neck.top_down_layers.0.1',
'model.17': 'neck.top_down_layers.1',
'model.18': 'neck.downsample_layers.0',
'model.20': 'neck.bottom_up_layers.0',
'model.21': 'neck.downsample_layers.1',
'model.23': 'neck.bottom_up_layers.1',
'model.24.m': 'bbox_head.head_module.convs_pred',
'model.24.proto': 'bbox_head.head_module.proto_preds',
}
convert_dict_p6 = {
'model.0': 'backbone.stem',
'model.1': 'backbone.stage1.0',
'model.2': 'backbone.stage1.1',
'model.3': 'backbone.stage2.0',
'model.4': 'backbone.stage2.1',
'model.5': 'backbone.stage3.0',
'model.6': 'backbone.stage3.1',
'model.7': 'backbone.stage4.0',
'model.8': 'backbone.stage4.1',
'model.9': 'backbone.stage5.0',
'model.10': 'backbone.stage5.1',
'model.11.cv1': 'backbone.stage5.2.conv1',
'model.11.cv2': 'backbone.stage5.2.conv2',
'model.12': 'neck.reduce_layers.3',
'model.15': 'neck.top_down_layers.0.0',
'model.16': 'neck.top_down_layers.0.1',
'model.19': 'neck.top_down_layers.1.0',
'model.20': 'neck.top_down_layers.1.1',
'model.23': 'neck.top_down_layers.2',
'model.24': 'neck.downsample_layers.0',
'model.26': 'neck.bottom_up_layers.0',
'model.27': 'neck.downsample_layers.1',
'model.29': 'neck.bottom_up_layers.1',
'model.30': 'neck.downsample_layers.2',
'model.32': 'neck.bottom_up_layers.2',
'model.33.m': 'bbox_head.head_module.convs_pred',
'model.33.proto': 'bbox_head.head_module.proto_preds',
}
def convert(src, dst):
"""Convert keys in pretrained YOLOv5 models to mmyolo style."""
if src.endswith('6.pt'):
convert_dict = convert_dict_p6
is_p6_model = True
print('Converting P6 model')
else:
convert_dict = convert_dict_p5
is_p6_model = False
print('Converting P5 model')
try:
yolov5_model = torch.load(src)['model']
blobs = yolov5_model.state_dict()
except ModuleNotFoundError:
raise RuntimeError(
'This script must be placed under the ultralytics/yolov5 repo,'
' because loading the official pretrained model need'
' `model.py` to build model.')
state_dict = OrderedDict()
for key, weight in blobs.items():
num, module = key.split('.')[1:3]
if (is_p6_model and
(num == '11' or num == '33')) or (not is_p6_model and
(num == '9' or num == '24')):
if module == 'anchors':
continue
prefix = f'model.{num}.{module}'
else:
prefix = f'model.{num}'
new_key = key.replace(prefix, convert_dict[prefix])
if '.m.' in new_key:
new_key = new_key.replace('.m.', '.blocks.')
new_key = new_key.replace('.cv', '.conv')
elif 'bbox_head.head_module.proto_preds.cv' in new_key:
new_key = new_key.replace(
'bbox_head.head_module.proto_preds.cv',
'bbox_head.head_module.proto_preds.conv')
else:
new_key = new_key.replace('.cv1', '.main_conv')
new_key = new_key.replace('.cv2', '.short_conv')
new_key = new_key.replace('.cv3', '.final_conv')
state_dict[new_key] = weight
print(f'Convert {key} to {new_key}')
# save checkpoint
checkpoint = dict()
checkpoint['state_dict'] = state_dict
torch.save(checkpoint, dst)
# Note: This script must be placed under the yolov5 repo to run.
def main():
parser = argparse.ArgumentParser(description='Convert model keys')
parser.add_argument(
'--src', default='yolov5s.pt', help='src yolov5 model path')
parser.add_argument('--dst', default='mmyolov5s.pt', help='save path')
args = parser.parse_args()
convert(args.src, args.dst)
if __name__ == '__main__':
main()
|