【tensorboard】保存训练日志并可视化

文章概览

生成log文件

from torch.utils.tensorboard import SummaryWriter 
writer = SummaryWriter(log_dir='logs')

# 写入loss
writer.add_scalar('Train/Loss', loss.item(), epoch)

# 写入网络结构图
writer.add_graph(model, input)

# 写入权重分布图
for name, param in model.named_parameters():
    writer.add_histogram(name, param, global_step)

log文件可视化

tensorboard --logdir "logs"，log是log文件所在文件夹

多想多做，发篇一作