-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain_chg_dim_macher_layer.py
36 lines (24 loc) · 1.01 KB
/
main_chg_dim_macher_layer.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
import argparse
import torch
def main():
args = get_args()
model = torch.load(args.inp_model)
model_params = model['params'] if 'params' in model.keys() else model
# find 'dim_matcher_layer' layers.
matcher_layers = [p for p in model_params.keys() if 'dim_matcher_layer' in p]
# pair weights and biases
matcher_layers = list(zip(matcher_layers[::2], matcher_layers[1::2]))
for w, b in matcher_layers:
linear = torch.nn.Linear(args.modulation_dim, args.nework_dim)
model_params[w] = torch.zeros_like(linear.weight.data)
model_params[b] = torch.zeros_like(linear.bias.data)
torch.save(model, args.out_model)
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('inp_model', type=str)
parser.add_argument('out_model', type=str)
parser.add_argument('-m', '--modulation_dim', type=int, default=1024)
parser.add_argument('-n', '--nework_dim', type=int, default=180)
return parser.parse_args()
if __name__ == '__main__':
main()