Browse Source

Fix yaml saving (don't sort keys), reorder --opt keys, bug fix hyp dict accessor

5.0
Alex Stoken 4 years ago
parent
commit
de191655e4
1 changed files with 7 additions and 6 deletions
  1. +7
    -6
      train.py

+ 7
- 6
train.py View File

else: else:
pg0.append(v) # all else pg0.append(v) # all else


if hyp.optimizer =='adam':
if hyp['optimizer'] =='adam':
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) #use default beta2, adjust beta1 for Adam momentum per momentum adjustments in https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) #use default beta2, adjust beta1 for Adam momentum per momentum adjustments in https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
else: else:
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)


#save hyperparamter and training options in run folder #save hyperparamter and training options in run folder
with open(os.path.join(log_dir, 'hyp.yaml'), 'w') as f: with open(os.path.join(log_dir, 'hyp.yaml'), 'w') as f:
yaml.dump(hyp, f)
yaml.dump(hyp, f, sort_keys=False)


with open(os.path.join(log_dir, 'opt.yaml'), 'w') as f: with open(os.path.join(log_dir, 'opt.yaml'), 'w') as f:
yaml.dump(vars(opt), f)
yaml.dump(vars(opt), f, sort_keys=False)
# Class frequency # Class frequency
labels = np.concatenate(dataset.labels, 0) labels = np.concatenate(dataset.labels, 0)
if __name__ == '__main__': if __name__ == '__main__':
check_git_status() check_git_status()
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--batch-size', type=int, default=16)
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model cfg path[*.yaml]') parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model cfg path[*.yaml]')
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data cfg path [*.yaml]') parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data cfg path [*.yaml]')
parser.add_argument('--hyp', type=str, default='',help='hyp cfg path [*.yaml].')
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--batch-size', type=int, default=16)
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes. Assumes square imgs.') parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes. Assumes square imgs.')
parser.add_argument('--rect', action='store_true', help='rectangular training') parser.add_argument('--rect', action='store_true', help='rectangular training')
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
parser.add_argument('--hyp', type=str, default='', help ='hyp cfg path [*.yaml].')
opt = parser.parse_args() opt = parser.parse_args()


opt.cfg = check_file(opt.cfg) # check file opt.cfg = check_file(opt.cfg) # check file

Loading…
Cancel
Save