PyTorch使用TensorboardX进行训练可视化

demo代码

if __name__ == '__main__':
    out="/home/zengxh/workspace/zdata"
    summarywriter=run_tensorboard(out,port=6006) # port为0的话是生成不冲突的随机端口

    # 记录模型
    model = torchvision.models.shufflenet_v2_x0_5()
    dummy_input = torch.randn(1, 3, 224, 224)
    summarywriter.add_graph(model, (dummy_input))

    # 记录训练时权重的值
    epoch = 10
    for name, param in model.named_parameters():
        summarywriter.add_histogram(name, param.clone().data, epoch)

    # 保存的是训练时候的图片
    train_inputs_make_grid = torchvision.utils.make_grid(dummy_input.to("cpu"), normalize=True,scale_each=True)
    summarywriter.add_image('Train Image{}'.format(epoch), train_inputs_make_grid,epoch)

    line1="this is ..."
    summarywriter.add_text('line1', str(line1), epoch)

    lr = 0.1
    for e in range(epoch):
        lr = lr * 0.1
        summarywriter.add_scalar("train/lr", lr, global_step=e)
        summarywriter.add_scalar("train/lr2", lr, global_step=e)
        summarywriter.add_scalar("val/lr2", lr, global_step=e)

    summarywriter.close()  # 关闭tensorboardX 日志

运行后会打印启动tensorboard的命令

或者是直接访问在浏览器中输入http://127.0.0.1:6006