[英]How do you alter the size of a Pytorch Dataset?
假设我正在从 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)
需要注意的是,当您创建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.
切片数据集的另一种快速方法是使用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
设置为训练和测试数据集所需的分割长度。
您可以使用torch.utils.data.Subset()
例如前 10,000 个元素:
import torch.utils.data as data_utils
indices = torch.arange(10000)
tr_10k = data_utils.Subset(tr, indices)
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.