问题描述
在使用 yolov5 训练的过程中,想查看训练过程曲线。
问题解决
- 找到
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
- 进入项目的根目录,然后执行如下命令:
tensorboard --logdir=./runs
结果图
然后通过浏览器访问 http://localhost:6006/