Machine Learning: Entrenamiento y Prueba en Python

La división train/test es el paso más fundamental al construir un modelo de machine learning. La idea es simple: mantén una parte de tus datos oculta del modelo durante el entrenamiento y luego mide qué tan bien funciona el modelo en esa porción oculta. Sin esta separación, no tienes una forma honesta de saber si tu modelo ha aprendido genuinamente un patrón o simplemente ha memorizado los ejemplos de entrenamiento.

Este capítulo explica cómo funciona train_test_split de scikit-learn, qué hace cada parámetro y cómo evaluar el modelo resultante, tanto para problemas de regresión como de clasificación.

¿Por qué separar los datos de entrenamiento de los de prueba?

Un modelo que entrena y evalúa con los mismos datos parecerá mucho más preciso de lo que realmente es. Esto se llama fuga de datos u overfitting: el modelo ha memorizado los ejemplos de entrenamiento en lugar de aprender un patrón generalizable.

Imagina un estudiante que estudia 100 preguntas de práctica y luego hace un examen con esas mismas 100 preguntas. Puede obtener 100% — pero esa puntuación no dice nada sobre si entiende la materia.

La división train/test evita esto al:

• Entrenar el modelo en una porción de los datos para que pueda aprender patrones.
• Probar el modelo en una porción separada y no vista para medir el rendimiento en el mundo real.

Una división típica es 80% entrenamiento / 20% prueba, aunque la proporción correcta depende de cuántos datos tengas.

La función train_test_split

scikit-learn proporciona train_test_split en su módulo model_selection. Mezcla aleatoriamente tu conjunto de datos y lo divide en dos partes:

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, random_state=42
)

Los cuatro valores de retorno siempre están en este orden: características de entrenamiento, características de prueba, etiquetas de entrenamiento, etiquetas de prueba.

Parámetros clave

Efecto de test_size

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris

X, y = load_iris(return_X_y=True)   # 150 samples

for ts in [0.1, 0.2, 0.3]:
    X_tr, X_te, _, _ = train_test_split(X, y, test_size=ts, random_state=42)
    print(f"test_size={ts}: train={len(X_tr)}, test={len(X_te)}")

Salida:

test_size=0.1: train=135, test=15
test_size=0.2: train=120, test=30
test_size=0.3: train=105, test=45

random_state y reproducibilidad

Sin un random_state, la división es diferente cada vez que ejecutas el script. Establécelo en cualquier entero para obtener resultados reproducibles:

import numpy as np
from sklearn.model_selection import train_test_split

X = np.arange(10).reshape(-1, 1)
y = np.arange(10)

_, X_te1, _, _ = train_test_split(X, y, test_size=0.3, random_state=42)
_, X_te2, _, _ = train_test_split(X, y, test_size=0.3, random_state=42)

print("Same random_state → same split:", list(X_te1.ravel()) == list(X_te2.ravel()))
# True

La elección del entero (42, 0, 1, etc.) no importa — siempre que uses el mismo valor de forma consistente.

Ejemplo de regresión: predicción de precios de casas

El siguiente ejemplo genera un conjunto de datos sintético de tamaños y precios de casas, entrena un modelo de regresión lineal y lo evalúa en el conjunto de prueba reservado.

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score

# Generate synthetic data: house size (sq ft) → price
np.random.seed(42)
n = 200
sqft = np.random.randint(500, 3500, n).astype(float)
price = 150 * sqft + np.random.randn(n) * 20000

X = sqft.reshape(-1, 1)
y = price

# Split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
print("Training samples:", len(X_train))   # 160
print("Testing samples:", len(X_test))     # 40

# Train
model = LinearRegression()
model.fit(X_train, y_train)

# Evaluate on the test set
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
r2  = r2_score(y_test, y_pred)

print(f"Mean Squared Error: {mse:,.0f}")
print(f"R² Score:           {r2:.4f}")
print(f"Coefficient:        {model.coef_[0]:.2f}")
print(f"Intercept:          {model.intercept_:.2f}")

Salida:

Training samples: 160
Testing samples: 40
Mean Squared Error: 489,271,263
R² Score:           0.9651
Coefficient:        147.55
Intercept:          5237.02

Interpretación de las métricas:

• MSE (Error cuadrático medio) es la diferencia cuadrática promedio entre las predicciones y los valores reales. Menor es mejor, pero la escala depende de tu variable objetivo (aquí, dólares).
• R² va de 0 a 1. Un valor de 0.965 significa que el modelo explica aproximadamente el 96.5% de la varianza en los precios de las casas — un buen ajuste para este conjunto de datos simple.

Para más información sobre regresión lineal, consulta Regresión Lineal en Python.

Ejemplo de clasificación: especies de flores Iris

Para una tarea de clasificación, la exactitud y un informe de clasificación son más informativos que el MSE.

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report

X, y = load_iris(return_X_y=True)

# stratify=y ensures each class is proportionally represented in both splits
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print("Train class distribution:", np.bincount(y_train))  # [40 40 40]
print("Test class distribution: ", np.bincount(y_test))   # [10 10 10]

model = LogisticRegression(max_iter=200)
model.fit(X_train, y_train)

y_pred = model.predict(X_test)
print("\nAccuracy:", round(accuracy_score(y_test, y_pred), 4))
print()
print(classification_report(y_test, y_pred,
      target_names=["setosa", "versicolor", "virginica"]))

Salida:

Train class distribution: [40 40 40]
Test class distribution:  [10 10 10]

Accuracy: 0.9667

              precision    recall  f1-score   support

      setosa       1.00      1.00      1.00        10
  versicolor       1.00      0.90      0.95        10
   virginica       0.91      1.00      0.95        10

    accuracy                           0.97        30
   macro avg       0.97      0.97      0.97        30
weighted avg       0.97      0.97      0.97        30

Por qué importa stratify=y: Sin él, una división aleatoria de un conjunto de datos desbalanceado podría poner la mayoría de las muestras de una clase rara en el entrenamiento, sin dejar ninguna en el conjunto de prueba. stratify=y garantiza que cada clase aparezca en las mismas proporciones en ambas divisiones.

Para más información sobre clasificación, consulta Regresión Logística en Python y Matriz de Confusión en Python.

Errores comunes

¿Preprocesamiento antes o después de la división?

Siempre divide antes de ajustar cualquier paso de preprocesamiento (escalado, codificación, imputación). Si escala todos tus datos y luego divides, el conjunto de prueba ha influido en el escalador — una forma de fuga de datos.

El orden correcto:

from sklearn.preprocessing import StandardScaler
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, random_state=42)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)  # fit ONLY on training data
X_test_scaled  = scaler.transform(X_test)       # transform test with same parameters

Consulta Escalado de características en Python para una guía completa.

Mezcla y datos de series temporales

train_test_split mezcla los datos por defecto. Para datos de series temporales esto es incorrecto — estarías entrenando con datos futuros para predecir el pasado. Establece shuffle=False y asegúrate de que tus datos estén ordenados cronológicamente antes de dividir:

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, shuffle=False
)

Cuando una sola división no es suficiente

Una única división 80/20 te da una estimación del rendimiento del modelo que depende de qué muestras terminaron en el conjunto de prueba. La validación cruzada repite la división varias veces y promedia las puntuaciones, dando una estimación mucho más estable, especialmente en conjuntos de datos pequeños.

Elección de una métrica de evaluación

La métrica correcta depende del tipo de problema:

Para clasificación binaria, consulta Curva AUC-ROC en Python y Matriz de Confusión en Python. Para ajuste de hiperparámetros tras tener una división funcional, consulta Grid Search en Python.

Resumen

• Divide tus datos antes de cualquier preprocesamiento para evitar la fuga de datos.
• Usa test_size=0.2 como valor predeterminado razonable; ajusta según el tamaño del conjunto de datos.
• Establece random_state en cualquier entero para obtener divisiones reproducibles.
• Usa stratify=y para tareas de clasificación, especialmente con datos desbalanceados.
• Establece shuffle=False para datos de series temporales.
• Una sola división train/test es una línea base rápida; usa la validación cruzada para estimaciones de rendimiento más confiables.