Hui Zhang
4 years ago
parent
fa6e333422
commit
156ccb947b
@ -0,0 +1,96 @@
|
|||||||
|
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
import argparse
|
||||||
|
import glob
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import paddle
|
||||||
|
|
||||||
|
|
||||||
|
def main(args):
|
||||||
|
checkpoints = []
|
||||||
|
val_scores = []
|
||||||
|
|
||||||
|
if args.val_best:
|
||||||
|
jsons = glob.glob(f'{args.ckpt_dir}/[!train]*.json')
|
||||||
|
for y in jsons:
|
||||||
|
dic_json = json.load(y)
|
||||||
|
loss = dic_json['valid_loss']
|
||||||
|
epoch = dic_json['epoch']
|
||||||
|
if epoch >= args.min_epoch and epoch <= args.max_epoch:
|
||||||
|
val_scores.append((epoch, loss))
|
||||||
|
|
||||||
|
val_scores = np.array(val_scores)
|
||||||
|
sort_idx = np.argsort(val_scores[:, 1])
|
||||||
|
sorted_val_scores = val_scores[sort_idx]
|
||||||
|
path_list = [
|
||||||
|
args.ckpt_dir + '/{}.pdparams'.format(int(epoch))
|
||||||
|
for epoch in sorted_val_scores[:args.num, 0]
|
||||||
|
]
|
||||||
|
print("best val scores = " + str(sorted_val_scores[:args.num, 1]))
|
||||||
|
print("selected epochs = " + str(sorted_val_scores[:args.num, 0].astype(
|
||||||
|
np.int64)))
|
||||||
|
else:
|
||||||
|
path_list = glob.glob(f'{args.ckpt_dir}/[!avg][!final]*.pdparams')
|
||||||
|
path_list = sorted(path_list, key=os.path.getmtime)
|
||||||
|
path_list = path_list[-args.num:]
|
||||||
|
|
||||||
|
print(path_list)
|
||||||
|
|
||||||
|
avg = None
|
||||||
|
num = args.num
|
||||||
|
assert num == len(path_list)
|
||||||
|
for path in path_list:
|
||||||
|
print(f'Processing {path}')
|
||||||
|
states = paddle.load(path)
|
||||||
|
if avg is None:
|
||||||
|
avg = states
|
||||||
|
else:
|
||||||
|
for k in avg.keys():
|
||||||
|
avg[k] += states[k]
|
||||||
|
# average
|
||||||
|
for k in avg.keys():
|
||||||
|
if avg[k] is not None:
|
||||||
|
avg[k] = paddle.divide(avg[k], num)
|
||||||
|
|
||||||
|
paddle.save(avg, args.dst_model)
|
||||||
|
print(f'Saving to {args.dst_model}')
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
parser = argparse.ArgumentParser(description='average model')
|
||||||
|
parser.add_argument('--dst_model', required=True, help='averaged model')
|
||||||
|
parser.add_argument(
|
||||||
|
'--ckpt_dir', required=True, help='ckpt model dir for average')
|
||||||
|
parser.add_argument(
|
||||||
|
'--val_best', action="store_true", help='averaged model')
|
||||||
|
parser.add_argument(
|
||||||
|
'--num', default=5, type=int, help='nums for averaged model')
|
||||||
|
parser.add_argument(
|
||||||
|
'--min_epoch',
|
||||||
|
default=0,
|
||||||
|
type=int,
|
||||||
|
help='min epoch used for averaging model')
|
||||||
|
parser.add_argument(
|
||||||
|
'--max_epoch',
|
||||||
|
default=65536, # Big enough
|
||||||
|
type=int,
|
||||||
|
help='max epoch used for averaging model')
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
print(args)
|
||||||
|
|
||||||
|
main(args)
|
||||||
Loading…
Reference in new issue