示例代码:
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
# 准备的测试数据集
test_data = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor())
# batch_size=4 取test_data[0]到test_data[3] 返回 打包好的img0-3, 打包好的target0-3(shuffle=True随机抓取)
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=True)
# 测试数据集中第一张图片及target
img, target = test_data[0]
print(img.shape)
print(target)
writer = SummaryWriter("dataloader")
step = 0
for data in test_loader:
imgs, targets = data
# print(imgs.shape)
# print(targets)
writer.add_images("test_data_drop_last", imgs, step)
step = step+1
writer.close()
# batch_size=4 取test_data[0]到test_data[3] 返回 打包好的img0-3, 打包好的target0-3(随机抓取)
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=False)
设置drop_last=False 结果,最后一步不足64张任然进行了保留
设置drop_last=True后
最后一步不足64张进行了舍去,所以只有155步