From b7cd1f540d5815b0a1cf2e23ce82a5fdb8f6b525 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 19 May 2021 13:51:55 +0200 Subject: [PATCH] TensorBoard add_graph() fix (#3236) --- train.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train.py b/train.py index 5330e4a..78d91fa 100644 --- a/train.py +++ b/train.py @@ -330,9 +330,9 @@ def train(hyp, opt, device, tb_writer=None): if plots and ni < 3: f = save_dir / f'train_batch{ni}.jpg' # filename Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start() - # if tb_writer: - # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) - # tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph + if tb_writer: + tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph + # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) elif plots and ni == 10 and wandb_logger.wandb: wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg') if x.exists()]})