updates

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)

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