diff --git a/test.py b/test.py index 1461c714..8b9776fe 100644 --- a/test.py +++ b/test.py @@ -209,7 +209,7 @@ if __name__ == '__main__': parser.add_argument('--cfg', type=str, default='cfg/yolov3-spp.cfg', help='*.cfg path') parser.add_argument('--data', type=str, default='data/coco2014.data', help='*.data path') parser.add_argument('--weights', type=str, default='weights/yolov3-spp.weights', help='path to weights file') - parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch') + parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch') parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold') parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression') diff --git a/train.py b/train.py index 5c939931..ee9caad2 100644 --- a/train.py +++ b/train.py @@ -213,7 +213,7 @@ def train(): rect=True, cache_labels=True, cache_images=opt.cache_images), - batch_size=batch_size, + batch_size=batch_size * 2, num_workers=nw, pin_memory=True, collate_fn=dataset.collate_fn) @@ -326,7 +326,7 @@ def train(): 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, + batch_size=batch_size * 2, img_size=opt.img_size, model=model, conf_thres=0.001 if final_epoch else 0.1, # 0.1 for speed