Update COCO train postprocessing (#1702)

pull/1705/head
Glenn Jocher 2020-12-15 21:50:28 -08:00 committed by GitHub
parent 035ac82ed0
commit 7947c86b57
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 9 deletions

View File

@ -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()