glenn-jocher commited on
Commit
7947c86
·
unverified ·
1 Parent(s): 035ac82

Update COCO train postprocessing (#1702)

Browse files
Files changed (1) hide show
  1. train.py +12 -9
train.py CHANGED
@@ -402,15 +402,18 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
402
 
403
  # Test best.pt
404
  if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
405
- results, _, _ = test.test(opt.data,
406
- batch_size=total_batch_size,
407
- imgsz=imgsz_test,
408
- model=attempt_load(best if best.exists() else last, device).half(),
409
- single_cls=opt.single_cls,
410
- dataloader=testloader,
411
- save_dir=save_dir,
412
- save_json=True, # use pycocotools
413
- plots=False)
 
 
 
414
 
415
  else:
416
  dist.destroy_process_group()
 
402
 
403
  # Test best.pt
404
  if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
405
+ for conf, iou, save_json in ([0.25, 0.45, False], [0.001, 0.65, True]): # speed, mAP tests
406
+ results, _, _ = test.test(opt.data,
407
+ batch_size=total_batch_size,
408
+ imgsz=imgsz_test,
409
+ conf_thres=conf,
410
+ iou_thres=iou,
411
+ model=attempt_load(best if best.exists() else last, device).half(),
412
+ single_cls=opt.single_cls,
413
+ dataloader=testloader,
414
+ save_dir=save_dir,
415
+ save_json=save_json,
416
+ plots=False)
417
 
418
  else:
419
  dist.destroy_process_group()