Skip to content
Temas
Python
Sklearn Train Test Split: Guía Completa para Dividir Datos en Python

Sklearn Train Test Split: Guía Completa para Dividir Datos en Python

Updated on

Entrenar un modelo de machine learning en todo tu conjunto de datos y luego evaluarlo en los mismos datos conduce a un problema crítico: tu modelo parecerá funcionar bien, pero simplemente ha memorizado los datos en lugar de aprender patrones. Este sobreajuste significa que tu modelo fallará estrepitosamente cuando encuentre datos nuevos y no vistos. Los científicos de datos necesitan una forma confiable de evaluar el rendimiento del modelo en datos que el modelo nunca ha visto durante el entrenamiento.

La solución es la división train-test. Al reservar una porción de tus datos para evaluación, obtienes una evaluación honesta de cómo funcionará tu modelo en el mundo real. La función train_test_split de sklearn hace que este proceso sea sencillo, pero usarla incorrectamente aún puede conducir a fugas de datos, mala generalización y métricas de rendimiento engañosas.

Esta guía cubre todo lo que necesitas saber sobre train_test_split de sklearn, desde el uso básico hasta técnicas avanzadas para datos de series temporales, clases desbalanceadas y problemas de múltiples salidas.

📚

¿Qué es Train Test Split?

Train test split es la técnica fundamental para evaluar modelos de machine learning. Divides tu conjunto de datos en dos partes: un conjunto de entrenamiento usado para ajustar el modelo, y un conjunto de prueba usado para evaluar el rendimiento del modelo en datos no vistos.

La función train_test_split de scikit-learn (sklearn) automatiza este proceso, manejando el mezclado y la división aleatoria con solo una línea de código.

from sklearn.model_selection import train_test_split
 
# Uso básico
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

En este ejemplo, X contiene tus características (variables de entrada) e y contiene tu variable objetivo (lo que quieres predecir). La función devuelve cuatro arrays: características de entrenamiento, características de prueba, etiquetas de entrenamiento y etiquetas de prueba.

Sintaxis Básica de train_test_split

El uso más simple de train_test_split requiere solo dos argumentos: tus características y tu variable objetivo.

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
 
# Cargar datos de ejemplo
iris = load_iris()
X = iris.data
y = iris.target
 
# Dividir los datos (80% entrenamiento, 20% prueba)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
 
print(f"Training samples: {len(X_train)}")
print(f"Test samples: {len(X_test)}")

Esto divide tus datos aleatoriamente, con 80% yendo al entrenamiento y 20% a las pruebas. Sin embargo, este uso básico tiene un defecto crítico: la división es diferente cada vez que ejecutas el código, haciendo los resultados irreproducibles.

Parámetros Esenciales

test_size y train_size

El parámetro test_size controla cuántos datos van al conjunto de prueba. Puedes especificarlo como:

  • Un float entre 0.0 y 1.0 (proporción del conjunto de datos)
  • Un entero (número absoluto de muestras de prueba)
# Conjunto de prueba del 30%
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)
 
# 50 muestras en el conjunto de prueba
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=50)
 
# Alternativamente, especificar train_size
X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8)

Si especificas tanto test_size como train_size, deben sumar 1.0 (o el tamaño total del conjunto de datos si usas enteros). En la mayoría de los casos, especificar solo test_size es suficiente.

random_state para Reproducibilidad

El parámetro random_state es crucial para resultados reproducibles. Sin él, obtienes una división diferente cada vez que ejecutas tu código, haciendo imposible depurar o comparar experimentos.

# Sin random_state - división diferente cada vez
X_train1, X_test1, y_train1, y_test1 = train_test_split(X, y, test_size=0.2)
X_train2, X_test2, y_train2, y_test2 = train_test_split(X, y, test_size=0.2)
 
print(f"Same split? {(X_train1 == X_train2).all()}")  # False
 
# Con random_state - misma división cada vez
X_train1, X_test1, y_train1, y_test1 = train_test_split(X, y, test_size=0.2, random_state=42)
X_train2, X_test2, y_train2, y_test2 = train_test_split(X, y, test_size=0.2, random_state=42)
 
print(f"Same split? {(X_train1 == X_train2).all()}")  # True

Usa cualquier entero para random_state. El número específico no importa; lo que importa es usar el mismo número consistentemente en tu proyecto.

Parámetro shuffle

Por defecto, train_test_split mezcla los datos antes de dividir. Para la mayoría de las tareas de machine learning, esto es exactamente lo que quieres. Sin embargo, para datos de series temporales o cuando el orden importa, deberías deshabilitar el mezclado.

# Mezclado habilitado (predeterminado)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=True)
 
# Mezclado deshabilitado (para series temporales)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

Cuando shuffle=False, la función simplemente toma la primera porción para entrenamiento y la última porción para prueba, manteniendo el orden original.

Tabla de Referencia de Parámetros

ParámetroTipoPredeterminadoDescripción
test_sizefloat o intNoneProporción (0.0-1.0) o número de muestras para conjunto de prueba
train_sizefloat o intNoneProporción (0.0-1.0) o número de muestras para conjunto de entrenamiento
random_stateintNoneSemilla aleatoria para reproducibilidad
shuffleboolTrueSi se deben mezclar los datos antes de dividir
stratifyarray-likeNoneDatos para usar en división estratificada

División Estratificada para Datos Desbalanceados

Cuando tu conjunto de datos tiene clases desbalanceadas (algunas clases tienen muchas menos muestras que otras), la división aleatoria puede crear conjuntos de entrenamiento o prueba que representen mal la distribución general. Esto es especialmente problemático para tareas de clasificación.

El parámetro stratify asegura que la distribución de clases en los conjuntos de entrenamiento y prueba coincida con el conjunto de datos original.

import numpy as np
from sklearn.model_selection import train_test_split
 
# Crear conjunto de datos desbalanceado (90% clase 0, 10% clase 1)
X = np.random.randn(1000, 5)
y = np.array([0] * 900 + [1] * 100)
 
# Sin estratificación
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
print(f"Train distribution: Class 0: {sum(y_train == 0)}, Class 1: {sum(y_train == 1)}")
print(f"Test distribution: Class 0: {sum(y_test == 0)}, Class 1: {sum(y_test == 1)}")
 
# Con estratificación
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"\nStratified train distribution: Class 0: {sum(y_train == 0)}, Class 1: {sum(y_train == 1)}")
print(f"Stratified test distribution: Class 0: {sum(y_test == 0)}, Class 1: {sum(y_test == 1)}")

Con estratificación, tanto los conjuntos de entrenamiento como de prueba mantienen la distribución de clases 90/10. Sin ella, podrías tener suerte y obtener una división representativa, o podrías terminar con un conjunto de prueba que tiene solo 5% de la clase 1, lo que lleva a métricas de evaluación poco confiables.

Dividir Múltiples Arrays

Puedes dividir múltiples arrays a la vez, y sklearn asegurará que se dividan de la misma manera (mismos índices para todos los arrays).

import numpy as np
 
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)
sample_weights = np.random.rand(100)
 
# Dividir los tres arrays
X_train, X_test, y_train, y_test, weights_train, weights_test = train_test_split(
    X, y, sample_weights, test_size=0.2, random_state=42
)
 
print(f"X_train shape: {X_train.shape}")
print(f"y_train shape: {y_train.shape}")
print(f"weights_train shape: {weights_train.shape}")

Esto es particularmente útil cuando tienes pesos de muestras, múltiples variables objetivo o metadatos adicionales que necesitan dividirse consistentemente.

Train/Test Split vs Validación Cruzada vs Holdout

Diferentes estrategias de validación sirven para diferentes propósitos. Aquí hay una comparación:

MétodoUso de DatosCosto ComputacionalMejor ParaLimitaciones
Train/Test Split70-80% entrenamiento, 20-30% pruebaBajoEvaluación rápida de modelos, conjuntos de datos grandesEvaluación única, podría tener suerte/mala suerte con la división
Validación Cruzada100% usado para entrenamiento/prueba (k-fold)Alto (k veces más lento)Conjuntos de datos pequeños, estimación confiable de rendimientoComputacionalmente costoso, no para series temporales
Train/Val/Test (Holdout)60% entrenamiento, 20% validación, 20% pruebaMedioAjuste de hiperparámetros, evaluación finalSe necesitan más datos, flujo de trabajo más complejo
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
 
iris = load_iris()
X, y = iris.data, iris.target
model = RandomForestClassifier(random_state=42)
 
# Método 1: División simple train/test
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
model.fit(X_train, y_train)
print(f"Train/Test Split Score: {model.score(X_test, y_test):.3f}")
 
# Método 2: Validación cruzada de 5 pliegues
scores = cross_val_score(model, X, y, cv=5)
print(f"Cross-Validation Score: {scores.mean():.3f} (+/- {scores.std():.3f})")
 
# Método 3: División Train/Val/Test
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25, random_state=42)
model.fit(X_train, y_train)
print(f"Validation Score: {model.score(X_val, y_val):.3f}")
print(f"Test Score: {model.score(X_test, y_test):.3f}")

Para la mayoría de los proyectos, comienza con una simple división train/test. Usa validación cruzada cuando tengas datos limitados o necesites una estimación de rendimiento más robusta. Usa train/val/test cuando necesites ajustar hiperparámetros.

Técnicas Avanzadas de División

División de Series Temporales

Para datos de series temporales, el mezclado aleatorio destruye el orden temporal, lo que puede llevar a fugas de datos (usar información futura para predecir el pasado). Usa TimeSeriesSplit en su lugar:

from sklearn.model_selection import TimeSeriesSplit
import numpy as np
 
X = np.random.randn(100, 5)
y = np.random.randn(100)
 
tscv = TimeSeriesSplit(n_splits=5)
 
for train_index, test_index in tscv.split(X):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]
    print(f"Train size: {len(train_index)}, Test size: {len(test_index)}")

TimeSeriesSplit crea múltiples divisiones train/test donde cada conjunto de entrenamiento incluye todos los datos pasados hasta cierto punto, y el conjunto de prueba incluye el período inmediatamente siguiente. Esto simula pronósticos del mundo real donde solo tienes datos pasados para predecir el futuro.

GroupShuffleSplit para Datos Agrupados

Cuando tus datos tienen grupos (por ejemplo, múltiples mediciones del mismo paciente, múltiples transacciones del mismo cliente), necesitas asegurar que grupos enteros permanezcan juntos en el conjunto de entrenamiento o de prueba para evitar fugas de datos.

from sklearn.model_selection import GroupShuffleSplit
import numpy as np
 
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)
groups = np.array([0] * 25 + [1] * 25 + [2] * 25 + [3] * 25)  # 4 grupos
 
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
 
for train_idx, test_idx in gss.split(X, y, groups):
    X_train, X_test = X[train_idx], X[test_idx]
    y_train, y_test = y[train_idx], y[test_idx]
    print(f"Train groups: {np.unique(groups[train_idx])}")
    print(f"Test groups: {np.unique(groups[test_idx])}")

Esto asegura que todas las muestras de un grupo dado estén en el conjunto de entrenamiento o en el conjunto de prueba, nunca en ambos.

División Estratificada Multi-Output

Para problemas de clasificación multi-output, no puedes usar directamente stratify con un array 2D. En su lugar, crea una sola etiqueta que represente la combinación de todas las salidas:

import numpy as np
from sklearn.model_selection import train_test_split
 
X = np.random.randn(1000, 10)
y = np.random.randint(0, 2, (1000, 3))  # 3 salidas binarias
 
# Crear etiquetas combinadas para estratificación
y_combined = y[:, 0] * 4 + y[:, 1] * 2 + y[:, 2]
 
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y_combined
)
 
print(f"Original distribution: {np.unique(y_combined, return_counts=True)[1]}")
print(f"Train distribution: {np.unique(y_combined[:(len(X) - len(X_test))], return_counts=True)[1]}")

Mejores Prácticas para Train Test Split

Elegir la Proporción de División Correcta

Las proporciones de división más comunes son:

  • 80/20: Elección estándar para conjuntos de datos medianos a grandes (10,000+ muestras)
  • 70/30: Mejor para conjuntos de datos más pequeños (1,000-10,000 muestras) para tener una evaluación de prueba más robusta
  • 90/10: Para conjuntos de datos muy grandes (100,000+ muestras) donde incluso el 10% proporciona amplias muestras de prueba
  • 60/20/20: Para train/validation/test al ajustar hiperparámetros
import numpy as np
 
def recommend_split_ratio(n_samples):
    if n_samples < 1000:
        return "Consider cross-validation instead of simple split"
    elif n_samples < 10000:
        return "70/30 split recommended"
    elif n_samples < 100000:
        return "80/20 split recommended"
    else:
        return "90/10 or 80/20 split recommended"
 
sample_sizes = [500, 5000, 50000, 500000]
for size in sample_sizes:
    print(f"{size} samples: {recommend_split_ratio(size)}")

Evitar Fugas de Datos

Las fugas de datos ocurren cuando información del conjunto de prueba influye en el proceso de entrenamiento. Fuentes comunes:

  1. Preprocesamiento antes de dividir: Siempre divide primero, luego preprocesa
  2. Escalado de características en datos combinados: Ajusta el escalador solo en datos de entrenamiento
  3. Selección de características en datos combinados: Selecciona características usando solo datos de entrenamiento
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import numpy as np
 
X = np.random.randn(1000, 10)
y = np.random.randint(0, 2, 1000)
 
# INCORRECTO: Escalar antes de dividir (¡fuga de datos!)
scaler_wrong = StandardScaler()
X_scaled_wrong = scaler_wrong.fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X_scaled_wrong, y, test_size=0.2)
 
# CORRECTO: Dividir primero, luego escalar
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)  # Ajustar en datos de entrenamiento
X_test_scaled = scaler.transform(X_test)  # Transformar datos de prueba usando estadísticas de entrenamiento

El enfoque incorrecto usa información de todo el conjunto de datos (incluidas las muestras de prueba) para calcular parámetros de escalado, lo que filtra información sobre el conjunto de prueba al proceso de entrenamiento.

Estratificar Siempre que Sea Posible

Para problemas de clasificación, siempre usa división estratificada a menos que tengas una razón específica para no hacerlo. Esto es especialmente crítico para:

  • Conjuntos de datos desbalanceados
  • Conjuntos de datos pequeños
  • Problemas multi-clase con clases raras
from sklearn.model_selection import train_test_split
import numpy as np
 
# Conjunto de datos de enfermedad rara: 1% casos positivos
X = np.random.randn(1000, 20)
y = np.array([0] * 990 + [1] * 10)
 
# Sin estratificación - ¡podría no haber casos positivos en conjunto de prueba!
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=123)
print(f"Non-stratified test positives: {sum(y_test)}")
 
# Con estratificación - garantiza representación proporcional
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.1, random_state=123, stratify=y
)
print(f"Stratified test positives: {sum(y_test)}")

Errores Comunes a Evitar

1. Olvidar random_state

Sin random_state, tus resultados cambian cada vez que ejecutas el código. Esto hace que la depuración sea imposible y los experimentos irreproducibles.

# MALO: Sin random_state
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
 
# BUENO: Establecer random_state
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

2. No Estratificar Clases Desbalanceadas

Para conjuntos de datos desbalanceados, la división aleatoria puede crear conjuntos de prueba altamente no representativos, llevando a métricas de rendimiento poco confiables.

# MALO: Sin estratificación para datos desbalanceados
X_train, X_test, y_train, y_test = train_test_split(X, y_imbalanced, test_size=0.2)
 
# BUENO: Usar estratificación
X_train, X_test, y_train, y_test = train_test_split(
    X, y_imbalanced, test_size=0.2, stratify=y_imbalanced, random_state=42
)

3. Dividir Datos de Series Temporales con Shuffle

Los modelos de series temporales dependen del orden temporal. Mezclar destruye esta estructura y puede llevar a fugas de datos severas.

# MALO: Mezclar datos de series temporales
X_train, X_test, y_train, y_test = train_test_split(
    X_timeseries, y_timeseries, test_size=0.2, shuffle=True
)
 
# BUENO: Deshabilitar mezclado o usar TimeSeriesSplit
X_train, X_test, y_train, y_test = train_test_split(
    X_timeseries, y_timeseries, test_size=0.2, shuffle=False
)

4. Preprocesar Antes de Dividir

Ajustar preprocesadores (escaladores, imputadores, codificadores) en todo el conjunto de datos antes de dividir causa fugas de datos.

# MALO: Preprocesar antes de dividir
X_scaled = StandardScaler().fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2)
 
# BUENO: Dividir primero, luego preprocesar
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

5. Usar Conjunto de Prueba para Ajuste de Hiperparámetros

El conjunto de prueba solo debe usarse para evaluación final. Si lo usas para elegir hiperparámetros, esencialmente estás entrenando en tus datos de prueba.

# MALO: Ajustar en conjunto de prueba
from sklearn.ensemble import RandomForestClassifier
 
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
best_score = 0
best_params = None
for n_estimators in [10, 50, 100]:
    model = RandomForestClassifier(n_estimators=n_estimators)
    model.fit(X_train, y_train)
    score = model.score(X_test, y_test)  # ¡Usando conjunto de prueba!
    if score > best_score:
        best_score = score
        best_params = n_estimators
 
# BUENO: Usar conjunto de validación o validación cruzada
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25, random_state=42)
 
best_score = 0
best_params = None
for n_estimators in [10, 50, 100]:
    model = RandomForestClassifier(n_estimators=n_estimators)
    model.fit(X_train, y_train)
    score = model.score(X_val, y_val)  # Usando conjunto de validación
    if score > best_score:
        best_score = score
        best_params = n_estimators

Ejemplo Práctico: Flujo de Trabajo Completo

Aquí hay un flujo de trabajo completo de machine learning usando train_test_split correctamente:

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix
 
# Cargar datos
np.random.seed(42)
X = np.random.randn(1000, 10)
y = (X[:, 0] + X[:, 1] > 0).astype(int)  # Clasificación binaria
 
# Paso 1: Dividir datos (estratificado para conjunto de prueba balanceado)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# Paso 2: Preprocesar (ajustar solo en datos de entrenamiento)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
 
# Paso 3: Entrenar modelo
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train_scaled, y_train)
 
# Paso 4: Evaluar
y_pred = model.predict(X_test_scaled)
print("Classification Report:")
print(classification_report(y_test, y_pred))
print("\nConfusion Matrix:")
print(confusion_matrix(y_test, y_pred))
 
# Paso 5: Verificar sobreajuste
train_score = model.score(X_train_scaled, y_train)
test_score = model.score(X_test_scaled, y_test)
print(f"\nTrain accuracy: {train_score:.3f}")
print(f"Test accuracy: {test_score:.3f}")
print(f"Overfitting gap: {train_score - test_score:.3f}")

Usar RunCell para División Interactiva de Datos

Al trabajar en notebooks Jupyter, experimentar con diferentes proporciones de división y parámetros puede ser tedioso. RunCell (opens in a new tab) proporciona un agente de IA diseñado específicamente para flujos de trabajo de ciencia de datos en Jupyter. Puede ayudarte a:

  • Probar automáticamente múltiples proporciones de división y comparar resultados
  • Detectar fugas de datos en tu pipeline de preprocesamiento
  • Sugerir estrategias óptimas de estratificación para tu conjunto de datos específico
  • Generar curvas de validación para elegir la proporción train/test correcta

RunCell se integra directamente en tu entorno Jupyter, facilitando la iteración en tu estrategia de división de datos sin escribir código repetitivo.

Visualizar tus Datos con PyGWalker

Después de dividir tus datos, es crucial verificar que tus conjuntos de entrenamiento y prueba tengan distribuciones similares. PyGWalker (opens in a new tab) convierte tus DataFrames de pandas en visualizaciones interactivas estilo Tableau, facilitando:

  • Comparar distribuciones de características entre conjuntos de entrenamiento y prueba
  • Identificar posibles sesgos de muestreo en tus divisiones
  • Visualizar desbalances de clases y verificar que la estratificación funcionó correctamente
  • Explorar relaciones entre características en tus datos de entrenamiento
import pygwalker as pyg
import pandas as pd
 
# Convertir a DataFrames para visualización
train_df = pd.DataFrame(X_train, columns=[f'feature_{i}' for i in range(X_train.shape[1])])
train_df['dataset'] = 'train'
test_df = pd.DataFrame(X_test, columns=[f'feature_{i}' for i in range(X_test.shape[1])])
test_df['dataset'] = 'test'
 
combined = pd.concat([train_df, test_df])
 
# Crear visualización interactiva
pyg.walk(combined)

Esto te permite explorar interactivamente si tus distribuciones de entrenamiento y prueba coinciden, lo cual es crítico para una evaluación confiable del modelo.

FAQ

¿Cómo elijo entre división 80/20 y 70/30?

Usa 80/20 para conjuntos de datos mayores a 10,000 muestras, y 70/30 para conjuntos de datos más pequeños (1,000-10,000 muestras). La clave es asegurar que tu conjunto de prueba tenga suficientes muestras para una evaluación confiable, típicamente al menos 200-500 muestras para problemas de clasificación. Para conjuntos de datos muy grandes (100,000+ muestras), puedes usar 90/10 o incluso 95/5 ya que incluso el 5% proporciona miles de muestras de prueba.

¿Qué es random_state y por qué importa?

random_state es la semilla para el generador de números aleatorios que mezcla tus datos antes de dividir. Usar el mismo valor de random_state asegura que obtengas la división idéntica cada vez que ejecutes tu código, lo cual es esencial para reproducibilidad y depuración. Sin él, obtendrás diferentes divisiones train/test cada vez, haciendo imposible determinar si los cambios de rendimiento se deben a mejoras del modelo o solo a divisiones de datos afortunadas/desafortunadas.

¿Cuándo debo usar el parámetro stratify?

Usa stratify=y para todos los problemas de clasificación, especialmente cuando tienes clases desbalanceadas o conjuntos de datos pequeños. La estratificación asegura que la distribución de clases en tus conjuntos de entrenamiento y prueba coincida con la distribución general. Por ejemplo, si el 10% de tus datos son casos positivos, la estratificación garantiza que tanto los conjuntos de entrenamiento como de prueba tengan aproximadamente 10% de casos positivos, previniendo sesgos de evaluación por divisiones no representativas.

¿Puedo usar train_test_split para datos de series temporales?

No, no debes usar train_test_split con shuffle=True para datos de series temporales, ya que destruye el ordenamiento temporal y causa fugas de datos (usar datos futuros para predecir el pasado). En su lugar, usa train_test_split con shuffle=False para una división cronológica simple, o usa TimeSeriesSplit para validación cruzada que respete el orden temporal. Para series temporales, siempre asegura que los datos de entrenamiento vengan antes que los datos de prueba cronológicamente.

¿En qué se diferencia train_test_split de la validación cruzada?

train_test_split crea una sola partición train/test (típicamente 80/20), dándote una estimación de rendimiento. La validación cruzada (como k-fold) crea múltiples divisiones train/test y promedia los resultados, proporcionando una estimación de rendimiento más robusta. Usa train_test_split para evaluación rápida y conjuntos de datos grandes. Usa validación cruzada para conjuntos de datos pequeños (menos de 1,000 muestras) o cuando necesites estimaciones de rendimiento más confiables. La validación cruzada es k veces más lenta (por ejemplo, 5-fold es 5× más lento) pero reduce la varianza en tus métricas de rendimiento.

Conclusión

El train_test_split de sklearn es la herramienta fundamental para evaluar modelos de machine learning. Al dividir tus datos correctamente, obtienes estimaciones honestas de rendimiento que predicen el comportamiento del modelo en el mundo real. Recuerda los principios clave: siempre establece random_state para reproducibilidad, usa stratify para problemas de clasificación, evita preprocesar antes de dividir y elige tu proporción de división según el tamaño del conjunto de datos.

Domina estos fundamentos y evitarás las trampas más comunes que llevan a modelos sobreajustados y métricas de rendimiento engañosas. Ya sea que estés construyendo un clasificador simple o un sistema complejo de deep learning, la división train-test adecuada es el primer paso hacia machine learning confiable.

📚