TensorBoard add_graph() fix (#3236)
Browse files
train.py
CHANGED
@@ -330,9 +330,9 @@ def train(hyp, opt, device, tb_writer=None):
|
|
330 |
if plots and ni < 3:
|
331 |
f = save_dir / f'train_batch{ni}.jpg' # filename
|
332 |
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
|
333 |
-
|
334 |
-
|
335 |
-
|
336 |
elif plots and ni == 10 and wandb_logger.wandb:
|
337 |
wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
|
338 |
save_dir.glob('train*.jpg') if x.exists()]})
|
|
|
330 |
if plots and ni < 3:
|
331 |
f = save_dir / f'train_batch{ni}.jpg' # filename
|
332 |
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
|
333 |
+
if tb_writer:
|
334 |
+
tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph
|
335 |
+
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
|
336 |
elif plots and ni == 10 and wandb_logger.wandb:
|
337 |
wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
|
338 |
save_dir.glob('train*.jpg') if x.exists()]})
|