XGBoost: split data in train and test

Question

I am using the python interface for XGBoost for building models. I have a dataset that I am reading in using xgb.DMatrix(data_path). I need to split this data into train and test (and validation, if required). But most of the implementations I have seen are of the form

dtrain = xgb.DMatrix('')
dtest = xgb.DMatrix('')

I couldn't find a way to where we can read in the dataset and then split 'em into train, test (and validation) sets.

Furthermore, is it possible to perform stratified sampling while splitting into train and test?

I need to know this because I have slightly larger datasets and currently I am reading it in using spark, splitting them up, storing on disk and then reading from there. Is there a way I can do it without having to go through Pyspark and reading from the hdfs?

Answer 1

I would use sklearn's train_test_split, which also has a stratify parameter, and then put the results into dtrain and dtest.

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

import xgboost as xgb
dtrain = xgb.DMatrix(X_train, label=y_train)
dtest = xgb.DMatrix(X_test, label=y_test)

See implementation here: A Simple XGBoost Tutorial Using the Iris Dataset.