128
$\begingroup$

How could I randomly split a data matrix and the corresponding label vector into a X_train, X_test, X_val, y_train, y_test, y_val with scikit-learn?

As far as I know, sklearn.cross_validation.train_test_split is only capable of splitting into two not into three...

Improve this question
$\endgroup$

12 Answers 12

Active Oldest Votes
163
$\begingroup$

You could just use sklearn.model_selection.train_test_split twice. First to split to train, test and then split train again into validation and train. Something like this:

 X_train, X_test, y_train, y_test 
    = train_test_split(X, y, test_size=0.2, random_state=1)

 X_train, X_val, y_train, y_val 
    = train_test_split(X_train, y_train, test_size=0.25, random_state=1) # 0.25 x 0.8 = 0.2
Improve this answer
$\endgroup$
7
  • 4
    $\begingroup$ Yes, this works of course but I hoped for something more elegant ;) Never mind, I accept this answer. $\endgroup$ – Hendrik Nov 17 '16 at 8:10
  • 1
    $\begingroup$ I wanted to add that if you want to use the validation set to search for the best hyper-parameters you can do the following after the split: gist.github.com/albertotb/1bad123363b186267e3aeaa26610b54b $\endgroup$ – skd Jun 6 '18 at 16:34
  • 16
    $\begingroup$ So what is the final train, test, validation proportion in this example? Because on the second train_test_split , you are doing this over the previous 80/20 split. So your val is 20% of 80%. The split proportions aren't very straightforward in this way. $\endgroup$ – Monica Heddneck Jun 14 '18 at 19:22
  • 1
    $\begingroup$ I agree with @Monica Heddneck that the 64% train, 16% validation and 20% test splt could be clearer. It's an annoying inference you have to make with this solution. $\endgroup$ – Perry Jun 25 '19 at 8:00
  • 1
    $\begingroup$ if test_size is an integer number this function will take test_size number of elements for test, so you can pre-compute the number of elements in each subsets given your proportion and use these values to do a double split $\endgroup$ – GJCode Nov 10 '19 at 10:39
57
$\begingroup$

There is a great answer to this question over on SO that uses numpy and pandas.

The command (see the answer for the discussion):

train, validate, test = np.split(df.sample(frac=1), [int(.6*len(df)), int(.8*len(df))])

produces a 60%, 20%, 20% split for training, validation and test sets.

Improve this answer
$\endgroup$
4
  • 2
    $\begingroup$ I can see the .6 meaning 60%... but what does the .8 mean? $\endgroup$ – Tom Hale May 11 '19 at 5:02
  • 1
    $\begingroup$ @TomHale np.split will split at 60% of the length of the shuffled array, then 80% of length (which is an additional 20% of data), thus leaving a remaining 20% of the data. This is due to the definition of the function. You can test/play with: x = np.arange(10.0), followed by np.split(x, [ int(len(x)*0.6), int(len(x)*0.8)]) $\endgroup$ – 0_0 May 14 '19 at 13:35
  • 1
    $\begingroup$ This is fantastic, such a simple, straightforward method. I always tried shuffling the indexes, then selecting a first X%, a.s.o. Just great! $\endgroup$ – devplayer Mar 11 '20 at 11:24
  • 5
    $\begingroup$ Major benefit of train_test_split is stratification $\endgroup$ – Kalanos Oct 5 '20 at 1:16
16
$\begingroup$

Adding to @hh32's answer, while respecting any predefined proportions such as (75, 15, 10):

train_ratio = 0.75
validation_ratio = 0.15
test_ratio = 0.10

# train is now 75% of the entire data set
# the _junk suffix means that we drop that variable completely
x_train, x_test, y_train, y_test = train_test_split(dataX, dataY, test_size=1 - train_ratio)

# test is now 10% of the initial data set
# validation is now 15% of the initial data set
x_val, x_test, y_val, y_test = train_test_split(x_test, y_test, test_size=test_ratio/(test_ratio + validation_ratio)) 

print(x_train, x_val, x_test)
Improve this answer
$\endgroup$
2
  • $\begingroup$ I think this is the best answer and should be accepted. What do you mean by "# the _junk suffix means that we drop that variable completely" though? $\endgroup$ – PascalIv Jun 12 '20 at 7:52
  • $\begingroup$ And I think the shuffle argument should be set to False in the second call, simply because there is no reason to shuffle again. $\endgroup$ – PascalIv Jun 12 '20 at 8:01
9
$\begingroup$

You can use train_test_split twice. I think this is most straightforward.

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=1)
X_train, X_val, y_train, y_val = train_test_split(
    X_train, y_train, test_size=0.25, random_state=1)

In this way, train, val, test set will be 60%, 20%, 20% of the dataset respectively.

Improve this answer
$\endgroup$
0
5
$\begingroup$

Most often you will find yourself not splitting it once but in a first step you will split your data in a training and test set. Subsequently you will perform a parameter search incorporating more complex splittings like cross-validation with a 'split k-fold' or 'leave-one-out(LOO)' algorithm.

Improve this answer
$\endgroup$
3
$\begingroup$

Best answer above does not mention that by separating two times using train_test_split not changing partition sizes won`t give initially intended partition:

x_train, x_remain = train_test_split(x, test_size=(val_size + test_size))

Then the portion of validation and test sets in the x_remain change and could be counted as

new_test_size = np.around(test_size / (val_size + test_size), 2)
# To preserve (new_test_size + new_val_size) = 1.0 
new_val_size = 1.0 - new_test_size

x_val, x_test = train_test_split(x_remain, test_size=new_test_size)

In this occasion all initial partitions are saved.

Improve this answer
$\endgroup$
3
$\begingroup$

Extension of @hh32's answer with preserved ratios.

# Defines ratios, w.r.t. whole dataset.
ratio_train = 0.8
ratio_val = 0.1
ratio_test = 0.1

# Produces test split.
x_remaining, x_test, y_remaining, y_test = train_test_split(
    x, y, test_size=ratio_test)

# Adjusts val ratio, w.r.t. remaining dataset.
ratio_remaining = 1 - ratio_test
ratio_val_adjusted = ratio_val / ratio_remaining

# Produces train and val splits.
x_train, x_val, y_train, y_val = train_test_split(
    x_remaining, y_remaining, test_size=ratio_val_adjusted)

Since the remaining dataset is reduced after the first split, new ratios with respect to the reduced dataset must be calculated by solving the equation:

$ R_{remaining} \cdot R_{new} = R_{old}$

Improve this answer
$\endgroup$
1
  • $\begingroup$ This is a correct implementation! Thank you! @Jorge Barrios $\endgroup$ – amc Sep 10 '20 at 17:26
2
$\begingroup$

How about using numpy random choice

import numpy as np
from sklearn.datasets import load_iris

def ttv_split(X, y = None, train_size = .6, test_size = .2, validation_size = .2, random_state = 42):
    """
    Basic approach using np random choice
    """
    np.random.seed(random_state)
    X = pd.DataFrame(X, columns = ["col_" + str(i) for i in range(X.shape[1])])
    size = sum((train_size,test_size,validation_size))
    n_samples = X.shape[0]
    if  size != 1:
        return f"Size of the dataset must sum up to 100% instead: {size} correct and try again"
    else:
        split_series = np.random.choice(a = ["train","test","validation"], p = [train_size, test_size, validation_size], size = n_samples)
        split_series = pd.Series(split_series)
        
        
        X_train, X_test, X_validation = X.iloc[split_series[split_series == "train"].index,:], X.iloc[split_series[split_series == "test"].index,:], X.iloc[split_series[split_series == "validation"].index,:]
        
        if not y is None:
            y = pd.DataFrame(y,columns=["target"])
            
            y_train, y_test, y_validation = y.iloc[split_series[split_series == "train"].index,:], y.iloc[split_series[split_series == "test"].index,:], y.iloc[split_series[split_series == "validation"].index,:]
            
            return X_train,X_test,X_validation,y_train,y_test,y_validation
        else:
            return X_train,X_test,X_validation
            

X,y = load_iris(return_Xy = True)

X_train,X_test,X_validation,y_train,y_test,y_validation = ttv_split(X, y)
Improve this answer
$\endgroup$
1
$\begingroup$

Here's another approach (assumes equal three-way split):

# randomly shuffle the dataframe
df = df.reindex(np.random.permutation(df.index))

# how many records is one-third of the entire dataframe
third = int(len(df) / 3)

# Training set (the top third from the entire dataframe)
train = df[:third]

# Testing set (top half of the remainder two third of the dataframe)
test = df[third:][:third]

# Validation set (bottom one third)
valid = df[-third:]

This can be made more concise but I kept it verbose for explanation purposes.

Improve this answer
$\endgroup$
1
$\begingroup$

Given train_frac=0.8, this function creates a 80% / 10% / 10% split:

import sklearn

def data_split(examples, labels, train_frac, random_state=None):
    ''' https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html
    param data:       Data to be split
    param train_frac: Ratio of train set to whole dataset

    Randomly split dataset, based on these ratios:
        'train': train_frac
        'valid': (1-train_frac) / 2
        'test':  (1-train_frac) / 2

    Eg: passing train_frac=0.8 gives a 80% / 10% / 10% split
    '''

    assert train_frac >= 0 and train_frac <= 1, "Invalid training set fraction"

    X_train, X_tmp, Y_train, Y_tmp = sklearn.model_selection.train_test_split(
                                        examples, labels, train_size=train_frac, random_state=random_state)

    X_val, X_test, Y_val, Y_test   = sklearn.model_selection.train_test_split(
                                        X_tmp, Y_tmp, train_size=0.5, random_state=random_state)

    return X_train, X_val, X_test,  Y_train, Y_val, Y_test
Improve this answer
$\endgroup$
0
$\begingroup$ 
import numpy as np
import pandas as pd

#length of data 
N = 10
scale=2


#generated random data
X, y = np.arange(N*scale).reshape((N, scale)), np.arange(N)

#Works for pandas dataframe too
#You can download titanic.csv from here 
#https://github.com/fuwiak/faster_ds/blob/master/sample_data/titanic.csv

#df = pd.read_csv("titanic.csv", sep="\t")
#X=df[df.columns.difference(["Survived"])]
#y=df["Survived"]



def train_test_val(X, y, train_ratio, test_ratio, val_ratio):
    assert sum([train_ratio, test_ratio, val_ratio])==1.0, "wrong given ratio, all ratios have to sum to 1.0"
    assert X.shape[0]==len(y), "X and y shape mismatch"

    ind_train = int(round(X.shape[0]*train_ratio))
    ind_test = int(round(X.shape[0]*(train_ratio+test_ratio)))

    X_train = X[:ind_train]
    X_test = X[ind_train:ind_test]
    X_val = X[ind_test:]

    y_train = y[:ind_train]
    y_test = y[ind_train:ind_test]
    y_val = y[ind_test:]

    return X_train, X_test, X_val, y_train, y_test, y_val
# put ratio as you wish
X_train, X_test, X_val, y_train, y_test, y_val=train_test_val(X, y, 0.8, 0.1, 0.1)
Improve this answer
$\endgroup$
1
  • $\begingroup$ You do not randomize the choice of the training set / testing set. You just put a given share on the full dataset. The model will not learn from a representative dataset as soon as the dataset is not fully randomly distributed, which is likely in such datasets. $\endgroup$ – questionto42 May 22 at 21:04
0
$\begingroup$

Run it twice. Here is the math for the 2nd test_size.

Let's say I want {train:0.67, validation:0.13, test:0.20}

The first test_size is 20% which leaves 80% of the original data to be split into validation and training data.

(1.0/(1.0-test_size))*validation_size = second_test_size

# (1.0/(1.0-0.20))*0.13 = 0.1625

Also, look into the stratify parameter as that is the real reason to use train_test_split as opposed to selecting random row indices.

Improve this answer
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.