Skip to content

Commit

Permalink
Update train.py with yaml.SafeLoader (ultralytics#1972)
Browse files Browse the repository at this point in the history
  • Loading branch information
Anon-Artist authored Jan 18, 2021
1 parent a913314 commit 1513d12
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -479,7 +479,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
apriori = opt.global_rank, opt.local_rank
with open(Path(ckpt).parent.parent / 'opt.yaml') as f:
opt = argparse.Namespace(**yaml.load(f, Loader=yaml.FullLoader)) # replace
opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader)) # replace
opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = '', ckpt, True, opt.total_batch_size, *apriori # reinstate
logger.info('Resuming training from %s' % ckpt)
else:
Expand All @@ -503,7 +503,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):

# Hyperparameters
with open(opt.hyp) as f:
hyp = yaml.load(f, Loader=yaml.FullLoader) # load hyps
hyp = yaml.load(f, Loader=yaml.SafeLoader) # load hyps

# Train
logger.info(opt)
Expand Down

0 comments on commit 1513d12

Please sign in to comment.