updates
This commit is contained in:
parent
f8b57a59ef
commit
2ee16e8280
25
test.py
25
test.py
|
@ -56,7 +56,8 @@ def test(cfg,
|
|||
model.eval()
|
||||
coco91class = coco80_to_coco91_class()
|
||||
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 = [], [], [], []
|
||||
for batch_i, (imgs, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
|
||||
targets = targets.to(device)
|
||||
|
@ -72,7 +73,7 @@ def test(cfg,
|
|||
|
||||
# Compute loss
|
||||
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
|
||||
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
|
||||
for i, c in enumerate(ap_class):
|
||||
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__':
|
||||
|
@ -204,12 +205,12 @@ if __name__ == '__main__':
|
|||
print(opt)
|
||||
|
||||
with torch.no_grad():
|
||||
mAP = test(opt.cfg,
|
||||
opt.data,
|
||||
opt.weights,
|
||||
opt.batch_size,
|
||||
opt.img_size,
|
||||
opt.iou_thres,
|
||||
opt.conf_thres,
|
||||
opt.nms_thres,
|
||||
opt.save_json)
|
||||
results = test(opt.cfg,
|
||||
opt.data,
|
||||
opt.weights,
|
||||
opt.batch_size,
|
||||
opt.img_size,
|
||||
opt.iou_thres,
|
||||
opt.conf_thres,
|
||||
opt.nms_thres,
|
||||
opt.save_json)
|
||||
|
|
6
train.py
6
train.py
|
@ -279,7 +279,7 @@ def train(cfg,
|
|||
mem = torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available() else 0 # (GB)
|
||||
s = ('%10s' * 2 + '%10.3g' * 7) % (
|
||||
'%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)
|
||||
if not (opt.notest or (opt.nosave and epoch < 10)) or epoch == epochs - 1:
|
||||
|
@ -289,10 +289,10 @@ def train(cfg,
|
|||
|
||||
# Write epoch results
|
||||
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
|
||||
fitness = results[2]
|
||||
fitness = results[2] # mAP
|
||||
if fitness > best_fitness:
|
||||
best_fitness = fitness
|
||||
|
||||
|
|
Loading…
Reference in New Issue