Skip to content

Sklearn train_test_split: Dividir datos para Machine Learning en Python

Soren Atelier
Name
Soren Atelier

Updated on

Entrenar un modelo de machine learning con los mismos datos que usas para evaluarlo produce una precisión engañosamente alta. El modelo memoriza los datos de entrenamiento en lugar de aprender patrones generalizables -- un problema llamado sobreajuste (overfitting). Necesitas un conjunto de prueba separado que el modelo nunca vea durante el entrenamiento para obtener métricas de rendimiento honestas.

La función train_test_split() de Scikit-learn es la forma estándar de dividir conjuntos de datos en porciones de entrenamiento y prueba. Maneja arrays, DataFrames y matrices dispersas, con opciones para estratificación, reproducibilidad y ratios de división personalizados.

📚

    Uso básico

    from sklearn.model_selection import train_test_split
import numpy as np
 
# Datos de ejemplo: 100 muestras, 5 características
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)  # Etiquetas binarias
 
# División: 80% entrenamiento, 20% prueba
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
 
print(f"Set de entrenamiento: {X_train.shape[0]} muestras")
print(f"Set de prueba:        {X_test.shape[0]} muestras")
# Set de entrenamiento: 80 muestras
# Set de prueba:        20 muestras

    Parámetros clave

    ParámetroValor predeterminadoDescripción
    test_size0.25Fracción (0.0-1.0) o número absoluto de muestras de prueba
    train_sizeNoneFracción o número de muestras de entrenamiento (complemento de test_size)
    random_stateNoneSemilla para divisiones reproducibles
    shuffleTrueSi se barajan los datos antes de dividir
    stratifyNoneArray para división estratificada

    Con Pandas DataFrames

    from sklearn.model_selection import train_test_split
import pandas as pd
 
df = pd.DataFrame({
    'age': [25, 30, 35, 40, 45, 50, 55, 60, 28, 33],
    'income': [40, 50, 60, 70, 80, 90, 100, 110, 45, 55],
    'purchased': [0, 0, 1, 1, 1, 1, 1, 1, 0, 0]
})
 
X = df[['age', 'income']]
y = df['purchased']
 
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42
)
 
print(f"Train shape: {X_train.shape}")  # (7, 2)
print(f"Test shape:  {X_test.shape}")   # (3, 2)
print(type(X_train))  # <class 'pandas.core.frame.DataFrame'>

    Los DataFrames siguen siendo DataFrames después de la división -- los nombres de columnas e índices se preservan.

    random_state: Divisiones reproducibles

    Sin random_state, obtienes una división diferente cada vez:

    from sklearn.model_selection import train_test_split
import numpy as np
 
X = np.arange(10).reshape(5, 2)
y = np.array([0, 0, 1, 1, 1])
 
# Sin random_state: división diferente en cada ejecución
_, X_test1, _, _ = train_test_split(X, y, test_size=0.4)
_, X_test2, _, _ = train_test_split(X, y, test_size=0.4)
print(np.array_equal(X_test1, X_test2))  # Probablemente False
 
# Con random_state: misma división cada vez
_, X_test3, _, _ = train_test_split(X, y, test_size=0.4, random_state=42)
_, X_test4, _, _ = train_test_split(X, y, test_size=0.4, random_state=42)
print(np.array_equal(X_test3, X_test4))  # True

    Siempre establece random_state para reproducibilidad. Usa cualquier entero -- 42 es convencional, pero el número específico no importa.

    División estratificada

    Para conjuntos de datos desbalanceados, una división aleatoria podría poner la mayoría de las muestras minoritarias en un set. La estratificación asegura que ambos sets tengan las mismas proporciones de clase:

    from sklearn.model_selection import train_test_split
import numpy as np
from collections import Counter
 
# Conjunto de datos desbalanceado: 90% clase 0, 10% clase 1
np.random.seed(42)
X = np.random.randn(200, 4)
y = np.array([0] * 180 + [1] * 20)
 
# Sin estratificación
_, _, y_train_bad, y_test_bad = train_test_split(
    X, y, test_size=0.2, random_state=42
)
print("Sin estratificación:")
print(f"  Train: {Counter(y_train_bad)}")
print(f"  Test:  {Counter(y_test_bad)}")
 
# Con estratificación
_, _, y_train_good, y_test_good = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print("\nCon stratify=y:")
print(f"  Train: {Counter(y_train_good)}")
print(f"  Test:  {Counter(y_test_good)}")
# Train: Counter({0: 144, 1: 16})  -- 10% clase 1
# Test:  Counter({0: 36, 1: 4})    -- 10% clase 1

    Cuándo usar estratificación

    Escenario¿Usar stratify?
    Clases balanceadas (50/50)Opcional
    Clases desbalanceadas (90/10)Sí
    Clasificación multiclaseSí
    Regresión (objetivo continuo)No (no soportado)
    Conjuntos pequeños (< 100 muestras)Sí (previene clases vacías)

    División Train/Validación/Test

    Para el ajuste de hiperparámetros, necesitas tres conjuntos: entrenamiento, validación y prueba. Aplica train_test_split dos veces:

    from sklearn.model_selection import train_test_split
import numpy as np
 
X = np.random.randn(1000, 10)
y = np.random.randint(0, 3, 1000)
 
# Primera división: 80% train+val, 20% test
X_temp, X_test, y_temp, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# Segunda división: 75% train, 25% val (del 80% = 60/20/20 en total)
X_train, X_val, y_train, y_val = train_test_split(
    X_temp, y_temp, test_size=0.25, random_state=42, stratify=y_temp
)
 
print(f"Entrenamiento: {X_train.shape[0]} muestras (60%)")
print(f"Validación:    {X_val.shape[0]} muestras (20%)")
print(f"Prueba:        {X_test.shape[0]} muestras (20%)")

    Ejemplo completo de pipeline ML

    from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report
import numpy as np
 
# Generar datos de ejemplo
np.random.seed(42)
X = np.random.randn(500, 8)
y = (X[:, 0] + X[:, 1] > 0).astype(int)
 
# Dividir datos
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# Escalar características (¡ajustar solo con datos de entrenamiento!)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)  # Usar estadísticas del entrenamiento
 
# Entrenar modelo
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train_scaled, y_train)
 
# Evaluar
y_pred = model.predict(X_test_scaled)
print(f"Precisión: {accuracy_score(y_test, y_pred):.4f}")
print(classification_report(y_test, y_pred))

    Regla crítica: Ajusta el scaler (y cualquier preprocesamiento) solo con el set de entrenamiento. Aplica la misma transformación al set de prueba. Ajustar con el conjunto completo causa fuga de datos (data leakage).

    Ratios de división comunes

    DivisiónEntrenamientoPruebaCuándo usar
    80/2080%20%Opción predeterminada, la mayoría de conjuntos
    70/3070%30%Conjuntos pequeños, se necesita mayor set de prueba
    90/1090%10%Conjuntos grandes (10k+ muestras)
    60/20/2060%20% val + 20% testAl ajustar hiperparámetros

    Explorar resultados del modelo

    Después de dividir y entrenar tu modelo, PyGWalker (opens in a new tab) te permite explorar interactivamente predicciones vs valores reales, distribuciones de características entre sets de train/test y patrones de error en Jupyter:

    import pandas as pd
import pygwalker as pyg
 
results = pd.DataFrame({
    'actual': y_test,
    'predicted': y_pred,
    'correct': y_test == y_pred
})
walker = pyg.walk(results)

    FAQ

    ¿Qué hace train_test_split en sklearn?

    train_test_split() divide aleatoriamente arrays o DataFrames en dos subconjuntos: uno para entrenamiento y otro para prueba. Asegura que la evaluación del modelo use datos que no ha visto durante el entrenamiento, proporcionando estimaciones de rendimiento honestas.

    ¿Cuál es el mejor ratio de división train/test?

    80/20 es el valor predeterminado estándar. Usa 70/30 para conjuntos más pequeños donde necesitas un set de prueba confiable, y 90/10 para conjuntos grandes (10k+ muestras) donde el 10% sigue siendo sustancial. Para ajuste de hiperparámetros, usa 60/20/20 (train/val/test).

    ¿Qué hace random_state en train_test_split?

    random_state establece la semilla aleatoria para el barajado que ocurre antes de la división. Usar el mismo valor de random_state produce la misma división cada vez, haciendo tus resultados reproducibles. Cualquier entero funciona.

    ¿Cuándo debo usar stratify en train_test_split?

    Usa stratify=y cuando tu variable objetivo esté desbalanceada (ej., 95% negativo, 5% positivo) o cuando tengas un conjunto de datos pequeño. La estratificación asegura que tanto el set de entrenamiento como el de prueba tengan la misma proporción de cada clase.

    ¿Cómo divido datos en sets de entrenamiento, validación y prueba?

    Llama a train_test_split dos veces. Primero divide en train+val y test (ej., 80/20). Luego divide train+val en train y val (ej., 75/25 del 80%, dando 60/20/20 en total). Alternativamente, usa sklearn.model_selection.KFold para validación cruzada.

    Conclusión

    train_test_split() es la base de todo flujo de trabajo de machine learning en Python. Úsalo siempre antes de entrenar -- nunca evalúes con datos de entrenamiento. Establece random_state para reproducibilidad, usa stratify=y para clases desbalanceadas, y recuerda ajustar los pasos de preprocesamiento solo con el set de entrenamiento. Para selección de modelos, divide en tres partes (train/validación/test) o usa validación cruzada para estimaciones más robustas.

    📚