[머신러닝] 사이킷런(Scikit-learn) 데이터 train data와 test data로 나누는 방법

 

사이킷런 데이터를 이용하여 머신러닝 모델을 학습시키려고 하면, 이 데이터를 train data와 test data로 나눠야 하는데요, 코드를 어떻게 작성해야 하는지 아래 코드를 보며 살펴보도록 하겠습니다.

 

 

1. 데이터 불러오기

먼저, 데이터를 불러오는 코드부터 살펴볼텐데요, 사이킷런 데이터에서 대표적으로 사용되는 iris data를 불러오는 코드를 작성해보겠습니다.

 

from sklearn import datasets
iris_dataset = datasets.load_iris()

 

이와 같이 입력한 뒤, 아래와 같이 iris_dataset.keys() 를 입력하고 실행하면, 'data'와 'target'이라는 array가 있는 것을 확인할 수 있습니다.

 

iris_dataset.keys()

# output :
dict_keys(['data', 'target', 'target_names', 'DESCR', 'feature_names', 'filename'])

 

여기서 'data'에 있는 값을 학습모델의 입력값(X), 'target'에 있는 값을 결과값(Y)으로 보도록 하겠습니다.

 

 

2. 데이터 나누기

그다음, 데이터를 나누려면 train_test_split 라이브러리를 사용해야 하는데요, 아래와 같이 코드를 작성하여 import 해줍니다.

 

from sklearn.model_selection import train_test_split

 

그리고나서, 아래와 같이 train_test_split 함수를 이용하여 데이터를 x_train(학습에 사용할 입력값), x_test(테스트에 사용할 입력값), y_train(학습에 사용할 결과값), y_test(테스트에 사용할 결과값) 으로 나눠줍니다.

 

x_train, x_test, y_train, y_test = train_test_split(iris_dataset['data'],
                                                   iris_dataset['target'],
                                                   random_state=0)

 

위 코드를 보면, train_test_split() 함수 안에는 먼저 어떤 입력데이터, 결과 데이터를 사용할지 인자로 넣어줘야 하는데요, 각각 iris_Dataset['data'], iris_dataset['target'] 데이터를 사용해주었습니다.

 

그리고나서, 데이터를 나눌 때 랜덤값이 일정하도록 random_state = 0으로 고정을 해주었습니다.

이렇게 입력한 뒤, 입력값의 train data와 test data, 결과값의 train data와 test data로 반환되도록 각각 x_train, x_test, y_train, y_test에 저장되도록 하였습니다.

 

이렇게 하고, 예시로 x_train을 출력해보면 array 배열로 입력값에 쓰일 데이터가 잘 나눠졌음을 살펴볼 수 있답니다.

 

한편, train_test_split 함수는 기본적으로 train 데이터와 test 데이터의 비율을 각각 75%, 25% 로 나눠주니, 이 점 참고하시기 바랍니다 :D