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
| import os
| import sys
| import torch
| import argparse
|
| if __name__ == '__main__':
| parser = argparse.ArgumentParser()
| parser.add_argument(
| "--model_dir",
| required=True,
| default=None,
| type=str,
| help="Director contains saved models."
| )
| parser.add_argument(
| "--average_epochs",
| nargs="+",
| type=int,
| default=[],
| )
| parser.add_argument(
| "--metric_name",
| type=str,
| default="der",
| help="The metric name of best models, only used for name."
| )
| args = parser.parse_args()
|
| root_path = args.model_dir
| idx_list = args.average_epochs
| n_models = len(idx_list)
| metric = args.metric_name
|
| if n_models > 0:
| avg = None
| for idx in idx_list:
| model_file = os.path.join(root_path, "{}epoch.pth".format(str(idx)))
| states = torch.load(model_file, map_location="cpu")
| if avg is None:
| avg = states
| else:
| for k in avg:
| avg[k] = avg[k] + states[k]
|
| for k in avg:
| if str(avg[k].dtype).startswith("torch.int"):
| pass
| else:
| avg[k] = avg[k] / n_models
|
| output_file = os.path.join(root_path, "valid.{}.ave_{}best.pth".format(metric, n_models))
| torch.save(avg, output_file)
| else:
| print("Number of models to average is 0, skip.")
|
|
