如何使用 pytorch 将数据集拆分为自定义训练集和自定义验证集？答案

【问题标题】：How to split a dataset into a custom training set and a custom validation set with pytorch?如何使用 pytorch 将数据集拆分为自定义训练集和自定义验证集？
【发布时间】：2023-04-02 16:36:01
【问题描述】：

我使用的是非 Torchvision 数据集，并且我使用 ImageFolder 方法提取了它。我正在尝试将数据集拆分为 20% 的验证集和 80% 的训练集。我只能从允许拆分数据集的 PyTorch 库中找到这种方法（random_split）。但是，这每次都是随机的。我想知道有没有办法在 PyTorch 库中将数据集拆分为特定数量？

这是我提取数据集并随机拆分的代码。

transformations = transforms.Compose([
    transforms.Resize(255),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

TrafficSignSet = datasets.ImageFolder(root='./train/', transform=transformations)

####### split data
train_size = int(0.8 * len(TrafficSignSet))
test_size = len(TrafficSignSet) - train_size
train_dataset_split, test_dataset_split = torch.utils.data.random_split(TrafficSignSet, [train_size, test_size])

#######put into a Dataloader
train_dataset = torch.utils.data.DataLoader(train_dataset_split, batch_size=32, shuffle=True)
test_dataset = torch.utils.data.DataLoader(test_dataset_split, batch_size=32, shuffle=True)

【问题讨论】：

标签： python machine-learning neural-network pytorch

【解决方案1】：

如果您查看random_split 的“幕后”，您会发现它使用torch.utils.data.Subset 进行实际拆分。您可以使用固定索引自己执行此操作：

import random

indices = list(range(len(TrafficSignSet))
random.seed(310)  # fix the seed so the shuffle will be the same everytime
random.shuffle(indices)
train_dataset_split = torch.utils.data.Subset(TrafficSignSet, indices[:train_size])
val_dataset_split = torch.utils.data.Subset(TrafficSignSet, indices[train_size:])

【讨论】：