mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

Update train.py

Browse Source 
This updates the PR to a one-liner to minimize additions. Perhaps we can include opt in the future but let's start with this for now.
This commit is contained in:
Glenn Jocher 2020-07-14 12:32:08 -07:00 committed by GitHub
parent 38acc5f3c5
commit 120d40c06a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
train.py
View File

@ -58,14 +58,6 @@ def train(hyp):
with open(Path(log_dir) / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Log hyperparameters in tensorboard
if tb_writer:
tb_hparams_dict = hyp
tb_hparams_dict.update(vars(opt))
tb_hparams_dict['img_size_train'], tb_hparams_dict['img_size_test'] = tb_hparams_dict['img_size']
del tb_hparams_dict['img_size']
tb_writer.add_hparams(tb_hparams_dict, {})
epochs = opt.epochs # 300
batch_size = opt.batch_size # 64
weights = opt.weights # initial training weights
@ -194,6 +186,7 @@ def train(hyp):
# model._initialize_biases(cf.to(device))
plot_labels(labels, save_dir=log_dir)
if tb_writer:
tb_writer.add_hparams(hyp, {})
tb_writer.add_histogram('classes', c, 0)
# Check anchors