在PyTorch中使用DataLoaders进行验证数据集。

Question

我想在PyTorch和Torchvision中加载MNIST数据集，并将其分为训练集、验证集和测试集。目前，我已经有了以下代码：

def load_dataset():
    train_loader = torch.utils.data.DataLoader(
        torchvision.datasets.MNIST(
            '/data/', train=True, download=True,
            transform=torchvision.transforms.Compose([
                torchvision.transforms.ToTensor()])),
        batch_size=batch_size_train, shuffle=True)

    test_loader = torch.utils.data.DataLoader(
        torchvision.datasets.MNIST(
            '/data/', train=False, download=True,
            transform=torchvision.transforms.Compose([
                torchvision.transforms.ToTensor()])),
        batch_size=batch_size_test, shuffle=True)

如果我的训练数据集在 DataLoader 中，我该如何将其分成训练和验证集？我想将训练数据集的最后10000个样本作为验证数据集（我知道为了获得更准确的结果应该进行交叉验证，但这里我只想快速验证）。

- qalis

---

这个回答解决了你的问题吗？如何将自定义数据集分成训练集和测试集？ - kHarshit

@kHarshit 是的和不是的，经过进一步阅读，似乎这是一个更复杂的问题。我已经找到了解决方案，但它比普通库中应该有的要复杂。 - qalis

Answer 1

将训练数据集在PyTorch中拆分为训练集和验证集，比它应该的要困难得多。

首先，将训练集拆分为训练子集和验证子集（使用Subset类），它们不是数据集（Dataset类）：

train_subset, val_subset = torch.utils.data.random_split(
        train, [50000, 10000], generator=torch.Generator().manual_seed(1))

然后从这些数据集中获取实际数据：

X_train = train_subset.dataset.data[train_subset.indices]
y_train = train_subset.dataset.targets[train_subset.indices]

X_val = val_subset.dataset.data[val_subset.indices]
y_val = val_subset.dataset.targets[val_subset.indices]

注意，这种方式没有Dataset对象，因此我们不能使用DataLoader对象进行批量训练。如果要使用DataLoaders，它们直接与Subsets一起使用：

train_loader = DataLoader(dataset=train_subset, shuffle=True, batch_size=BATCH_SIZE)
val_loader = DataLoader(dataset=val_subset, shuffle=False, batch_size=BATCH_SIZE)

- qalis

---

是的...就是这样！ - Monica Heddneck

val_loader 应该将 val_subset 作为数据集参数，而不是 train_subset，对吗？ - R. Rayl

Answer 2

如果您想确保拆分的类别平衡，可以使用来自sklearn的train_test_split。

import torchvision
from torch.utils.data import DataLoader, Subset
from sklearn.model_selection import train_test_split

VAL_SIZE = 0.1
BATCH_SIZE = 64

mnist_train = torchvision.datasets.MNIST(
    '/data/',
    train=True,
    download=True,
    transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
)
mnist_test = torchvision.datasets.MNIST(
    '/data/',
    train=False,
    download=True,
    transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
)

# generate indices: instead of the actual data we pass in integers instead
train_indices, val_indices, _, _ = train_test_split(
    range(len(mnist_train)),
    mnist_train.targets,
    stratify=mnist_train.targets,
    test_size=VAL_SIZE,
)

# generate subset based on indices
train_split = Subset(mnist_train, train_indices)
val_split = Subset(mnist_train, val_indices)

# create batches
train_batches = DataLoader(train_split, batch_size=BATCH_SIZE, shuffle=True)
val_batches = DataLoader(val_split, batch_size=BATCH_SIZE, shuffle=True)
test_batches = DataLoader(mnist_test, batch_size=BATCH_SIZE, shuffle=True)

- Eric