Pytorch：使用 DataLoader 加载图像样本

Question

我使用来自 torch.utils.data 的标准 DataLoader。我创建数据集类，然后以这种方式构建 DataLoader：

train_dataset = LandmarksDataset(os.path.join(args.data, 'train'), train_transforms, split="train")
train_dataloader = data.DataLoader(train_dataset, batch_size=args.batch_size, num_workers=2,
                                   pin_memory=True, shuffle=True, drop_last=True)

它运行完美，但数据集足够大 - 300k 图像。因此，使用 DataLoader 读取图像需要花费大量时间。所以在debug阶段搭建这么大的DataLoader实在是太可悲了！我只是想测试一些我的假设并想快速完成！我不需要为此加载整个数据集。

我正在尝试如何在不在整个数据集上构建 dataLoader 的情况下仅加载数据集的一小部分固定部分？ 目前我所有的想法只是创建另一个文件夹，在此处复制部分图像并在其上使用管道。但我想，Pytorch 足够聪明，有一些内置方法可以从大数据集中加载一部分图像。你能给我建议吗？

Answer 1

据我所知，没有任何机制可以为您做到这一点。您的问题出在 LandmarksDataset 类中，您正在读取火车数据文件夹的路径。我假设os.listdir(train_data_folder)。

相反，您可以使用更有效的方式 os.scandir(train_data_folder) 这将返回一个生成器并在其上调用 next() 将为您提供火车数据中图像的路径。这样，您可以多次调用 next()，而无需更改训练数据文件夹的结构并构建它的子集。