Tensorboard step和图片加载不完全处理办法

西雅图彭于晏

已于 2024-07-31 21:59:32 修改

阅读量782

点赞数 9

CC 4.0 BY-SA版权

文章标签：深度学习人工智能 pytorch

于 2024-07-31 21:39:30 首次发布

本文链接：https://blog.youkuaiyun.com/aozhoupengyuyan/article/details/140832926

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

# 加载CIFAR10测试数据集
# 参数说明:
# "./dataset": 数据集保存路径
# train=False: 加载测试集而非训练集
# transform: 将图像转换为PyTorch张量
test_data = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor())

# 创建DataLoader
# dataset=test_data: 使用上面加载的测试数据集
# batch_size=4: 每批加载4张图片
# shuffle=True: 随机打乱数据顺序
# num_workers=0: 不使用多进程加载数据
# drop_last=False: 不丢弃最后一个不完整的批次
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=False)

# 获取第一个数据样本
img, target = test_data[0]
print(img.shape)  # 打印图像的形状
print(target)  # 打印目标标签

# 创建TensorBoard的SummaryWriter对象
writer = SummaryWriter("dataloader")
for epoch in range(2):
    step = 0
    for data in test_loader:
        imgs, targets = data
        # 将每批图像添加到TensorBoard
        # "test_data": 图像的标识名
        # imgs: 图像数据
        # step: 当前步骤
        writer.add_images("test_data", imgs, step)
        step = step + 1

# 关闭SummaryWriter
writer.close()

小土堆Pytorch P15学习的时候发现直接在terminal中使用这行代码，tensorboard加载出来的step是不完全的

tensorboard --logdir=dataloader

因为在默认情况下，TensorBoard 可能会限制显示的图像样本数量以节约资源，如果想要所有的图片和step都可以完整的显示需要使用这一行代码

tensorboard --logdir=dataloader --samples_per_plugin=images=10000

通过在terminal执行这个代码，就可以看到完成的step

--samples_per_plugin 参数可以控制每个插件展示的最大样本数。默认情况下，TensorBoard 可能会限制显示的图像样本数量以节约资源。而使用 --samples_per_plugin=images=10000 可以显著增加 TensorBoard 显示的图像样本数量。