-
Notifications
You must be signed in to change notification settings - Fork 22
/
model_split.py
119 lines (92 loc) · 4.28 KB
/
model_split.py
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
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')
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
# remove unnecessary part in the state dict
for j in range(5):
# remove associate part
for tensor_name in weight_names:
new_ckpt['state_dict'].pop(tensor_name.replace('keypoint_head',
f'associate_keypoint_heads.{j}'))
# remove expert part
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]]
# remove unnecessary part in the state dict
for j in range(5):
# remove associate part
for tensor_name in weight_names:
new_ckpt['state_dict'].pop(tensor_name.replace('keypoint_head',
f'associate_keypoint_heads.{j}'))
# remove expert part
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()