fix: use room sr to choose the best model

This commit is contained in:
Ting-Jun Wang 2024-07-16 13:51:00 +08:00
parent 59dfbf9c30
commit 8b5a7438a0
Signed by: snsd0805
GPG Key ID: 48D331A3D6160354

View File

@ -138,7 +138,7 @@ def train(args, train_env, val_envs, aug_env=None, rank=-1):
'\nListener training starts, start iteration: %s' % str(start_iter), record_file
)
best_val = {'val_unseen': {"spl": 0., "sr": 0., "state":"", "sspl": 0., 'found_sr': 0.}}
best_val = {'val_unseen': {"spl": 0., "sr": 0., "room_sr": 0., "state":"", "sspl": 0., 'found_sr': 0.}}
for idx in range(start_iter, start_iter+args.iters, args.log_every):
listner.logs = defaultdict(list)
@ -203,11 +203,12 @@ def train(args, train_env, val_envs, aug_env=None, rank=-1):
# select model by spl
if env_name in best_val:
if score_summary['sspl'] >= best_val[env_name]['sspl']:
if score_summary['room_sr'] >= best_val[env_name]['room_sr']:
best_val[env_name]['spl'] = score_summary['spl']
best_val[env_name]['sspl'] = score_summary['sspl']
best_val[env_name]['sr'] = score_summary['sr']
best_val[env_name]['found_sr'] = score_summary['found_sr']
best_val[env_name]['room_sr'] = score_summary['room_sr']
best_val[env_name]['state'] = 'Iter %d %s' % (iter, loss_str)
listner.save(idx, os.path.join(args.ckpt_dir, "best_%s" % (env_name)))