yolov5 tensorboard 可视化查看

问题描述

在使用 yolov5 训练的过程中，想查看训练过程曲线。

问题解决

1. 找到 models/yolo.py 文件，将最下面有关 Tensorboard 的注释打开：

    # Create model
    model = Model(opt.cfg).to(device)
    model.train()

    # Profile
    if opt.profile:
        img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
        y = model(img, profile=True)

    # Test all models
    if opt.test:
        for cfg in Path(ROOT / 'models').rglob('yolo*.yaml'):
            try:
                _ = Model(cfg)
            except Exception as e:
                print(f'Error in {cfg}: {e}')

    # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898)
    from torch.utils.tensorboard import SummaryWriter
    tb_writer = SummaryWriter('.')
    LOGGER.info("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

2. 进入项目的根目录，然后执行如下命令：

tensorboard --logdir=./runs

结果图

然后通过浏览器访问 http://localhost:6006/