在自己代码上使用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,