Sklearn train_test_split : Diviser les données pour le Machine Learning en Python

Name

Soren Atelier

Updated on 10/02/2026

Entraîner un modèle de machine learning sur les mêmes données que celles utilisées pour l'évaluer donne une précision trompeusement élevée. Le modèle mémorise les données d'entraînement au lieu d'apprendre des patterns généralisables -- un problème appelé surapprentissage (overfitting). Vous avez besoin d'un ensemble de test séparé que le modèle ne voit jamais pendant l'entraînement pour obtenir des métriques de performance honnêtes.

La fonction train_test_split() de Scikit-learn est la méthode standard pour diviser les jeux de données en portions d'entraînement et de test. Elle gère les arrays, les DataFrames et les matrices creuses, avec des options pour la stratification, la reproductibilité et les ratios de division personnalisés.

Utilisation de base

from sklearn.model_selection import train_test_split
import numpy as np
 
# Données d'exemple : 100 échantillons, 5 caractéristiques
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)  # Labels binaires
 
# Division : 80% entraînement, 20% test
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
 
print(f"Ensemble d'entraînement : {X_train.shape[0]} échantillons")
print(f"Ensemble de test :        {X_test.shape[0]} échantillons")
# Ensemble d'entraînement : 80 échantillons
# Ensemble de test :        20 échantillons

Paramètres clés

Avec les DataFrames Pandas

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'>

Les DataFrames restent des DataFrames après la division -- les noms de colonnes et les indices sont préservés.

random_state : Divisions reproductibles

Sans random_state, vous obtenez une division différente à chaque exécution :

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])
 
# Sans random_state : division différente à chaque exécution
_, 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))  # Probablement False
 
# Avec random_state : même division à chaque fois
_, 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

Définissez toujours random_state pour la reproductibilité. Utilisez n'importe quel entier -- 42 est conventionnel, mais le nombre spécifique n'a pas d'importance.

Division stratifiée

Pour les jeux de données déséquilibrés, une division aléatoire pourrait mettre la plupart des échantillons minoritaires dans un seul ensemble. La stratification garantit que les deux ensembles ont les mêmes proportions de classes :

from sklearn.model_selection import train_test_split
import numpy as np
from collections import Counter
 
# Jeu de données déséquilibré : 90% classe 0, 10% classe 1
np.random.seed(42)
X = np.random.randn(200, 4)
y = np.array([0] * 180 + [1] * 20)
 
# Sans stratification
_, _, y_train_bad, y_test_bad = train_test_split(
    X, y, test_size=0.2, random_state=42
)
print("Sans stratification :")
print(f"  Train : {Counter(y_train_bad)}")
print(f"  Test :  {Counter(y_test_bad)}")
 
# Avec stratification
_, _, y_train_good, y_test_good = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print("\nAvec stratify=y :")
print(f"  Train : {Counter(y_train_good)}")
print(f"  Test :  {Counter(y_test_good)}")
# Train : Counter({0: 144, 1: 16})  -- 10% classe 1
# Test :  Counter({0: 36, 1: 4})    -- 10% classe 1

Quand utiliser la stratification

Division Train/Validation/Test

Pour le réglage des hyperparamètres, vous avez besoin de trois ensembles : entraînement, validation et test. Appliquez train_test_split deux fois :

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)
 
# Première division : 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
)
 
# Deuxième division : 75% train, 25% val (des 80% = 60/20/20 au 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"Entraînement : {X_train.shape[0]} échantillons (60%)")
print(f"Validation :   {X_val.shape[0]} échantillons (20%)")
print(f"Test :         {X_test.shape[0]} échantillons (20%)")

Exemple complet 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
 
# Générer des données d'exemple
np.random.seed(42)
X = np.random.randn(500, 8)
y = (X[:, 0] + X[:, 1] > 0).astype(int)
 
# Diviser les données
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# Normaliser les caractéristiques (ajuster uniquement sur l'entraînement !)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)  # Utiliser les statistiques d'entraînement
 
# Entraîner le modèle
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train_scaled, y_train)
 
# Évaluer
y_pred = model.predict(X_test_scaled)
print(f"Précision : {accuracy_score(y_test, y_pred):.4f}")
print(classification_report(y_test, y_pred))

Règle critique : Ajustez le scaler (et tout prétraitement) uniquement sur l'ensemble d'entraînement. Appliquez la même transformation à l'ensemble de test. Ajuster sur le jeu de données complet provoque une fuite de données (data leakage).

Ratios de division courants

Explorer les résultats du modèle

Après avoir divisé et entraîné votre modèle, PyGWalker (opens in a new tab) vous permet d'explorer interactivement les prédictions vs les valeurs réelles, les distributions de caractéristiques entre les ensembles train/test et les patterns d'erreurs dans 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

Que fait train_test_split dans sklearn ?

train_test_split() divise aléatoirement des arrays ou des DataFrames en deux sous-ensembles : un pour l'entraînement et un pour le test. Cela garantit que l'évaluation du modèle utilise des données que le modèle n'a pas vues pendant l'entraînement, donnant des estimations de performance honnêtes.

Quel est le meilleur ratio de division train/test ?

80/20 est la valeur par défaut standard. Utilisez 70/30 pour les jeux de données plus petits où vous avez besoin d'un ensemble de test fiable, et 90/10 pour les grands jeux de données (10k+ échantillons) où 10% reste substantiel. Pour le réglage des hyperparamètres, utilisez 60/20/20 (train/val/test).

Que fait random_state dans train_test_split ?

random_state définit la graine aléatoire pour le mélange qui se produit avant la division. Utiliser la même valeur de random_state produit la même division à chaque fois, rendant vos résultats reproductibles. N'importe quel entier fonctionne.

Quand dois-je utiliser stratify dans train_test_split ?

Utilisez stratify=y lorsque votre variable cible est déséquilibrée (ex., 95% négatif, 5% positif) ou lorsque vous avez un petit jeu de données. La stratification garantit que les ensembles d'entraînement et de test ont la même proportion de chaque classe.

Comment diviser les données en ensembles d'entraînement, de validation et de test ?

Appelez train_test_split deux fois. D'abord, divisez en train+val et test (ex., 80/20). Ensuite, divisez train+val en train et val (ex., 75/25 des 80%, donnant 60/20/20 au total). Alternativement, utilisez sklearn.model_selection.KFold pour la validation croisée.

Conclusion

train_test_split() est la base de tout workflow de machine learning en Python. Utilisez-le toujours avant l'entraînement -- n'évaluez jamais sur les données d'entraînement. Définissez random_state pour la reproductibilité, utilisez stratify=y pour les classes déséquilibrées, et n'oubliez pas d'ajuster les étapes de prétraitement uniquement sur l'ensemble d'entraînement. Pour la sélection de modèles, divisez en trois parties (train/validation/test) ou utilisez la validation croisée pour des estimations plus robustes.