最新版本yolox训练得到的结果,怎么用tensorboard查看total_loss曲线,需要怎么修改yolox/core/trainer.py代码
时间: 2023-06-17 12:05:29 浏览: 155
要使用 TensorBoard 查看 YOLOX 训练时的总损失曲线,您需要在 `yolox/core/trainer.py` 中添加以下内容:
1. 在导入的包中添加 TensorBoard 的包:
```
from torch.utils.tensorboard import SummaryWriter
```
2. 在 `__init__` 函数中添加 SummaryWriter:
```
def __init__(self, exp, args):
self.exp = exp
self.args = args
self.start_epoch = 0
self.best_map = 0.0
self.best_loss = float('inf')
self.writer = SummaryWriter(log_dir=self.exp.logdir)
```
3. 在 `train` 函数中添加记录 loss 的代码:
```
for i, (images, targets, _) in enumerate(dataloader):
self.scheduler.step(self.cur_iter)
self.cur_iter += 1
images = images.to(self.device)
targets = [target.to(self.device) for target in targets]
outputs = self.model(images)
loss, loss_items = self.loss_fn(outputs, targets)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
if self.rank == 0:
# 记录 loss
for k, v in loss_items.items():
self.writer.add_scalar('train/{}'.format(k), v, self.cur_iter)
self.writer.add_scalar('train/total_loss', loss, self.cur_iter)
if self.cur_iter % self.args.eval_interval == 0:
self.evaluate(self.cur_iter)
```
在上述代码中,`self.writer.add_scalar('train/total_loss', loss, self.cur_iter)` 记录了每个迭代的总损失,您可以在 TensorBoard 中查看该曲线。
最后,您需要在命令行中运行以下命令启动 TensorBoard:
```
tensorboard --logdir=path/to/logdir
```
其中 `path/to/logdir` 是您在 `exp.logdir` 中定义的路径。然后在浏览器中打开 TensorBoard 网址即可查看损失曲线。
阅读全文