본문 바로가기
AI/관련 자료

[Numpy] 사이킷런 train_test_split()으로 훈련 세트와 테스트 세트 분리하기.

by 뒹굴거리는프로도 2024. 2. 15.
반응형

 


 

사이킷런은 사용자 리스트나 배열을 적절히 섞은 뒤 비율에 맞게 훈련 세트와 테스트 세트로 나누어준다.

 

우선 함수를 import 한다.

from sklearn.model_selection import train_test_split

 

그리고 훈련 세트와 테스트 세트로 나눌 user_data와 user_target을 준비한다.


train_data, test_data, train_target, test_target = train_test_split(

	user_data, user_target, stratify = user_target, random_state  = 1 )

 

train_test_split 함수에 넣은 인자를 보자.

첫 번째 인자로 전체 데이터 세트인 user_data를 넣고,

두 번째 인자로 데이터 라벨 리스트에 해당하는 user_target을 넣는다.

세 번째 인자인 stratify 에 user_target을 넣어주면, 라벨 비율에 따라 적절하게 데이터를 나눠준다.

네 번째 인자인 random_state에 값을 지정하면 항상 같은 형태로 데이터를 나눌 수 있다.

 

사이킷런은 머신러닝 모델을 위한 다양한 유틸리티 도구를 제공하니 잘 숙지해두면 좋겠다.

반응형