This commit is contained in:
Glenn Jocher 2020-03-04 09:00:48 -08:00
parent eb81c0b9ae
commit e482392161
1 changed files with 11 additions and 1 deletions

12
test.py
View File

@ -67,7 +67,7 @@ def test(cfg,
model.eval() model.eval()
coco91class = coco80_to_coco91_class() coco91class = coco80_to_coco91_class()
s = ('%20s' + '%10s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@0.5', 'F1') s = ('%20s' + '%10s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@0.5', 'F1')
p, r, f1, mp, mr, map, mf1 = 0., 0., 0., 0., 0., 0., 0. p, r, f1, mp, mr, map, mf1, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3) loss = torch.zeros(3)
jdict, stats, ap, ap_class = [], [], [], [] jdict, stats, ap, ap_class = [], [], [], []
for batch_i, (imgs, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)): for batch_i, (imgs, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
@ -84,14 +84,18 @@ def test(cfg,
# Disable gradients # Disable gradients
with torch.no_grad(): with torch.no_grad():
# Run model # Run model
t = time.time()
inf_out, train_out = model(imgs) # inference and training outputs inf_out, train_out = model(imgs) # inference and training outputs
t0 += time.time() - t
# Compute loss # Compute loss
if hasattr(model, 'hyp'): # if model has loss hyperparameters if hasattr(model, 'hyp'): # if model has loss hyperparameters
loss += compute_loss(train_out, targets, model)[1][:3].cpu() # GIoU, obj, cls loss += compute_loss(train_out, targets, model)[1][:3].cpu() # GIoU, obj, cls
# Run NMS # Run NMS
t = time.time()
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres) output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres)
t1 += time.time() - t
# Statistics per image # Statistics per image
for si, pred in enumerate(output): for si, pred in enumerate(output):
@ -177,6 +181,11 @@ def test(cfg,
for i, c in enumerate(ap_class): for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap[i], f1[i])) print(pf % (names[c], seen, nt[c], p[i], r[i], ap[i], f1[i]))
# Print profile results
if opt.profile:
t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1))
print('Profile results: %.1f/%.1f/%.1f ms inference/NMS/total per image' % t)
# Save JSON # Save JSON
if save_json and map and len(jdict): if save_json and map and len(jdict):
imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files] imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files]
@ -220,6 +229,7 @@ if __name__ == '__main__':
parser.add_argument('--task', default='test', help="'test', 'study', 'benchmark'") parser.add_argument('--task', default='test', help="'test', 'study', 'benchmark'")
parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu') parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu')
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
parser.add_argument('--profile', action='store_true', help='profile inference and NMS times')
opt = parser.parse_args() opt = parser.parse_args()
opt.save_json = opt.save_json or any([x in opt.data for x in ['coco.data', 'coco2014.data', 'coco2017.data']]) opt.save_json = opt.save_json or any([x in opt.data for x in ['coco.data', 'coco2014.data', 'coco2017.data']])
print(opt) print(opt)