add multi_scale train option to argparser

This commit is contained in:
Glenn Jocher 2018-11-05 09:07:15 +01:00
parent b352f93f19
commit dc7b58bb3c
2 changed files with 12 additions and 12 deletions

View File

@ -209,9 +209,9 @@ class YOLOLayer(nn.Module):
class Darknet(nn.Module): class Darknet(nn.Module):
"""YOLOv3 object detection model""" """YOLOv3 object detection model"""
def __init__(self, config_path, img_size=416): def __init__(self, cfg_path, img_size=416):
super(Darknet, self).__init__() super(Darknet, self).__init__()
self.module_defs = parse_model_config(config_path) self.module_defs = parse_model_config(cfg_path)
self.module_defs[0]['height'] = img_size self.module_defs[0]['height'] = img_size
self.hyperparams, self.module_list = create_modules(self.module_defs) self.hyperparams, self.module_list = create_modules(self.module_defs)
self.img_size = img_size self.img_size = img_size

View File

@ -6,12 +6,13 @@ from utils.datasets import *
from utils.utils import * from utils.utils import *
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('-epochs', type=int, default=68, help='number of epochs') 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=16, 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('-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('-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 * 13, help='size of each image dimension')
parser.add_argument('-resume', default=False, help='resume training flag') 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')
opt = parser.parse_args() opt = parser.parse_args()
print(opt) print(opt)
@ -39,7 +40,7 @@ def main(opt):
train_path = '../coco/trainvalno5k.part' train_path = '../coco/trainvalno5k.part'
# Initialize model # Initialize model
model = Darknet(opt.cfg, opt.img_size) model = Darknet(opt.cfg, opt.img_size if opt.multi_scale is False else 608)
# Get dataloader # Get dataloader
dataloader = load_images_and_labels(train_path, batch_size=opt.batch_size, img_size=opt.img_size, augment=True) dataloader = load_images_and_labels(train_path, batch_size=opt.batch_size, img_size=opt.img_size, augment=True)
@ -100,20 +101,19 @@ def main(opt):
epoch += start_epoch epoch += start_epoch
# Multi-Scale YOLO Training # Multi-Scale YOLO Training
# img_size = random.choice(range(10, 20)) * 32 # 320 - 608 pixels if opt.multi_scale:
# dataloader = load_images_and_labels(train_path, batch_size=opt.batch_size, img_size=img_size, augment=True) img_size = random.choice(range(10, 20)) * 32 # 320 - 608 pixels
# print('Running this epoch with image size %g' % img_size) 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) # Update scheduler (automatic)
# scheduler.step() # scheduler.step()
# Update scheduler (manual) # Update scheduler (manual) at 0, 54, 61 epochs to 1e-3, 1e-4, 1e-5
if epoch < 54: if epoch < 50:
lr = 1e-3 lr = 1e-3
elif epoch < 61:
lr = 1e-4
else: else:
lr = 1e-5 lr = 1e-3
for g in optimizer.param_groups: for g in optimizer.param_groups:
g['lr'] = lr g['lr'] = lr