PyTorch笔记之TensorBoard可视化

Training With TensorBoard

TensorBoard setup

from torch.utils.tensorboard import SummaryWriter

# default 'log_dir' is 'runs'
write = SummaryWriter("runs/xx_experiment")

代码创建了 runs/xx_experiment 文件夹

Writing to TensorBoard

• 写入一张图片

# PyTorch 专门处理 CV的库
import torchvision

# create grid of images
img_grid = torchvision.utils.make_grid(images)

# write to tensorbard
writer.add_image("image", img_grid)

运行

tensorboard --logdir=runs

• 写入模型

# x is input data
writer.add_graph(model, x)

• 加入 Projector
  通过 add_embedding 可视化高维度数据

使用 TensorBoard 追踪模型训练过程

for epoch in range(epochs):
    for i, data in enumerate(db):
        ...
        if i % 1000 == 999:

            # ... log the running loss
            writer.add_scalar(
                    "Training loss",
                     running_loss / 1000,
                     epoch * len(db) + i
                )