From 77469a5268fc3f08fbef7f0211a1cc72b39199ba Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 5 Nov 2018 23:17:53 +0100 Subject: [PATCH] update multi-scale training --- train.py | 13 +++---------- utils/datasets.py | 8 +++++++- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/train.py b/train.py index 59ca365e..e558882a 100644 --- a/train.py +++ b/train.py @@ -7,12 +7,11 @@ from utils.utils import * parser = argparse.ArgumentParser() parser.add_argument('-epochs', type=int, default=100, help='number of epochs') -parser.add_argument('-batch_size', type=int, default=16, help='size of each image batch') +parser.add_argument('-batch_size', type=int, default=2, help='size of each image batch') parser.add_argument('-data_config_path', type=str, default='cfg/coco.data', help='data config file path') parser.add_argument('-cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path') -parser.add_argument('-img_size', type=int, default=32 * 13, help='size of each image dimension') +parser.add_argument('-img_size', type=int, default=32 * 19, help='size of each image dimension') parser.add_argument('-resume', default=False, help='resume training flag') -parser.add_argument('-multi_scale', default=True, help='train at random img_size 320-608') # ensure memory for 608 size opt = parser.parse_args() print(opt) @@ -40,7 +39,7 @@ def main(opt): train_path = '../coco/trainvalno5k.part' # Initialize model - model = Darknet(opt.cfg, opt.img_size if opt.multi_scale is False else 608) + model = Darknet(opt.cfg, opt.img_size) # Get dataloader dataloader = load_images_and_labels(train_path, batch_size=opt.batch_size, img_size=opt.img_size, augment=True) @@ -100,12 +99,6 @@ def main(opt): for epoch in range(opt.epochs): epoch += start_epoch - # Multi-Scale YOLO Training - if opt.multi_scale: - img_size = random.choice(range(10, 20)) * 32 # 320 - 608 pixels - dataloader = load_images_and_labels(train_path, batch_size=opt.batch_size, img_size=img_size, augment=True) - print('Running Epoch %g at multi_scale img_size %g' % (epoch, img_size)) - # Update scheduler (automatic) # scheduler.step() diff --git a/utils/datasets.py b/utils/datasets.py index 37982fbf..4038d738 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -100,7 +100,13 @@ class load_images_and_labels(): # for training ia = self.count * self.batch_size ib = min((self.count + 1) * self.batch_size, self.nF) - height = self.height + if self.augment is True: + # Multi-Scale YOLO Training + height = random.choice(range(10, 20)) * 32 # 320 - 608 pixels + else: + # Fixed-Scale YOLO Training + height = self.height + print(height) img_all = [] labels_all = []