Cómo dividir datos en entrenamiento y prueba en Machine Learning

Dividir los datos en entrenamiento y prueba es uno de los primeros pasos de cualquier proyecto de Machine Learning: se entrena el modelo en un conjunto y se evalúa en otro que nunca ha visto, evitando resultados demasiado optimistas. A continuación verás cómo dividir datos en entrenamiento y prueba en Machine Learning con la función train_test_split de scikit-learn, con ejemplos sencillos que puedes adaptar a tu caso.

Requisitos previos

• Python 3 instalado (versión 3.9 o superior).
• Las bibliotecas scikit-learn y pandas instaladas (pip install scikit-learn pandas).
• Un conjunto de datos con columnas de entrada (features) y una columna objetivo (target).
• Conocimientos básicos de Python, como variables e importación de módulos.

Paso 1: Preparar los datos

Antes de dividir, separamos las features (las columnas que el modelo usa para predecir, guardadas en X) de la columna objetivo (lo que queremos predecir, guardado en y). Vamos a usar un ejemplo pequeño con pandas.

import pandas as pd

# Example dataset
data = pd.DataFrame({
    "study_hours": [1, 2, 3, 4, 5, 6, 7, 8],
    "absences":    [8, 6, 5, 4, 3, 2, 1, 0],
    "passed":      [0, 0, 0, 1, 1, 1, 1, 1]
})

X = data[["study_hours", "absences"]]  # features
y = data["passed"]                      # target

Por convención, X es una tabla con varias columnas e y es una sola columna. Mantener esta separación clara evita que la columna objetivo entre por error en las features — un fallo que dejaría que el modelo "vea" la respuesta.

Paso 2: Dividir con train_test_split

La función train_test_split del módulo sklearn.model_selection baraja los datos y los separa en dos conjuntos. Solo tienes que pasar X e y e indicar el porcentaje para prueba.

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,     # 20% for testing, 80% for training
    random_state=42    # makes the split reproducible
)

Fíjate en el orden de las variables a la izquierda: primero los dos conjuntos de X (entrenamiento y prueba) y solo después los dos de y. Cambiar este orden es un error común que provoca resultados extraños más adelante.

Paso 3: Entender test_size y random_state

Dos parámetros controlan el comportamiento de la división:

• test_size: la fracción de datos reservada para prueba. 0.2 significa 20%; como alternativa puedes indicar train_size.
• random_state: fija la semilla aleatoria. Usando siempre el mismo número (por ejemplo 42), obtienes exactamente la misma división en cada ejecución, haciendo que tus experimentos sean reproducibles.

Por defecto, la función baraja las filas antes de dividir (shuffle=True). Esto es importante porque muchos conjuntos de datos vienen ordenados — por ejemplo, con todos los casos positivos al final — y dividir sin barajar dejaría el conjunto de prueba sesgado.

Paso 4: Mantener las proporciones con stratify

En problemas de clasificación, entrenamiento y prueba deben tener la misma proporción de clases. Si el 30% de los casos son de la clase "aprobado", queremos ese 30% en ambos conjuntos. El parámetro stratify lo garantiza:

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    random_state=42,
    stratify=y         # keep class proportions of y
)

Consejo: usa stratify=y siempre que la variable objetivo sea categórica. En problemas de regresión (objetivo continuo), deja stratify fuera.

Verificar el resultado

Para confirmar que la división salió bien, comprueba el número de filas de cada conjunto con el atributo .shape y las proporciones de clases:

print("Train:", X_train.shape, "Test:", X_test.shape)
print(y_train.value_counts(normalize=True))

La suma de las filas de entrenamiento y prueba debe ser igual al total original, y las proporciones de clases deben quedar similares en ambos conjuntos. Como fijamos random_state, estos números se repiten cada vez que ejecutas el script.

Conclusión

Con pocas líneas dividiste los datos en entrenamiento y prueba — el paso que garantiza una evaluación honesta de un modelo de Machine Learning. A partir de aquí puedes entrenar un algoritmo con X_train e y_train y medir el rendimiento con X_test e y_test, por ejemplo con una matriz de confusión. Como siguiente paso, prueba a cambiar test_size a 0.3 y observa cómo cambia el tamaño de los conjuntos. ¿Cuál será el porcentaje de prueba ideal para tu problema?