在Pytorch中使用Tensorboard

Tensorboard是一款深度学习可视化软件，目前主要使用了它的可视化模型, 可视化模型权重和可视化损失函数功能。

x.1 tensorboard初始化

tensorboard初始化需要导入SummaryWriter包并指定存储位置和开放端口号。

from torch.utils.tensorboard import SummaryWriter

print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/')
# 实例化SummaryWriter对象
tb_writer = SummaryWriter(log_dir="/home/yingmuzhi/_learning/src/tensorboard/runs/flower_experiment")

x.2 可视化模型

在模型初始化后做该操作 在可视化模型时，我们需要传入和输入大小一样的Tensor，使用add_graph来看网络模型大小：

    # 将模型写入tensorboard
    init_img = torch.zeros((1, 3, 224, 224), device=device)
    tb_writer.add_graph(model, init_img)

x.3 可视化模型权重在每次epoch中变化

在每一个epoch中添加该方法 通过add_histogram方法来添加需要查看的网络的结构，如conv1。

        # add conv1 weights into tensorboard
        tb_writer.add_histogram(tag="conv1",
                                values=model.conv1.weight,
                                global_step=epoch)
        tb_writer.add_histogram(tag="layer1/block0/conv1",
                                values=model.layer1[0].conv1.weight,
                                global_step=epoch)

x.4 可视化损失函数在每次epoch中变化

在每一个epoch中添加该方法 通过add_scalar方法来添加需要绘制的损失函数，如mean_loss。

        # add loss, acc and lr into tensorboard
        print("[epoch {}] accuracy: {}".format(epoch, round(acc, 3)))
        tags = ["train_loss", "accuracy", "learning_rate"]
        tb_writer.add_scalar(tags[0], mean_loss, epoch)
        tb_writer.add_scalar(tags[1], acc, epoch)
        tb_writer.add_scalar(tags[2], optimizer.param_groups[0]["lr"], epoch)