如何将训练数据集拆分为训练、验证和测试数据集？答案

【问题标题】：How do I split the training dataset into training, validation and test datasets?如何将训练数据集拆分为训练、验证和测试数据集？
【发布时间】：2019-05-01 02:56:43
【问题描述】：

我有一个自定义图像数据集及其目标。我在 PyTorch 中创建了一个训练数据集。我想把它分成三部分：训练、验证和测试。我该怎么做？

【问题讨论】：

标签： machine-learning dataset conv-neural-network pytorch

【解决方案1】：

给定参数train_frac=0.8，此函数会将dataset拆分为80%、10%、10%：

import torch, itertools
from torch.utils.data import TensorDataset

def dataset_split(dataset, train_frac):
    '''
    param dataset:    Dataset object to be split
    param train_frac: Ratio of train set to whole dataset

    Randomly split dataset into a dictionary with keys, based on these ratios:
        'train': train_frac
        'valid': (1-split_frac) / 2
        'test': (1-split_frac) / 2
    '''
    assert split_frac >= 0 and split_frac <= 1, "Invalid training set fraction"

    length = len(dataset)

    # Use int to get the floor to favour allocation to the smaller valid and test sets    
    train_length = int(length * train_frac)
    valid_length = int((length - train_length) / 2)
    test_length  = length - train_length - valid_length

    dataset = random_split(dataset, (train_length, valid_length, test_length))
    dataset = {name: set for name, set in zip(('train', 'valid', 'test'), sets)}
    return dataset

【讨论】：

【解决方案2】：

拥有“主”数据集后，您可以使用data.Subset 对其进行拆分。
这是一个随机分割的例子

import torch
from torch.utils import data
import random

master = data.Dataset( ... )  # your "master" dataset
n = len(master)  # how many total elements you have
n_test = int( n * .05 )  # number of test/val elements
n_train = n - 2 * n_test
idx = list(range(n))  # indices to all elements
random.shuffle(idx)  # in-place shuffle the indices to facilitate random splitting
train_idx = idx[:n_train]
val_idx = idx[n_train:(n_train + n_test)]
test_idx = idx[(n_train + n_test):]

train_set = data.Subset(master, train_idx)
val_set = data.Subset(master, val_idx)
test_set = data.Subset(master, test_idx)

这也可以使用data.random_split实现：

train_set, val_set, test_set = data.random_split(master, (n_train, n_val, n_test))

【讨论】：