From 7947c86b576a35f75e192dd55f82dd72f47198a2 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 15 Dec 2020 21:50:28 -0800 Subject: [PATCH] Update COCO train postprocessing (#1702) --- train.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/train.py b/train.py index b88e44c..87879ed 100644 --- a/train.py +++ b/train.py @@ -402,15 +402,18 @@ def train(hyp, opt, device, tb_writer=None, wandb=None): # Test best.pt if opt.data.endswith('coco.yaml') and nc == 80: # if COCO - results, _, _ = test.test(opt.data, - batch_size=total_batch_size, - imgsz=imgsz_test, - model=attempt_load(best if best.exists() else last, device).half(), - single_cls=opt.single_cls, - dataloader=testloader, - save_dir=save_dir, - save_json=True, # use pycocotools - plots=False) + for conf, iou, save_json in ([0.25, 0.45, False], [0.001, 0.65, True]): # speed, mAP tests + results, _, _ = test.test(opt.data, + batch_size=total_batch_size, + imgsz=imgsz_test, + conf_thres=conf, + iou_thres=iou, + model=attempt_load(best if best.exists() else last, device).half(), + single_cls=opt.single_cls, + dataloader=testloader, + save_dir=save_dir, + save_json=save_json, + plots=False) else: dist.destroy_process_group()