OpenMMLab Detection Toolbox and Benchmark
https://mmdetection.readthedocs.io/
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
137 lines
4.8 KiB
137 lines
4.8 KiB
# 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
|
|
|