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

Question

假设我正在从 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

请务必注意，当您创建 DataLoader 对象时，它不会立即加载您的所有数据（这对于大型数据集不切实际）。它为您提供了一个迭代器，您可以使用它来访问每个样本。

遗憾的是，DataLoader 没有为您提供任何方法来控制您希望提取的样本数量。您将不得不使用切片迭代器的典型方法。

最简单的做法（没有任何库）是在达到所需的样本数量后停止。

nsamples = 10000
for i, image, label in enumerate(train_loader):
    if i > nsamples:
        break

    # Your training code here.

或者，您可以使用 itertools.islice 获取前 10k 个样本。像这样。

for image, label in itertools.islice(train_loader, stop=10000):

    # your training code here.

Answer 2

另一种快速切片数据集的方法是使用 torch.utils.data.random_split()（在 PyTorch v0.4.1+ 中支持）。它有助于将数据集随机拆分为给定长度的非重叠新数据集。

所以我们可以有类似下面的东西：

tr = datasets.MNIST('../data', train=True, download=True, transform=transform)
te = datasets.MNIST('../data', train=False, transform=transform)

part_tr = torch.utils.data.random_split(tr, [tr_split_len, len(tr)-tr_split_len])[0]
part_te = torch.utils.data.random_split(te, [te_split_len, len(te)-te_split_len])[0]

train_loader = DataLoader(part_tr, batch_size=args.batch_size, shuffle=True, num_workers=4, **kwargs)
test_loader = DataLoader(part_te, batch_size=args.batch_size, shuffle=True, num_workers=4, **kwargs)

这里可以分别设置tr_split_len和te_split_len作为训练和测试数据集所需的分割长度。

Answer 3

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

import torch.utils.data as data_utils

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

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

How do you alter the size of a Pytorch Dataset?

python

machine-learning

dataset

torch

pytorch