From 2bc6683325559ee39034a13489b5059056cb629e Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 20 Dec 2019 09:07:25 -0800
Subject: [PATCH] updates

---
 test.py  | 45 +++++++++++++++++++++++++++++++--------------
 train.py | 19 +++++++++----------
 2 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/test.py b/test.py
index 1c98e4bc..4b33d761 100644
--- a/test.py
+++ b/test.py
@@ -78,12 +78,15 @@ def test(cfg,
         if batch_i == 0 and not os.path.exists('test_batch0.jpg'):
             plot_images(imgs=imgs, targets=targets, paths=paths, fname='test_batch0.jpg')
 
-        # Run model
-        inf_out, train_out = model(imgs)  # inference and training outputs
+        # Disable gradients
+        with torch.no_grad():
 
-        # Compute loss
-        if hasattr(model, 'hyp'):  # if model has loss hyperparameters
-            loss += compute_loss(train_out, targets, model)[1][:3].cpu()  # GIoU, obj, cls
+            # Run model
+            inf_out, train_out = model(imgs)  # inference and training outputs
+
+            # Compute loss
+            if hasattr(model, 'hyp'):  # if model has loss hyperparameters
+                loss += compute_loss(train_out, targets, model)[1][:3].cpu()  # GIoU, obj, cls
 
         # Run NMS
         output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres)
@@ -220,12 +223,26 @@ if __name__ == '__main__':
     opt = parser.parse_args()
     print(opt)
 
-    with torch.no_grad():
-        test(opt.cfg,
-             opt.data,
-             opt.weights,
-             opt.batch_size,
-             opt.img_size,
-             opt.conf_thres,
-             opt.nms_thres,
-             opt.save_json or any([x in opt.data for x in ['coco.data', 'coco2014.data', 'coco2017.data']]))
+    # Test
+    test(opt.cfg,
+         opt.data,
+         opt.weights,
+         opt.batch_size,
+         opt.img_size,
+         opt.conf_thres,
+         opt.nms_thres,
+         opt.save_json or any([x in opt.data for x in ['coco.data', 'coco2014.data', 'coco2017.data']]))
+
+    # # Parameter study
+    # y = []
+    # x = np.arange(0.4, 0.81, 0.1)
+    # for v in x:
+    #     y.append(test(opt.cfg, opt.data, opt.weights, opt.batch_size, opt.img_size, 0.1, v, True)[0])
+    # y = np.stack(y, 0)
+    #
+    # fig, ax = plt.subplots(1, 1, figsize=(12, 6))
+    # ax.plot(x, y[:, 2], marker='.', label='mAP@0.5')
+    # ax.plot(x, y[:, 3], marker='.', label='mAP@0.5:0.95')
+    # ax.legend()
+    # fig.tight_layout()
+    # plt.savefig('parameters.jpg', dpi=200)
diff --git a/train.py b/train.py
index 3aeffe71..936ba39e 100644
--- a/train.py
+++ b/train.py
@@ -323,16 +323,15 @@ def train():
         if opt.prebias:
             print_model_biases(model)
         elif not opt.notest or final_epoch:  # Calculate mAP
-            with torch.no_grad():
-                is_coco = any([x in data for x in ['coco.data', 'coco2014.data', 'coco2017.data']]) and model.nc == 80
-                results, maps = test.test(cfg,
-                                          data,
-                                          batch_size=batch_size,
-                                          img_size=opt.img_size,
-                                          model=model,
-                                          conf_thres=0.001 if final_epoch else 0.1,  # 0.1 for speed
-                                          save_json=final_epoch and is_coco,
-                                          dataloader=testloader)
+            is_coco = any([x in data for x in ['coco.data', 'coco2014.data', 'coco2017.data']]) and model.nc == 80
+            results, maps = test.test(cfg,
+                                      data,
+                                      batch_size=batch_size,
+                                      img_size=opt.img_size,
+                                      model=model,
+                                      conf_thres=0.001 if final_epoch else 0.1,  # 0.1 for speed
+                                      save_json=final_epoch and is_coco,
+                                      dataloader=testloader)
 
         # Write epoch results
         with open(results_file, 'a') as f: