This commit is contained in:
Glenn Jocher 2019-08-04 00:12:46 +02:00
parent f8b57a59ef
commit 2ee16e8280
2 changed files with 16 additions and 15 deletions

View File

@ -56,7 +56,8 @@ def test(cfg,
model.eval() model.eval()
coco91class = coco80_to_coco91_class() coco91class = coco80_to_coco91_class()
s = ('%30s' + '%10s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP', 'F1') s = ('%30s' + '%10s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP', 'F1')
loss, p, r, f1, mp, mr, map, mf1 = 0., 0., 0., 0., 0., 0., 0., 0. p, r, f1, mp, mr, map, mf1 = 0., 0., 0., 0., 0., 0., 0.
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)):
targets = targets.to(device) targets = targets.to(device)
@ -72,7 +73,7 @@ def test(cfg,
# 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)[0].item() loss += compute_loss(train_out, targets, model)[1][[0, 2, 3]] # GIoU, obj, cls
# Run NMS # Run NMS
output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres) output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres)
@ -186,7 +187,7 @@ def test(cfg,
maps = np.zeros(nc) + map maps = np.zeros(nc) + map
for i, c in enumerate(ap_class): for i, c in enumerate(ap_class):
maps[c] = ap[i] maps[c] = ap[i]
return (mp, mr, map, mf1, loss / len(dataloader)), maps return (mp, mr, map, mf1, *(loss / len(dataloader)).tolist()), maps
if __name__ == '__main__': if __name__ == '__main__':
@ -204,7 +205,7 @@ if __name__ == '__main__':
print(opt) print(opt)
with torch.no_grad(): with torch.no_grad():
mAP = test(opt.cfg, results = test(opt.cfg,
opt.data, opt.data,
opt.weights, opt.weights,
opt.batch_size, opt.batch_size,

View File

@ -279,7 +279,7 @@ def train(cfg,
mem = torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available() else 0 # (GB) mem = torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available() else 0 # (GB)
s = ('%10s' * 2 + '%10.3g' * 7) % ( s = ('%10s' * 2 + '%10.3g' * 7) % (
'%g/%g' % (epoch, epochs - 1), '%.3gG' % mem, *mloss, len(targets), img_size) '%g/%g' % (epoch, epochs - 1), '%.3gG' % mem, *mloss, len(targets), img_size)
pbar.set_description(s) # print(s) pbar.set_description(s)
# Calculate mAP (always test final epoch, skip first 5 if opt.nosave) # Calculate mAP (always test final epoch, skip first 5 if opt.nosave)
if not (opt.notest or (opt.nosave and epoch < 10)) or epoch == epochs - 1: if not (opt.notest or (opt.nosave and epoch < 10)) or epoch == epochs - 1:
@ -289,10 +289,10 @@ def train(cfg,
# Write epoch results # Write epoch results
with open('results.txt', 'a') as file: with open('results.txt', 'a') as file:
file.write(s + '%11.3g' * 5 % results + '\n') # P, R, mAP, F1, test_loss file.write(s + '%11.3g' * 7 % results + '\n') # P, R, mAP, F1, test_losses=(GIoU, obj, cls)
# Update best map # Update best map
fitness = results[2] fitness = results[2] # mAP
if fitness > best_fitness: if fitness > best_fitness:
best_fitness = fitness best_fitness = fitness