[Sklearn] 파이썬 학습 데이터, 테스트 데이터 분리 : train_test_split

사이킷런 train / test 데이터 셋 분리 함수 사용법 정리

안녕하세요. 이번 글에서는 파이썬 scikit-learn 라이브러리에서

학습 데이터와 테스트 데이터를 원하는 조건으로 쉽게 분리 가능한

train_test_split 함수의 사용 방법에 대해서 정리해보도록 하겠습니다.

 

우선, 아주 간단한 1000개 행을 가진 데이터셋을 가정해보도록 하겠습니다.

feature는 3가지로, class label은 0과 1의 2가지로 설정해보았습니다.

import pandas as pd
import numpy as np

a = {'feature 1' : np.random.random(1000), 'feature 2' : np.random.random(1000), 
    'feature 3' : np.random.random(1000), 'class' : [0] * 700 + [1] * 300}

df = pd.DataFrame(a)
df

class label 개수의 분포도 미리 살펴보도록 하겠습니다.

df['class'].value_counts()
'''
0    700
1    300
Name: class, dtype: int64'''

0번 라벨이 700개, 1번 라벨이 300개로 불균등한 라벨 분포를 가지고 있음을 가정했습니다.

 

 

train_test_split 함수 기본 사용법

train_test_split 메소드는 기본적으로 학습 feature들을 모은 데이터프레임,

class label 칼럼을 input으로 받아 사용하게 됩니다.

 

만일 feature 1 ~ 3까지의 column을 전부 사용하고 싶다면 아래처럼 지정해주시면 됩니다.

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(df[['feature 1', 'feature 2', 'feature 3']],\
                                                    df['class'])

 

output은 4가지 형태인 X_train, X_test, y_train, y_test의 형태로 등장합니다.

 

X_train : 학습 데이터 셋의 feature 부분

X_test : 테스트 데이터 셋의 feature 부분

y_train : 학습 데이터 셋의 label 부분

y_test : 테스트 데이터 셋의 label 부분

 

을 의미하며, 형태는 기본적으로 feature 부분은 데이터프레임으로,

label 부분은 Series의 자료형으로 반환됩니다.

 

default로 설정된 학습 / 테스트 데이터셋의 비율은 75% / 25% 입니다.

 

X_train과 y_train의 형태를 직접 살펴보겠습니다.

(X_test, y_test도 길이만 짧은 동일한 형태입니다.)

750개 길이의 feature 데이터프레임, class label 정보 Series로 반환되었으며,

양쪽에서의 인덱스 순서(기존 데이터 내 위치)는 동일합니다.

반응형

학습, 테스트 데이터 셋 비율 설정

train, test 데이터 셋의 비율을 다르게 설정하고 싶은 경우는

train_size 인자 혹은 test_size 인자 중 1가지를 설정해주시면 됩니다.

 

예를 들어, 학습 / 테스트 셋의 크기 비율은 80% / 20%로 나누고 싶다면,

train_size = 0.8 혹은 test_size = 0.2로 지정해주시면 됩니다.

# 테스트 데이터 셋 비율 20%로 지정
X_train, X_test, y_train, y_test = train_test_split(df[['feature 1', 'feature 2', 'feature 3']],\
                                                    df['class'], test_size = 0.2)

 

 

학습, 테스트 셋에서 class 라벨 비율 동일하게 설정

train, test set 양쪽에서 class 라벨 비율을 같게 정하고 싶다면,

stratify 인자를 class 라벨이 들어있는 열로 지정해주시면 됩니다.

 

이해를 돕기 위하여, stratify 인자 설정 여부에 따른 두 결과의 차이를 살펴보겠습니다.

먼저, stratify를 미설정으로 나눈 결과입니다.

X_train, X_test, y_train, y_test = train_test_split(df[['feature 1', 'feature 2', 'feature 3']],\
                                                    df['class'], test_size = 0.2)

y_test.sum() # 62

200개의 y_test 중에서 0 라벨이 138개, 1 라벨이 62개로

기존 데이터에서의 70% : 30% 비율과 근소하게 다릅니다.

 

테스트 셋 내의 1 라벨 개수는 실행할때 마다 57개, 64개, 59개 등으로 달라질 수 있습니다.

 

이번에는 stratify를 df['class'] Series로 설정한 결과입니다.

X_train, X_test, y_train, y_test = train_test_split(df[['feature 1', 'feature 2', 'feature 3']],\
                                                    df['class'], test_size = 0.2, stratify = df['class'])

y_test.sum() # 60

200개의 y_test에서 0 라벨이 140개, 1 라벨이 60개로

항상 70% : 30%라는 기존 데이터 셋 내의 비율을 지키면서

데이터 셋을 나눌 수 있게 되었습니다.

 

 

기타 train_test_split 함수 arguments

이 외에도 random_state 인자를 원하는 정수로 설정하면

난수 추출의 seed를 정하여 실행시마다 분리 결과를 고정시킬 수 있으며,

 

shuffle 인자를 False로 설정하면 기존 데이터 셋 내에서의 순서를 기반으로

학습 / 테스트 데이터 셋을 분리할 수 있습니다.