train_test_split() - (4)

stratify: stratify 파라미터는 분류 문제를 다룰 때 매우 중요하게 활용되는 파라미터 값 입니다. stratify 값으로는 target 값을 지정해주면 됩니다.

stratify값을 target 값으로 지정해주면 target의 class 비율을 유지 한 채로 데이터 셋을 split 하게 됩니다. 만약 이 옵션을 지정해주지 않고 classification 문제를 다룬다면, 성능의 차이가 많이 날 수 있습니다.

# 라이브러리 로딩

from sklearn.model_selection import train_test_split



#train_test_split() 메소드를 이용해 train/validation 데이터 나누기 

# stratify 옵션을 활용하여 데이터 셋 split



x_train,x_valid, y_train, y_valid = train_test_split(train_x,train['category'],stratify = train['category'])



# y_train,y_valid 비율 확인 (value_counts())



print(y_train.value_counts())

print(y_valid.value_counts())





output :

2    10021

1    10003

0     9976

Name: category, dtype: int64



2    3341

1    3334

0    3325