R: How to split a data frame into training, validation, and test sets?

Question

I'm using R to do machine learning. Following standard machine learning methodology, I would like to randomly split my data into training, validation, and test data sets. How do I do that in R?

I know there are some related questions on how to split into 2 data sets (e.g. this post), but it is not obvious how to do it for 3 split data sets. By the way, the correct approach is to use 3 data sets (including a validation set to tune your hyperparameters).

Eal

Answer 1

This linked approach for two groups (using floor) doesn't extend naturally to three. I'd do

spec = c(train = .6, test = .2, validate = .2)

g = sample(cut(
  seq(nrow(df)), 
  nrow(df)*cumsum(c(0,spec)),
  labels = names(spec)
))

res = split(df, g)

To check the results:

sapply(res, nrow)/nrow(df)
#    train     test validate 
#  0.59375  0.18750  0.21875 
# or...
addmargins(prop.table(table(g)))
#    train     test validate      Sum 
#  0.59375  0.18750  0.21875  1.00000 

With set.seed(1) run just before, the result looks like

$train
                   mpg cyl  disp  hp drat    wt  qsec vs am gear carb
Mazda RX4         21.0   6 160.0 110 3.90 2.620 16.46  0  1    4    4
Mazda RX4 Wag     21.0   6 160.0 110 3.90 2.875 17.02  0  1    4    4
Datsun 710        22.8   4 108.0  93 3.85 2.320 18.61  1  1    4    1
Hornet Sportabout 18.7   8 360.0 175 3.15 3.440 17.02  0  0    3    2
Merc 240D         24.4   4 146.7  62 3.69 3.190 20.00  1  0    4    2
Merc 230          22.8   4 140.8  95 3.92 3.150 22.90  1  0    4    2
Merc 280          19.2   6 167.6 123 3.92 3.440 18.30  1  0    4    4
Merc 280C         17.8   6 167.6 123 3.92 3.440 18.90  1  0    4    4
Merc 450SE        16.4   8 275.8 180 3.07 4.070 17.40  0  0    3    3
Merc 450SL        17.3   8 275.8 180 3.07 3.730 17.60  0  0    3    3
Merc 450SLC       15.2   8 275.8 180 3.07 3.780 18.00  0  0    3    3
Fiat 128          32.4   4  78.7  66 4.08 2.200 19.47  1  1    4    1
Toyota Corolla    33.9   4  71.1  65 4.22 1.835 19.90  1  1    4    1
Dodge Challenger  15.5   8 318.0 150 2.76 3.520 16.87  0  0    3    2
AMC Javelin       15.2   8 304.0 150 3.15 3.435 17.30  0  0    3    2
Pontiac Firebird  19.2   8 400.0 175 3.08 3.845 17.05  0  0    3    2
Fiat X1-9         27.3   4  79.0  66 4.08 1.935 18.90  1  1    4    1
Porsche 914-2     26.0   4 120.3  91 4.43 2.140 16.70  0  1    5    2
Volvo 142E        21.4   4 121.0 109 4.11 2.780 18.60  1  1    4    2

$test
                    mpg cyl  disp  hp drat    wt  qsec vs am gear carb
Valiant            18.1   6 225.0 105 2.76 3.460 20.22  1  0    3    1
Cadillac Fleetwood 10.4   8 472.0 205 2.93 5.250 17.98  0  0    3    4
Toyota Corona      21.5   4 120.1  97 3.70 2.465 20.01  1  0    3    1
Camaro Z28         13.3   8 350.0 245 3.73 3.840 15.41  0  0    3    4
Ford Pantera L     15.8   8 351.0 264 4.22 3.170 14.50  0  1    5    4
Ferrari Dino       19.7   6 145.0 175 3.62 2.770 15.50  0  1    5    6

$validate
                     mpg cyl  disp  hp drat    wt  qsec vs am gear carb
Hornet 4 Drive      21.4   6 258.0 110 3.08 3.215 19.44  1  0    3    1
Duster 360          14.3   8 360.0 245 3.21 3.570 15.84  0  0    3    4
Lincoln Continental 10.4   8 460.0 215 3.00 5.424 17.82  0  0    3    4
Chrysler Imperial   14.7   8 440.0 230 3.23 5.345 17.42  0  0    3    4
Honda Civic         30.4   4  75.7  52 4.93 1.615 18.52  1  1    4    2
Lotus Europa        30.4   4  95.1 113 3.77 1.513 16.90  1  1    5    2
Maserati Bora       15.0   8 301.0 335 3.54 3.570 14.60  0  1    5    8

Data.frames can be accessed like res$test or res[["test"]].

cut is the standard tool for partitioning based on shares.

Shallot

This is nice because all the rows will always be used. Multiple uses of floor() make it possible for some rows to get lost. And split returning a list is of course very nice. – Exhibitioner

This is an outstanding answer and has been much easier to use and test vs. others. – Hemielytron

Answer 2

Following the approach shown in this post, here is working R code to divide a dataframe into three new dataframes for testing, validation, and test. The three subsets are non-overlapping.

# Create random training, validation, and test sets

# Set some input variables to define the splitting.
# Input 1. The data frame that you want to split into training, validation, and test.
df <- mtcars

# Input 2. Set the fractions of the dataframe you want to split into training, 
# validation, and test.
fractionTraining   <- 0.60
fractionValidation <- 0.20
fractionTest       <- 0.20

# Compute sample sizes.
sampleSizeTraining   <- floor(fractionTraining   * nrow(df))
sampleSizeValidation <- floor(fractionValidation * nrow(df))
sampleSizeTest       <- floor(fractionTest       * nrow(df))

# Create the randomly-sampled indices for the dataframe. Use setdiff() to
# avoid overlapping subsets of indices.
indicesTraining    <- sort(sample(seq_len(nrow(df)), size=sampleSizeTraining))
indicesNotTraining <- setdiff(seq_len(nrow(df)), indicesTraining)
indicesValidation  <- sort(sample(indicesNotTraining, size=sampleSizeValidation))
indicesTest        <- setdiff(indicesNotTraining, indicesValidation)

# Finally, output the three dataframes for training, validation and test.
dfTraining   <- df[indicesTraining, ]
dfValidation <- df[indicesValidation, ]
dfTest       <- df[indicesTest, ]

Eal

Thanks. Works perfectly! – Homoio

Can I use this to split a dataframe with 3 groups of the dependent variable? – Thresher

Answer 3

Some of these seem overly complex, here's a simple way using sample to split any dataset into 3, or even an arbitrary number of sets.

# Simple into 3 sets.
idx <- sample(seq(1, 3), size = nrow(iris), replace = TRUE, prob = c(.8, .2, .2))
train <- iris[idx == 1,]
test <- iris[idx == 2,]
cal <- iris[idx == 3,]

If you'd rather reusable code:

# Or a function to split into arbitrary number of sets
test_split <- function(df, cuts, prob, ...)
{
  idx <- sample(seq(1, cuts), size = nrow(df), replace = TRUE, prob = prob, ...)
  z = list()
  for (i in 1:cuts)
    z[[i]] <- df[idx == i,]
  z
}
z <- test_split(iris, 4, c(0.7, .1, .1, .1))

train <- z[1]
test <- z[2]
cal <- z[3]
other <- z[4]

Rocha

This does not guarantee the sizes of the subsets, since membership is independent across observations. In particular, a subset could end up entirely empty. – Shallot

Answer 4

Here is one solution with a 60, 20 , 20 split that also ensures that there is no overlapping. However it is a trouble to adapt the split. If anyone could help me out, I appreciate it

   # Draw a random, stratified sample including p percent of the data    
   idx.train <- createDataPartition(y = known$return_customer, p = 0.8, list = FALSE) 
   train <- known[idx.train, ] # training set with p = 0.8
   # test set with p = 0.2 (drop all observations with train indeces)
   test <-  known[-idx.train, ] 
   idx.validation <- createDataPartition(y = train$return_customer, p = 0.25, list = FALSE) # Draw a random, stratified sample of ratio p of the data
   validation <- train[idx.validation, ] #validation set with p = 0.8*0.25 = 0.2
   train60 <- train[-idx.validation, ] #final train set with p= 0.8*0.75 = 0.6

Cooncan

I don't think this is a very helpful answer, considering no one except you can run the code. Maybe you should post it as a question after reading some guidance on writing good R questions for this site: #5963769 – Shallot

you're right! I posted it here #41880953 – Cooncan

Ok, thanks. You'll probably want to delete this answer then. – Shallot

Answer 5

Caret also support data splitting with the function createDataPartition

if your outcome y is Unbalanced factor ( yes >>> No and vice versa), ideally the random sampling occurs within each class and should preserve the overall class distribution of the data. which is the case with createDataPartition

Example:

library(caret)
set.seed(123)
table(iris$Species=="setosa")

## 
## FALSE  TRUE 
##   100    50

Note our outcome is unbalanced

Splitting (80% train and 20% test):

trainIndex <- createDataPartition(iris$Species=="setosa", p = .8, 
                                  list = FALSE, 
                                  times = 1)
train = iris[ trainIndex,]
test = iris[-trainIndex,]

Verification:

table(train$Species == "setosa")

## 
## FALSE  TRUE 
##    80    40

table(test$Species == "setosa")
## 
## FALSE  TRUE 
##    20    10

Note we preserve the overall class distribution

Wellesley

Answer 6

I think my approach is the easiest one:

idxTrain <- sample(nrow(dat),as.integer(nrow(dat)*0.7))
idxNotTrain <- which(! 1:nrow(dat) %in% idxTrain )
idxVal <- sample(idxNotTrain,as.integer(length(idxNotTrain)*0.333))
idxTest <- idxNotTrain[which(! idxNotTrain %in% idxVal)]

First, it splits the data into 70% training data and the rest (idxNotTrain). Then, the rest is again splitted into a validation data set (33%, 10% of the total data) and the rest (the testing data, 66%, 20% of the total data).

Regardful

Answer 7

Let me know if this would work. Just a simplified version

sample_train<- sample(seq_len(nrow(mtcars)), size = floor(0.60*nrow(mtcars)))
sample_valid<- sample(seq_len(nrow(mtcars)), size = floor(0.20*nrow(mtcars)))
sample_test <- sample(seq_len(nrow(mtcars)), size = floor(0.20*nrow(mtcars)))

train     <- mtcars[sample_train, ]
validation<- mtcars[sample_valid, ]
test      <- mtcars[sample_test, ]

Markettamarkey