from sklearn.model_selection import train_test_split
train_test_split(arrays, test_size, train_size, random_state, shuffle, stratify)
Return
-
X_train, X_test, Y_train, Y_test : array에 데이터와 레이블 둘 다 넣었을 경우의 반환.
데이터와 레이블의 순서쌍은 유지됨. -
X_train, X_test : array에 레이블 없이 데이터만 넣었을 경우의 반환.
Parameter
-
arrays : 분할 시킬 데이터 (python list, numpy array, pandas datafrmae...)
-
test_size : 테스트 데이터셋의 비율이나 갯수 (default=0.25)
-
train_size : 학습 데이터셋의 비율이나 갯수(default=test_size의 나머지)
-
random_state : 데이터 분할시 셔플이 이루어지는데 이를 위한 시드값 (int나 RandomState로 입력)
-
shuffle : 셔플 여부 설정(default=True)
-
stratify : 지정한 data의 비율을 지정함. 예를 들어 Label Set인 Y가 25%의 0과 75%의 1로 이루어진 Binary Set일 때, stratify=Y로 설정하면 나누어진 데이터셋들도 0과 1을 각각 25%, 75%로 유지한 채 분할된다.
예시
import numpy as np
from sklearn.model_selection import train_test_split
X = [[0, 1], [2, 3], [4,5], [6,7], [8,9]]
Y = [0, 1, 2, 3, 4]
# 데이터(X)만 넣었을 경우
X_train, X_test = train_test_split(X, test_size=0.2, random_state=123)
# X_train : [[0, 1], [6,7], [8,9], [2,3]]
# X_test : [[4, 5]]
# 데이터(X)와 레이블(Y)을 넣었을 경우
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.33, random_state=321)
# X_train : [[4,5], [0,1], [6,7]]
# Y_train : [2,0,3]
# X_test : [[2,3],[8,9]]
# Y_test : [1,4]
- 사용 방법이 쉽고 셔플까지 적용되므로 간단하고 합리적인 함수이지만 Validation Set을 따로 만들어주지는 않는다. Validation Set이 필요하다면 분할한 데이터 중 한 덩어리에 해당 함수를 한 번 더 사용하면 된다.
