这篇文章介绍如何在 PyTorch 中使用 TensorBoard 记录训练数据。
记录数据
初始化
在程序启动时创建 SummaryWriter 对象用于写入日志数据。
from torch.utils.tensorboard import SummaryWriter
import datetime# 获取当前时间戳，一般以时间戳作为记录文件夹名称的一部分
timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
# 创建 SummaryWriter 对象，用于记录训练和测试数据
writer = SummaryWriter(f'runs/mnist_{timestamp}')
我们创建了一个 SummaryWriter 并将数据写入了 runs 目录。
记录模型结构
在创建模型实例和数据加载器后，记录模型结构：
# 获取一个 batch 的样本
sample_images, _ = next(iter(train_loader))
# 记录模型结构
writer.add_graph(model, sample_images.to(device))
可以在 TensorBoard 的 GRAPHS 板块看到记录的模型结构：
记录数据集样本
可以记录训练/测试过程中使用的数据集样本：
if epoch == 0:# 获取一个 batch 的样本images, _ = next(iter(test_loader))# 将前 25 张图像样本转换为网格图img_grid = torchvision.utils.make_grid(images[:25])# 将网格图添加到 TensorBoard 中writer.add_image('mnist_images', img_grid)
可以在 TensorBoard 的 IMAGES 板块看到记录的数据集样本：
记录模型指标
可以在每个 epoch 结束时记录 loss 和 accuracy：
writer.add_scalar('Training Loss/Epoch', avg_loss, epoch)
writer.add_scalar('Training Accuracy/Epoch', accuracy, epoch)
可以在 TensorBoard 的 SCALARS 板块看到记录的数据：
查看数据
安装依赖：
pip install tensorboard torch_tb_profiler
命令行启动：
tensorboard --logdir=runs # 这里的 logdir 就是创建 SummaryWriter 对象时指定的目录
在 VS Code 中，可以使用下面的命令启动 TensorBoard：
> Python: Launch TensorBoard
