PyTorch之tensorboard

最近新get的技能，可视化训练过程和网络结构，比之前看log日志直接多了，写个帖子记录下。

正文

TensorBoard是一个可视化工具，它可以用来展示网络图、张量的指标变化、张量的分布情况等。特别是在训练网络的时候，我们可以设置不同的参数（比如：权重W、偏置B、卷积层数、全连接层数等），使用TensorBoader可以很直观的帮我们进行参数的选择。

示例

首先导入包

from tensorboardX import SummaryWriter

创建tensorboard

writer = SummaryWriter(snapshot_path + '/log') # 括号内为保存路径

记录网络结构

writer.add_graph(model, torch.randn(1,3,img_size,img_size).cuda())

注意 如果网络已经放在gpu上，生成随机输入时要.cuda()如果放在cpu上则不需要

可视化结果

写入相关参数

writer.add_scalar('info/loss', loss, iter_num)
writer.add_scalar('info/lr', optimizer.state_dict()['param_groups'][0]['lr'], epoch_num)

可以根据迭代次数或者epoch写入相关的信息，如损失 或者动态调整的学习率。
可视化结果

写入训练过程中测试结果

image = image_batch[0, 0:1, :, :] # 读取批处理第一张
image = (image - image.min()) / (image.max()-image.min()) # 归一化
writer.add_image('train/Image', image, iter_num) # 写入输入数据
outputs = torch.argmax(torch.softmax(outputs, dim=1), dim=1, keepdim=True) # 测试结果
writer.add_image('train/Prediction', outputs[0, ...] * 50, iter_num) # 写入
labs = label_batch[0, ...].unsqueeze(0) * 50 # 标签
writer.add_image('train/GroundTruth', labs, iter_num) # 写入标签

可视化结果

启动tensorboard

在终端输入如下指令（注意 依赖于python torch环境）

tensorboard --logdir="log路径" --port 6006 

—logdir 后面是代码中保存的路径
—port 设置监听端口号 可以不输入就会分配一个默认的，因为6006可能被占用

远程查看

因为代码都在服务器端训练，所以日志也保存在服务器，方便本地访问可以通过ssh端口映射
在服务器端启动tensorboard后 另开一个终端输入如下指令

ssh -L 16006:127.0.0.1:6006 username@server_ip

可以把服务器的6006端口映射到本地16006，通过浏览器访问 http://127.0.0.1:16006/ 即可访问

Mac端 safari浏览器打开会不显示，使用chrome浏览器没有影响。

tensorboard是个很强大的可视化工具，以上内容仅是记录下简单使用。后期摸索出其他用法后会持续更新