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
                )