torch下可视化 训练+计算图

记录loss

安装 tensorboard

pip install tensorboard   # 版本 > 1.15

导入模块

from torch.utils.tensorboard import SummaryWriter

创建tb文件夹，创建writer

# tensorboard
utils.mkdir(opt['train']['tb_folder'])
tb_writer_tr = SummaryWriter(log_dir = opt['train']['tb_folder'])
tb_writer_te = SummaryWriter(log_dir = opt['train']['tb_folder'])

要记录的loss值

loss_EveryEpoch = 0
loss_EveryEpoch = loss_EveryEpoch + train_loss 
loss_EveryEpoch = loss_EveryEpoch + test_loss  

使用writer写入到本地tensorboard文件夹

tb_writer_tr.add_scalar('loss_EveryEpoch', loss_EveryEpoch / (step - 1), epoch)
tb_writer_te.add_scalar('loss_EveryEpoch', loss_EveryEpoch / (step - 1), epoch)

可视化

tensorboard --logdir=./runs

可视化计算图
1、使用tensorboard

#
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()
#
writer.add_graph(model, input_data)
writer.close()

2、使用 torchviz

from torchviz import make_dot
g = make_dot(tensor, params=dict(list(model.named_parameters())))
# tensor 为model 的output
# g.view()
g.render('wm_model', view=False) # 保存pdf

#然并卵， 计算图的两个可视化工具都没能成功可视化，，，，，弃坑