Comment créer un ensemble de train et de test à partir d’un DataFrame Pandas

Lors de l’ajustement de modèles d’apprentissage automatique à des ensembles de données, nous divisons souvent l’ensemble de données en deux ensembles :

1. Ensemble de formation : utilisé pour entraîner le modèle (70 à 80 % de l’ensemble de données d’origine)

2. Ensemble de tests : utilisé pour obtenir une estimation impartiale des performances du modèle (20 à 30 % de l’ensemble de données d’origine)

En Python, il existe deux manières courantes de diviser un DataFrame pandas en un ensemble d’entraînement et un ensemble de tests :

Méthode 1 : utilisez train_test_split() de sklearn

from sklearn.model_selection import train_test_split

train, test = train_test_split(df, test_size=0.2, random_state=0)

Méthode 2 : utiliser sample() à partir de pandas

train = df.sample(frac=0.8,random_state=0)
test = df.drop(train.index)

Les exemples suivants montrent comment utiliser chaque méthode avec le DataFrame pandas suivant :

import pandas as pd
import numpy as np

#make this example reproducible
np.random.seed(1)

#create DataFrame with 1,000 rows and 3 columns
df = pd.DataFrame({'x1': np.random.randint(30, size=1000),
                   'x2': np.random.randint(12, size=1000),
                   'y': np.random.randint(2, size=1000)})

#view first few rows of DataFrame
df.head()

        x1	x2	y
0	5	1	1
1	11	8	0
2	12	4	1
3	8	7	0
4	9	0	0

Exemple 1 : utilisez train_test_split() depuis sklearn

Le code suivant montre comment utiliser la fonction train_test_split() de sklearn pour diviser le DataFrame pandas en ensembles d’entraînement et de test :

from sklearn.model_selection import train_test_split

#split original DataFrame into training and testing sets
train, test = train_test_split(df, test_size=0.2, random_state=0)

#view first few rows of each set
print(train.head())

     x1  x2  y
687  16   2  0
500  18   2  1
332   4  10  1
979   2   8  1
817  11   1  0

print(test.head())

     x1  x2  y
993  22   1  1
859  27   6  0
298  27   8  1
553  20   6  0
672   9   2  1

#print size of each set
print(train.shape, test.shape)

(800, 3) (200, 3)

À partir du résultat, nous pouvons voir que deux ensembles ont été créés :

• Ensemble de formation : 800 lignes et 3 colonnes
• Ensemble de test : 200 lignes et 3 colonnes

Notez que test_size contrôle le pourcentage d’observations du DataFrame d’origine qui appartiendra à l’ensemble de test et que la valeur random_state rend la division reproductible.

Exemple 2 : utilisez sample() à partir de pandas

Le code suivant montre comment utiliser la fonction sample() de pandas pour diviser le DataFrame pandas en ensembles d’entraînement et de test :

#split original DataFrame into training and testing sets
train = df.sample(frac=0.8,random_state=0)
test = df.drop(train.index)

#view first few rows of each set
print(train.head())

     x1  x2  y
993  22   1  1
859  27   6  0
298  27   8  1
553  20   6  0
672   9   2  1

print(test.head())

    x1  x2  y
9   16   5  0
11  12  10  0
19   5   9  0
23  28   1  1
28  18   0  1

#print size of each set
print(train.shape, test.shape)

(800, 3) (200, 3)

À partir du résultat, nous pouvons voir que deux ensembles ont été créés :

• Ensemble de formation : 800 lignes et 3 colonnes
• Ensemble de test : 200 lignes et 3 colonnes

Notez que frac contrôle le pourcentage d’observations du DataFrame d’origine qui appartiendra à l’ensemble d’apprentissage et que la valeur random_state rend la division reproductible.

Ressources additionnelles

Les didacticiels suivants expliquent comment effectuer d’autres tâches courantes en Python :

Comment effectuer une régression logistique en Python
Comment créer une matrice de confusion en Python
Comment calculer la précision équilibrée en Python