Train Test Split in Python: train_test_split

Q: What is the default test size?

If neither test_size nor train_size is set, test_size is set to 0.25.

Q: Does train_test_split return the same split every time?

Only if you pass an integer random_state . With the default None, the shuffling can differ between calls.

A test train split in Python (also written train test split in Python) divides one dataset into a training part, used to fit a model, and a held-back test part, used once to estimate how the model performs on data it has not seen. In Python the standard tool is train_test_split from scikit-learn. This guide covers the function, how its parameters interact, the order of operations that keeps test information out of the model, and what to write in a methods section so the split can be reproduced.

What it is used for in research

A model scored on the same data it was fitted on says little about how it will behave on new data. Holding out part of the data for testing gives an estimate that is not tied to the fitted sample. scikit-learn’s cross-validation documentation describes the pattern directly: hold out a test set, and call train_test_split to do it (its example holds out 40% of the data with test_size=0.4, random_state=0).

The split is a research-design decision, not just a coding step. How much is held out, whether the split is shuffled, whether class proportions are preserved and which seed was used all change the reported number, so each belongs in the methods description. For the wider workflow around the function, see the CASRAI scikit-learn tutorial, and for the field context, what is machine learning.

Install or access

train_test_split is part of scikit-learn, not SciPy or statsmodels. Install scikit-learn from the project’s official channel (scikit-learn installation page), for example pip install -U scikit-learn, then import the function:

from sklearn.model_selection import train_test_split

Record the scikit-learn version you ran. The documentation notes version-specific behaviour (for example, since version 0.16 sparse input is returned as scipy.sparse.csr_matrix), so the version is part of the method. For a reproducible environment, see sharing a reproducible conda environment.

Step-by-step task

Train test split in Python with train_test_split

The signature in the current scikit-learn documentation is:

sklearn.model_selection.train_test_split(*arrays, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None)

Put your features and labels in arrays, lists, dataframes or scipy sparse matrices. All inputs must have a consistent first dimension (the same number of rows).
Split before any preprocessing, feature selection or tuning. This ordering is covered in the next section.

Call the function with a fixed integer seed:

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)

Fit on X_train, y_train only. Use X_test, y_test for the final evaluation.

The documentation’s own example uses X, y = np.arange(10).reshape((5, 2)), range(5) with the call above and returns y_train as [2, 0, 3] and y_test as [1, 4]. Calling train_test_split(y, shuffle=False) on the same y returns [[0, 1, 2], [3, 4]]: without shuffling, the first rows become the training set and the last rows the test set.

Test train split in Python: the same function

“Test train split” and “train test split” refer to the same operation. The function always returns the training part first and the test part second for each input array, so the four-variable unpacking above is the same whichever phrase you search for.

Output and how to read it

The function returns a list of length 2 * len(arrays), in the order train, test for each array you passed. Pass two arrays and you get four items; pass one and you get two. Output types follow the input, except that sparse input is returned as scipy.sparse.csr_matrix.

Parameter	What the documentation says	Interaction to check
`test_size`	Float between 0.0 and 1.0 is a proportion of the dataset; int is an absolute number of test samples.	If None, it is the complement of the train size. If `train_size` is also None, it is set to 0.25.
`train_size`	Float is a proportion; int is an absolute count.	If None, it is the complement of the test size.
`random_state`	Controls the shuffling applied before the split. Pass an int for reproducible output across calls.	Default None, so repeated runs can give different splits.
`shuffle`	Default True.	If `shuffle=False` then `stratify` must be None.
`stratify`	If not None, data is split in a stratified fashion, using this as the class labels.	Cannot be combined with `shuffle=False`.

Because the defaults leave test_size at 0.25 when neither size is given, a call with no size arguments is a 75/25 split. State the proportion explicitly in code and in the paper rather than relying on the default.

Split first, then fit anything that learns from the data

The scikit-learn common-pitfalls page defines data leakage as information that would not be available at prediction time being used when building the model. Its rules are short: never call fit on the test data; always split into train and test subsets first, particularly before any preprocessing steps; and do feature selection using training data only. Using all the data, for example fit(X) on the full dataset, can give overly optimistic scores. The same idea is defined in the CASRAI dictionary entry data leakage (training).

Reporting the split so it can be reproduced

For a methods section, report the scikit-learn version, the split proportions, whether shuffling and stratification were used, and the seed. On seeds, the scikit-learn pitfalls page says reproducible results across executions require removing any use of random_state=None, that an integer random_state makes repeated calls give the same result, and that it recommends declaring an rng variable at the top and passing it down. It does not recommend setting the global numpy seed with np.random.seed(0). These details support what the CASRAI dictionary calls computational reproducibility.

Troubleshooting

The split changes every time I run the script

random_state defaults to None. Pass an integer, as in the example above.

stratify is rejected with shuffle=False

The documentation requires stratify to be None when shuffle=False. Either allow shuffling or drop stratification.

Inputs have different lengths

The function expects indexable inputs with a consistent first dimension. Check that features and labels have the same number of rows.

Test scores look too good

Check for leakage: any fit, scaling, imputation or feature selection run on the full dataset before the split lets test information into the model. Move the split to the first step.

Test scores stopped meaning anything after tuning

The cross-validation documentation warns that tweaking hyperparameters against the test set risks overfitting on it: knowledge about the test set can leak into the model and the metrics no longer report generalization performance.

Free alternatives and related approaches

A single hold-out is not the only design the scikit-learn documentation describes. The cross-validation documentation describes these designs.

Approach	What the documentation says
Single hold-out with `train_test_split`	One call holds out a test set; the example holds out 40% with `test_size=0.4, random_state=0`.
Train, validation and test sets	Training proceeds on the training set, evaluation on the validation set, and when the experiment seems successful, final evaluation is done on the test set. The three-way partition drastically reduces the number of samples available for learning, and results can depend on the particular random choice.
Stratified split	Set `stratify` to the class labels so the split is stratified.

For cross-validation as a way around the cost of a separate validation set, see the cross-validation section of the CASRAI scikit-learn tutorial and the related elastic net regression guide, which covers cross-validation.