安装tensorboard:
pip install tensorboardX
在训练中导入:
from tensorboardX import SummaryWriter
在程序中调用:
writer = SummaryWriter(‘log’)
保存要做图的所有信息,在目录下建立一个log文件夹存放。
绘制reward, loss曲线:
writer.add_scalar('Train/reward', reward.data[0], niter)
writer.add_scalar('Train/Loss', loss.data[0], niter)
在训练过程中,依次写入:图像名称;reward, loss数值;n_interation。