🚀 Cómo crear un Test Set representativo en Machine Learning: Muestreo Estratificado con Scikit-Learn

Publicado en: MontesSTEAM.org | Categoría: Ciencia de Datos / Machine Learning

¡Hola a todos y todas!

En la entrega anterior aprendimos a inspeccionar el dataset de viviendas de California (California Housing Dataset) con Pandas. Antes de ponernos a entrenar modelos o transformar características, llega el paso más crítico del diseño: separar un conjunto de datos de prueba (Test Set) y guardarlo en un cajón sin volver a mirarlo.

En este tutorial aprenderemos:

1. Cómo instalar la librería Scikit-Learn en JupyterLab.
2. Por qué se debe separar el test set antes de profundizar en el análisis (para evitar el Data Snooping Bias).
3. Por qué usamos Muestreo Estratificado en lugar de un muestreo puramente aleatorio.
4. El código final paso a paso listo para ejecutar.

1. Paso previo: Instalación de Scikit-Learn

Para realizar la división de los datos utilizaremos Scikit-Learn, la librería estándar de aprendizaje automático en Python.

Si te ha aparecido el error ModuleNotFoundError: No module named 'sklearn', instala el paquete ejecutando esta línea en una celda de código de tu cuaderno en JupyterLab:

Python

%pip install scikit-learn

💡 Tip: Usar %pip install en lugar de !pip asegura que el paquete se instale exactamente en el entorno informático (kernel) donde se está ejecutando tu cuaderno.

2. ¿Por qué se crea el Test Set en esta fase? (Evitando el Data Snooping Bias)

Podría parecer apresurado apartar el 20% de los datos cuando apenas hemos empezado a explorar el dataset. Sin embargo, hay una razón fundamental: nuestro cerebro es una máquina experta en buscar patrones.

Si analizamos todo el conjunto de datos de antemano, nuestro cerebro detectará involuntariamente patrones específicos del grupo de prueba y tenderemos a elegir algoritmos o ajustar hiperparámetros sesgados hacia esos datos en particular.

Esto se conoce como Sesgo por Inspección de Datos (Data Snooping Bias). Si cometemos este error, cuando evaluemos el modelo final con el test set obtendremos una métrica engañosamente optimista, y el sistema fallará cuando lo despleguemos en producción.

3. ¿Por qué elegir este subconjunto mediante Muestreo Estratificado?

Un error muy común al empezar es dividir los datos de forma puramente aleatoria (por ejemplo, usando train_test_split estándar).

• El problema del muestreo aleatorio: Si el dataset no es gigantesco, elegir filas al azar corre el riesgo de introducir un sesgo de muestreo (Sampling Bias).
• El símil de las encuestas: Si una población tiene un 51.3% de mujeres y un 48.7% de hombres, una encuesta seria de 1.000 personas debe mantener exactamente esa misma proporción (513 mujeres y 487 hombres). Si se hiciera al azar, existe el riesgo de obtener una muestra sesgada que invalide el estudio.

Aplicación a nuestro Dataset de Viviendas

Los expertos nos indican que el ingreso mediano (median_income) es una variable determinante para predecir el precio de la vivienda (median_house_value).

Si dividiéramos al azar, corremos el riesgo de que en el conjunto de prueba queden sobrerrepresentados o infrarepresentados los distritos con ingresos extremadamente altos o muy bajos.

Por eso aplicamos Muestreo Estratificado (Stratified Sampling):

1. Convertimos la variable continua median_income en 5 categorías de ingresos (estratos) mediante pd.cut().
2. Usamos StratifiedShuffleSplit de Scikit-Learn para garantizar que tanto el conjunto de entrenamiento (80%) como el de prueba (20%) contengan exactamente la misma proporción de cada nivel de ingresos que el dataset original.

4. Código Python Paso a Paso

Aquí tenéis el bloque completo que podéis copiar y ejecutar en vuestro JupyterLab:

Python

import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedShuffleSplit

# 1. Crear la categoría de ingresos (5 estratos)
housing["income_cat"] = pd.cut(
    housing["median_income"],
    bins=[0.0, 1.5, 3.0, 4.5, 6.0, np.inf],
    labels=[1, 2, 3, 4, 5]
)

# 2. Dividir en train (80%) y test (20%) garantizando la misma proporción en cada estrato
split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)

for train_index, test_index in split.split(housing, housing["income_cat"]):
    strat_train_set = housing.loc[train_index]
    strat_test_set = housing.loc[test_index]

# 3. Eliminar la columna auxiliar 'income_cat' para dejar los DataFrames en su estado original
for set_ in (strat_train_set, strat_test_set):
    set_.drop("income_cat", axis=1, inplace=True)
```[cite: 2]

---

## 5. Comprobación de resultados

Al ejecutar la celda anterior no verás ninguna salida en pantalla (es normal en asignaciones de Python). Para comprobar que los subconjuntos han quedado guardados correctamente, ejecuta esto en una celda nueva:

```python
# Comprobamos la cantidad de registros
print("Filas en el conjunto de entrenamiento (train):", len(strat_train_set))
print("Filas en el conjunto de prueba (test):", len(strat_test_set))

# Verificamos que 'income_cat' ha sido borrada de la estructura
print("¿Existe 'income_cat' en el dataset final?:", "income_cat" in strat_train_set.columns)
```[cite: 2]

**Resultados esperados:**
```text
Filas en el conjunto de entrenamiento (train): 16512
Filas en el conjunto de prueba (test): 4128
¿Existe 'income_cat' en el dataset final?: False
```[cite: 2]

---

### 📝 Resumen final
Con estas tres líneas de Scikit-Learn y Pandas hemos separado `strat_test_set` (4,128 distritos)[cite: 2]. Ahora ya podemos trabajar libremente sobre `strat_train_set` sin miedo a sesgar las evaluaciones de nuestro futuro modelo[cite: 2].

¡En la siguiente entrada entraremos de lleno en las técnicas de visualización geográfica de datos! 🗺️