Cómo dividir un conjunto de datos en dos partes (train/test split) en Python

Una operación que es común en todos los modelos de aprendizaje supervisado es la división de nuestro conjunto de datos en -al menos- dos partes: una parte Train, de entrenamiento, que corresponderá a la mayor parte de nuestro dataset y que usaremos para entrenar nuestro modelo y un parte Test, de menor tamaño, sobre la que evaluaremos nuestro modelo entrenado.

Imagen vía Towards Data Science.

Esta división se puede realizar de forma aleatoria o lineal, y siempre elegimos qué porcentaje queremos para cada división. Vamos a ver cómo hacer esta división “manualmente” o usando la función train_test_split() de scikit-learn.

0. Dataset de ejemplo

Utilizaremos un sencillo y pequeño dataset con datos aleatorios para mostrar los resultados de los ejemplos utilizados.

import pandas as pd

df = pd.read_excel("ejemplo.xlsx")

Simplemente cargamos en una estructura pandas nuestro dataframe y los visualizamos:

Dataframe usado para los ejemplos.

Vemos que se trata de un pequeño dataset de 4 variables y 10 observaciones. Así que nos ponemos manos a la obra.

1. División lineal del conjunto de datos

Aunque la mayoría de las veces la división del conjunto de datos se realiza de forma aleatoria, en ciertas ocasiones nos interesa realizar una división lineal. Esto es especialmente recomendable cuando nuestros datos tienen un orden temporal (por ejemplo el valor diario de las acciones de cierta empresa) y no tendría sentido usar datos futuros para predecir valores pasados, sino más bien lo contrario.

Con la división lineal, usamos la primera parte del dataset para entrenar y la última parte (la más reciente si son datos temporales) para el test. Lo haríamos de la siguiente forma:

p_train = 0.80 # Porcentaje de train.

train = df[:int((len(df))p_train)] 
test = df[int((len(df))p_train):]

print("Ejemplos usados para entrenar: ", len(train))
print("Ejemplos usados para test: ", len(test))

A la variable p_train le asignamos el porcentaje que queremos para nuestro conjunto de entrenamiento (80% en el ejemplo) y simplemente guardamos en train y test la parte correspondiente.

Resultado del train/test split lineal.

Como se puede observar, hemos seleccionado los primeros 8 registros (80%) como train y los dos últimos registros (20%) como test.

2. División aleatoria del conjunto de datos

Si nuestro dataset no tiene un orden temporal, qué será lo más normal, es conveniente realizar la división de forma aleatoria para asegurarnos la máxima variabilidad de los datos en nuestros conjuntos de train/test. Lo podemos hacer de la siguiente forma:

import numpy as np

p_train = 0.70 # Porcentaje de train.

df['is_train'] = np.random.uniform(0, 1, len(df)) <= p_train
train, test = df[df['is_train']==True], df[df['is_train']==False]
df = df.drop('is_train', 1)

print("Ejemplos usados para entrenar: ", len(train))
print("Ejemplos usados para test: ", len(test))

En este caso, hemos seleccionado un 70% como porcentaje para el train. Creamos una nueva columna en df que se llamará is_train y que tendrá un valor True en el 70% de los casos y un valor de False en el 30% restante. De esta forma realizamos nuestro Train/Test split:

Resultado de de train/test split aleatorio.

Vemos como la división aleatoria ha escogido los registros 1, 7 y 9 como Test (30%) y el resto como Train (70%). Para obtener siempre la misma división para este conjunto de datos, debemos establecer una semilla al inicio de nuestro código con random.seed().

3. División aleatoria con train_test_split()

Hasta ahora hemos realizado las divisiones de nuestro conjunto de datos de forma “manual”. Ahora utilizaremos una útil herramienta que nos proporciona la librería Model Selection de scikit-learn: train_test_split().

Cargamos la librería:

from sklearn.model_selection import train_test_split 

Ejecutamos la división:

train, test = train_test_split(df, test_size = 0.30)

print("Ejemplos usados para entrenar: ", len(train))
print("Ejemplos usados para test: ", len(test))

Vemos como de una forma muy sencilla realizamos la división de nuestro dataframe usando la función train_test_split().

Resultado del train/test split aleatorio con SciKit Learn.

En este caso, las filas 6, 1 y 5 (30%) han sido seleccionadas aleatoriamente por train_test_split() como test y el resto (70%) como train.

Podríamos haber establecido la semilla aleatoria dentro de la misma función con el parámetro random_state de la siguiente forma:

train, test = train_test_split(df, test_size = 0.30, random_state = 100) 

4. División lineal con train_test_split()

train_test_split() nos permite también, de forma muy sencilla, realizar una división lineal de nuestros datos. Cargamos la librería:

 from sklearn.model_selection import train_test_split

Y ejecutamos el train/test split estableciendo el parámetro shuffle como False:

train, test = train_test_split(df, test_size = 0.30, shuffle = False)
Para saber más:
SciKit Learn Model Selection train_test_split()
Train/Test Split and Cross Validation in Python.
How to split your dataset to train and test datasets using SciKit Learn.

Deja una respuesta