在自己代码上使用tensorboard进行训练可视化

在自己代码上使用tensorboard进行训练可视化

因为自己在训练代码的过程中，有需要进行可视化的需求。所以就探究了下，如何进行可视化，并记录下来。
可视化的工具有很多，这里采用tensorboard进行训练可视化。

---

1. 导入tensorboard库

在这里看有些人都是下载tensorflow从而得到tensorboard，然后下载安装tensorboardX。
但是现在torch好像把tensorboard集成了进来，使用下面的代码导入，即可以使用tenosrboard。

from torch.utils.tensorboard import SummaryWriter

但是第一次使用可能会出现报错
出现ModuleNotFoundError: No module named ‘tensorboard‘错误
解决办法：
在自己运行的安装环境下，安装tb-nightly

pip install tb-nightly

2. 实例化SummaryWriter对象

在训练的main函数中，实例化SummaryWriter对象

#  实例化SummaryWriter对象
tb_writer = SummaryWriter(log_dir='runs')

这里的runs代表的是在当前目录下新建名为runs的文件夹，用于存储tensorboard的数据。

3. 将模型写入tensorboard中

这一步就是在实例化模型之后进行

    # 将模型写入tensorboard中
    init_img = torch.zeros((1,1