From e5589632446526f1db934991fcab71c2eb42464d Mon Sep 17 00:00:00 2001
From: Janko Ondras <jondras@cisco.com>
Date: Mon, 13 Jul 2020 12:55:43 +0200
Subject: [PATCH] Log hyperparameters in tensorboard

Log both hyperparameters and command line options in tensorboard.
---
 train.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/train.py b/train.py
index 1fca64c53..7b3494a19 100644
--- a/train.py
+++ b/train.py
@@ -58,6 +58,14 @@ 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_w'], tb_hparams_dict['img_size_h'] = 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