python - 如何改变 Pytorch 数据集的大小？

Question

这个问题已经有答案了:

Taking subsets of a pytorch dataset (2 个回答)

已关闭 去年。

假设我正在从 torchvision.datasets.MNIST 加载 MNIST ，但我只想加载总共 10000 个图像，我将如何切片数据以将其限制为仅某些数量的数据点？据我了解，DataLoader是生成指定批量大小的数据的生成器，但是如何对数据集进行切片？

tr = datasets.MNIST('../data', train=True, download=True, transform=transform)
te = datasets.MNIST('../data', train=False, transform=transform)
train_loader = DataLoader(tr, batch_size=args.batch_size, shuffle=True, num_workers=4, **kwargs)
test_loader = DataLoader(te, batch_size=args.batch_size, shuffle=True, num_workers=4, **kwargs)

Answer 1

您可以使用torch.utils.data.Subset()，例如对于前 10,000 个元素:

import torch.utils.data as data_utils

indices = torch.arange(10000)
tr_10k = data_utils.Subset(tr, indices)