|
import torch
|
|
import os
|
|
import argparse
|
|
import copy
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--source', type=str, required=True)
|
|
parser.add_argument('--prefix', type=str, required=True)
|
|
parser.add_argument('--target', type=str, default=None)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
def main():
|
|
|
|
args = parse_args()
|
|
|
|
if args.target is None:
|
|
args.target = '/'.join(args.source.split('/')[:-1])
|
|
|
|
ckpt = torch.load(args.source, map_location='cpu', weights_only=True)
|
|
|
|
experts = dict()
|
|
|
|
new_ckpt = copy.deepcopy(ckpt)
|
|
|
|
state_dict = new_ckpt['state_dict']
|
|
|
|
for key, value in state_dict.items():
|
|
if 'mlp.experts' in key:
|
|
experts[key] = value
|
|
|
|
keys = ckpt['state_dict'].keys()
|
|
|
|
weight_names = ['keypoint_head.deconv_layers.0.weight',
|
|
'keypoint_head.deconv_layers.1.weight',
|
|
'keypoint_head.deconv_layers.1.bias',
|
|
'keypoint_head.deconv_layers.1.running_mean',
|
|
'keypoint_head.deconv_layers.1.running_var',
|
|
'keypoint_head.deconv_layers.1.num_batches_tracked',
|
|
'keypoint_head.deconv_layers.3.weight',
|
|
'keypoint_head.deconv_layers.4.weight',
|
|
'keypoint_head.deconv_layers.4.bias',
|
|
'keypoint_head.deconv_layers.4.running_mean',
|
|
'keypoint_head.deconv_layers.4.running_var',
|
|
'keypoint_head.deconv_layers.4.num_batches_tracked',
|
|
'keypoint_head.final_layer.weight',
|
|
'keypoint_head.final_layer.bias']
|
|
|
|
target_expert = 0
|
|
new_ckpt = copy.deepcopy(ckpt)
|
|
|
|
for key in keys:
|
|
if 'mlp.fc2' in key:
|
|
value = new_ckpt['state_dict'][key]
|
|
value = torch.cat([value, experts[key.replace('fc2.', f'experts.{target_expert}.')]], dim=0)
|
|
new_ckpt['state_dict'][key] = value
|
|
|
|
|
|
for j in range(5):
|
|
|
|
for tensor_name in weight_names:
|
|
new_ckpt['state_dict'].pop(tensor_name.replace('keypoint_head',
|
|
f'associate_keypoint_heads.{j}'))
|
|
|
|
keys = new_ckpt['state_dict'].keys()
|
|
for key in list(keys):
|
|
if 'expert' in key:
|
|
new_ckpt['state_dict'].pop(key)
|
|
|
|
torch.save(new_ckpt, os.path.join(args.target, args.prefix + 'coco.pth'))
|
|
|
|
names = ['aic', 'mpii', 'ap10k', 'apt36k','wholebody']
|
|
num_keypoints = [14, 16, 17, 17, 133]
|
|
exist_range = True
|
|
|
|
for i in range(5):
|
|
|
|
new_ckpt = copy.deepcopy(ckpt)
|
|
|
|
target_expert = i + 1
|
|
|
|
for key in keys:
|
|
if 'mlp.fc2' in key:
|
|
expert_key = key.replace('fc2.', f'experts.{target_expert}.')
|
|
if expert_key in experts:
|
|
value = new_ckpt['state_dict'][key]
|
|
value = torch.cat([value, experts[expert_key]], dim=0)
|
|
else:
|
|
exist_range = False
|
|
|
|
new_ckpt['state_dict'][key] = value
|
|
|
|
if not exist_range:
|
|
break
|
|
|
|
for tensor_name in weight_names:
|
|
new_ckpt['state_dict'][tensor_name] = new_ckpt['state_dict'][tensor_name.replace('keypoint_head',
|
|
f'associate_keypoint_heads.{i}')]
|
|
|
|
for tensor_name in ['keypoint_head.final_layer.weight', 'keypoint_head.final_layer.bias']:
|
|
new_ckpt['state_dict'][tensor_name] = new_ckpt['state_dict'][tensor_name][:num_keypoints[i]]
|
|
|
|
|
|
for j in range(5):
|
|
|
|
for tensor_name in weight_names:
|
|
new_ckpt['state_dict'].pop(tensor_name.replace('keypoint_head',
|
|
f'associate_keypoint_heads.{j}'))
|
|
|
|
keys = new_ckpt['state_dict'].keys()
|
|
for key in list(keys):
|
|
if 'expert' in key:
|
|
new_ckpt['state_dict'].pop(key)
|
|
|
|
torch.save(new_ckpt, os.path.join(args.target, f'{args.prefix}{names[i]}.pth'))
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|
