可视化工具tensorboardX

安装

pip install tensorboardX

看曲线

tb_logger.add_scalar('loss_train', loss, curr_step)

loss

可视化中间过程

from torchvision.utils import make_grid
tb_logger.add_image('image', input[0].cpu(), curr_step)
tb_logger.add_image('channels', make_grid(feature_map[0].detach().cpu().unsqueeze(dim=1), nrow=5, padding=20, normalize=False, pad_value=1), curr_step)

featuremap
可视化model(完整代码)

dummy_input = Variable(torch.rand(13, 1, 28, 28))

model = Net1()
with SummaryWriter(comment='Net1') as w:
    w.add_graph(model, (dummy_input, ))

model = Net2()
with SummaryWriter(comment='Net2') as w:
    w.add_graph(model, (dummy_input, ))

在tensorboard中查看

tensorboard --logdir /path/log

参考：

• 官网
• pytorch可视化