112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import argparse
|
|
import re
|
|
from collections import OrderedDict, namedtuple
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
prog_description = """\
|
|
Convert OFA official models to MMPretrain format.
|
|
"""
|
|
|
|
MapItem = namedtuple(
|
|
'MapItem', 'pattern repl key_action value_action', defaults=[None] * 4)
|
|
|
|
|
|
def convert_by_mapdict(src_dict: dict, map_dict: Path):
|
|
dst_dict = OrderedDict()
|
|
convert_map_dict = dict()
|
|
|
|
for k, v in src_dict.items():
|
|
ori_k = k
|
|
for item in map_dict:
|
|
pattern = item.pattern
|
|
assert pattern is not None
|
|
match = next(re.finditer(pattern, k), None)
|
|
if match is None:
|
|
continue
|
|
match_group = match.groups()
|
|
repl = item.repl
|
|
|
|
key_action = item.key_action
|
|
if key_action is not None:
|
|
assert callable(key_action)
|
|
match_group = key_action(*match_group)
|
|
if isinstance(match_group, str):
|
|
match_group = (match_group, )
|
|
start, end = match.span(0)
|
|
if repl is not None:
|
|
k = k[:start] + repl.format(*match_group) + k[end:]
|
|
else:
|
|
for i, sub in enumerate(match_group):
|
|
start, end = match.span(i + 1)
|
|
k = k[:start] + str(sub) + k[end:]
|
|
|
|
value_action = item.value_action
|
|
if value_action is not None:
|
|
assert callable(value_action)
|
|
v = value_action(v)
|
|
|
|
if v is not None:
|
|
dst_dict[k] = v
|
|
convert_map_dict[k] = ori_k
|
|
return dst_dict, convert_map_dict
|
|
|
|
|
|
map_dict = [
|
|
# Encoder modules
|
|
MapItem(r'\.type_embedding\.', '.embed_type.'),
|
|
MapItem(r'\.layernorm_embedding\.', '.embedding_ln.'),
|
|
MapItem(r'\.patch_layernorm_embedding\.', '.image_embedding_ln.'),
|
|
MapItem(r'encoder.layer_norm\.', 'encoder.final_ln.'),
|
|
# Encoder layers
|
|
MapItem(r'\.attn_ln\.', '.attn_mid_ln.'),
|
|
MapItem(r'\.ffn_layernorm\.', '.ffn_mid_ln.'),
|
|
MapItem(r'\.final_layer_norm', '.ffn_ln'),
|
|
MapItem(r'encoder.*(\.self_attn\.)', key_action=lambda _: '.attn.'),
|
|
MapItem(
|
|
r'encoder.*(\.self_attn_layer_norm\.)',
|
|
key_action=lambda _: '.attn_ln.'),
|
|
# Decoder modules
|
|
MapItem(r'\.code_layernorm_embedding\.', '.code_embedding_ln.'),
|
|
MapItem(r'decoder.layer_norm\.', 'decoder.final_ln.'),
|
|
# Decoder layers
|
|
MapItem(r'\.self_attn_ln', '.self_attn_mid_ln'),
|
|
MapItem(r'\.cross_attn_ln', '.cross_attn_mid_ln'),
|
|
MapItem(r'\.encoder_attn_layer_norm', '.cross_attn_ln'),
|
|
MapItem(r'\.encoder_attn', '.cross_attn'),
|
|
MapItem(
|
|
r'decoder.*(\.self_attn_layer_norm\.)',
|
|
key_action=lambda _: '.self_attn_ln.'),
|
|
# Remove version key
|
|
MapItem(r'version', '', value_action=lambda _: None),
|
|
# Add model prefix
|
|
MapItem(r'^', 'model.'),
|
|
]
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description=prog_description)
|
|
parser.add_argument('src', type=str, help='The official checkpoint path.')
|
|
parser.add_argument('dst', type=str, help='The save path.')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
src = torch.load(args.src)
|
|
if 'extra_state' in src and 'ema' in src['extra_state']:
|
|
print('Use EMA weights.')
|
|
src = src['extra_state']['ema']
|
|
else:
|
|
src = src['model']
|
|
dst, _ = convert_by_mapdict(src, map_dict)
|
|
torch.save(dst, args.dst)
|
|
print('Done!!')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|