This commit is contained in:
Glenn Jocher 2019-04-02 18:50:55 +02:00
parent 7f220a14cb
commit f527d30ccd
1 changed files with 3 additions and 3 deletions

View File

@ -70,10 +70,10 @@ def train(
del chkpt del chkpt
else: # Initialize model with backbone (optional) else: # Initialize model with backbone (optional)
if cfg.endswith('yolov3.cfg'): if '-tiny.cfg' in cfg:
cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')
elif cfg.endswith('yolov3-tiny.cfg'):
cutoff = load_darknet_weights(model, weights + 'yolov3-tiny.conv.15') cutoff = load_darknet_weights(model, weights + 'yolov3-tiny.conv.15')
else:
cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')
# Set scheduler (reduce lr at epoch 250) # Set scheduler (reduce lr at epoch 250)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[250], gamma=0.1, last_epoch=start_epoch - 1) scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[250], gamma=0.1, last_epoch=start_epoch - 1)