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 显示的图像样本数量。