diff --git a/models.py b/models.py index 2c1f1085..1b12bd94 100755 --- a/models.py +++ b/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() diff --git a/test.py b/test.py index b2c11517..5b1fc3d9 100644 --- a/test.py +++ b/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') diff --git a/train.py b/train.py index 11ec9722..e523d2d4 100644 --- a/train.py +++ b/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)