138 lines
4.8 KiB
Python
138 lines
4.8 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
# This script consists of several convert functions which
|
|
# can modify the weights of model in original repo to be
|
|
# pre-trained weights.
|
|
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
|
|
|
|
def pvt_convert(ckpt):
|
|
new_ckpt = OrderedDict()
|
|
# Process the concat between q linear weights and kv linear weights
|
|
use_abs_pos_embed = False
|
|
use_conv_ffn = False
|
|
for k in ckpt.keys():
|
|
if k.startswith('pos_embed'):
|
|
use_abs_pos_embed = True
|
|
if k.find('dwconv') >= 0:
|
|
use_conv_ffn = True
|
|
for k, v in ckpt.items():
|
|
if k.startswith('head'):
|
|
continue
|
|
if k.startswith('norm.'):
|
|
continue
|
|
if k.startswith('cls_token'):
|
|
continue
|
|
if k.startswith('pos_embed'):
|
|
stage_i = int(k.replace('pos_embed', ''))
|
|
new_k = k.replace(f'pos_embed{stage_i}',
|
|
f'layers.{stage_i - 1}.1.0.pos_embed')
|
|
if stage_i == 4 and v.size(1) == 50: # 1 (cls token) + 7 * 7
|
|
new_v = v[:, 1:, :] # remove cls token
|
|
else:
|
|
new_v = v
|
|
elif k.startswith('patch_embed'):
|
|
stage_i = int(k.split('.')[0].replace('patch_embed', ''))
|
|
new_k = k.replace(f'patch_embed{stage_i}',
|
|
f'layers.{stage_i - 1}.0')
|
|
new_v = v
|
|
if 'proj.' in new_k:
|
|
new_k = new_k.replace('proj.', 'projection.')
|
|
elif k.startswith('block'):
|
|
stage_i = int(k.split('.')[0].replace('block', ''))
|
|
layer_i = int(k.split('.')[1])
|
|
new_layer_i = layer_i + use_abs_pos_embed
|
|
new_k = k.replace(f'block{stage_i}.{layer_i}',
|
|
f'layers.{stage_i - 1}.1.{new_layer_i}')
|
|
new_v = v
|
|
if 'attn.q.' in new_k:
|
|
sub_item_k = k.replace('q.', 'kv.')
|
|
new_k = new_k.replace('q.', 'attn.in_proj_')
|
|
new_v = torch.cat([v, ckpt[sub_item_k]], dim=0)
|
|
elif 'attn.kv.' in new_k:
|
|
continue
|
|
elif 'attn.proj.' in new_k:
|
|
new_k = new_k.replace('proj.', 'attn.out_proj.')
|
|
elif 'attn.sr.' in new_k:
|
|
new_k = new_k.replace('sr.', 'sr.')
|
|
elif 'mlp.' in new_k:
|
|
string = f'{new_k}-'
|
|
new_k = new_k.replace('mlp.', 'ffn.layers.')
|
|
if 'fc1.weight' in new_k or 'fc2.weight' in new_k:
|
|
new_v = v.reshape((*v.shape, 1, 1))
|
|
new_k = new_k.replace('fc1.', '0.')
|
|
new_k = new_k.replace('dwconv.dwconv.', '1.')
|
|
if use_conv_ffn:
|
|
new_k = new_k.replace('fc2.', '4.')
|
|
else:
|
|
new_k = new_k.replace('fc2.', '3.')
|
|
string += f'{new_k} {v.shape}-{new_v.shape}'
|
|
elif k.startswith('norm'):
|
|
stage_i = int(k[4])
|
|
new_k = k.replace(f'norm{stage_i}', f'layers.{stage_i - 1}.2')
|
|
new_v = v
|
|
else:
|
|
new_k = k
|
|
new_v = v
|
|
new_ckpt[new_k] = new_v
|
|
|
|
return new_ckpt
|
|
|
|
|
|
def swin_converter(ckpt):
|
|
|
|
new_ckpt = OrderedDict()
|
|
|
|
def correct_unfold_reduction_order(x):
|
|
out_channel, in_channel = x.shape
|
|
x = x.reshape(out_channel, 4, in_channel // 4)
|
|
x = x[:, [0, 2, 1, 3], :].transpose(1,
|
|
2).reshape(out_channel, in_channel)
|
|
return x
|
|
|
|
def correct_unfold_norm_order(x):
|
|
in_channel = x.shape[0]
|
|
x = x.reshape(4, in_channel // 4)
|
|
x = x[[0, 2, 1, 3], :].transpose(0, 1).reshape(in_channel)
|
|
return x
|
|
|
|
for k, v in ckpt.items():
|
|
if k.startswith('head'):
|
|
continue
|
|
elif k.startswith('layers'):
|
|
new_v = v
|
|
if 'attn.' in k:
|
|
new_k = k.replace('attn.', 'attn.w_msa.')
|
|
elif 'mlp.' in k:
|
|
if 'mlp.fc1.' in k:
|
|
new_k = k.replace('mlp.fc1.', 'ffn.layers.0.0.')
|
|
elif 'mlp.fc2.' in k:
|
|
new_k = k.replace('mlp.fc2.', 'ffn.layers.1.')
|
|
else:
|
|
new_k = k.replace('mlp.', 'ffn.')
|
|
elif 'downsample' in k:
|
|
new_k = k
|
|
if 'reduction.' in k:
|
|
new_v = correct_unfold_reduction_order(v)
|
|
elif 'norm.' in k:
|
|
new_v = correct_unfold_norm_order(v)
|
|
else:
|
|
new_k = k
|
|
new_k = new_k.replace('layers', 'stages', 1)
|
|
elif k.startswith('patch_embed'):
|
|
new_v = v
|
|
if 'proj' in k:
|
|
new_k = k.replace('proj', 'projection')
|
|
else:
|
|
new_k = k
|
|
else:
|
|
new_v = v
|
|
new_k = k
|
|
|
|
new_ckpt['backbone.' + new_k] = new_v
|
|
|
|
return new_ckpt
|