Update yolo.py (#2899)

This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-04-22 19:27:21 +02:00
1 parent b40dd99167
commit cdb678f418
1 file changed
+5 -5
+5 -5
View File
@@ -266,12 +266,12 @@ if __name__ == '__main__':
model.train() model.train()
# Profile # 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) # y = model(img, profile=True)
# Tensorboard # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898)
# from torch.utils.tensorboard import SummaryWriter # from torch.utils.tensorboard import SummaryWriter
# tb_writer = SummaryWriter() # tb_writer = SummaryWriter('.')
# print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") # print("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/")
# tb_writer.add_graph(model.model, img) # add model to tensorboard # 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 # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard