updates

models.py

@@ -184,7 +184,6 @@ class YOLOLayer(nn.Module):
 
             # Sum loss components
             balance_losses_flag = True
-
             if balance_losses_flag:
                 k = 1 / self.loss_means.clone()
                 loss = (lx * k[0] + ly * k[1] + lw * k[2] + lh * k[3] + lconf * k[4] + lcls * k[5]) / k.mean()

test.py

@@ -4,7 +4,7 @@ from models import *
 from utils.datasets import *
 from utils.utils import *
 
-parser = argparse.ArgumentParser()
+parser = argparse.ArgumentParser(prog='test.py')
 parser.add_argument('-batch_size', type=int, default=32, help='size of each image batch')
 parser.add_argument('-cfg', type=str, default='cfg/yolov3.cfg', help='path to model config file')
 parser.add_argument('-data_config_path', type=str, default='cfg/coco.data', help='path to data config file')

train.py

@@ -1,4 +1,5 @@
 import argparse
+import sys
 import time
 
 from models import *
@@ -16,6 +17,10 @@ parser.add_argument('-batch_report', default=False, help='report TP, FP, FN, P a
 opt = parser.parse_args()
 print(opt)
 
+# Import test.py to get mAP after each epoch
+sys.argv[1:] = []  # delete any command line arguments that might get picked up by test.py
+import test  # must follow sys.argv[1:] = []
+
 cuda = torch.cuda.is_available()
 device = torch.device('cuda:0' if cuda else 'cpu')
 
@@ -184,7 +189,6 @@ def main(opt):
             os.system('cp weights/latest.pt weights/backup' + str(epoch) + '.pt')
 
         # Calculate mAP
-        import test
         test.opt.weights_path = 'weights/latest.pt'
         mAP, R, P = test.main(test.opt)