Pytorch:如何获取子集的所有数据和目标

2024-03-28 10:49:13 发布

您现在位置:Python中文网/ 问答频道 /正文

我使用以下代码从特定文件夹读取数据集,并将其划分为训练和测试子集。我可以使用列表理解获得每个子集的所有数据和目标,但对于大数据来说,这非常缓慢。 有没有其他快速方法可以做到这一点

def train_test_dataset(dataset, test_split=0.20):
    train_idx, test_idx = train_test_split(list(range(len(dataset))), test_size=test_split, stratify=dataset.targets)
    datasets = {}
    train_dataset = Subset(dataset, train_idx)
    test_dataset = Subset(dataset, test_idx)

    return train_dataset, test_dataset


dataset = dset.ImageFolder("/path_to_folder", transform = transform)
    
train_set, test_set = train_test_dataset(dataset)

train_data = [data for data, _ in train_set]
train_labels = [label for _, label in train_set]

我使用DataLoader尝试过这种方法,它更好,但也需要一些时间: PyTorch Datasets: Converting entire Dataset to NumPy

多谢各位


Tags: to数据方法intestfordatatransform
1条回答
网友
1楼 · 发布于 2024-03-28 10:49:13

您提供的链接中的answer基本上与拥有数据加载器的目的背道而驰:数据加载器意味着将数据逐块加载到内存中。这有一个明显的优点,即不必在给定时刻加载整个数据集

^{}数据集中,可以使用^{}函数拆分数据:

>>> def train_test_dataset(dataset, test_split=.2):
...    test_len = int(len(dataset)*test_split)
...    train_len = len(dataset) - test_len 
...    return random_split(dataset, [train_len, test_len])

然后,您可以将这些数据集插入到单独的^{}中:

>>> train_set, test_set = train_test_dataset(dataset)

>>> train_dl = DataLoader(train_set, batch_size=16, shuffle=True)
>>> test_dl  = DataLoader(train_set, batch_size=32 shuffle=False)

相关问题 更多 >