From 0964c95eecdc4a427bb8128f5fe05e2181a6bed0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 22 Apr 2021 19:26:47 +0200 Subject: [PATCH] Update yolo.py --- models/yolo.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/models/yolo.py b/models/yolo.py index 7db0e7da2629..36fa27e89134 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -264,14 +264,14 @@ def parse_model(d, ch): # model_dict, input_channels(3) # Create model model = Model(opt.cfg).to(device) model.train() - + # Profile - # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) + # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 320, 320).to(device) # y = model(img, profile=True) - # Tensorboard + # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898) # from torch.utils.tensorboard import SummaryWriter - # tb_writer = SummaryWriter() - # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") - # tb_writer.add_graph(model.model, img) # add model to tensorboard + # tb_writer = SummaryWriter('.') + # print("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/") + # tb_writer.add_graph(torch.jit.trace(model, img, strict=False), []) # add model graph # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard