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

时间:2017-07-01 01:52:41

标签: python machine-learning dataset torch pytorch

假设我从torchvision.datasets.MNIST加载MNIST,但我只想加载10000个图像,我如何切片数据以将其限制为只有一些数据点?我知道DataLoader是一个生成器,产生的数据大小与指定的批量大小相同,但是如何对数据集进行切片?

ScrollView

3 个答案:

答案 0 :(得分:7)

重要的是要注意,当您创建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.

答案 1 :(得分:7)

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

import torch.utils.data as data_utils

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

答案 2 :(得分:0)

切片数据集的另一种快速方法是使用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_lente_split_len分别设置为训练和测试数据集所需的分割长度。