train_test_split의 파라미터에서 stratify는 언제 사용할까?

안녕하세요,

오늘은 모델의 성능을 평가하기 위해 훈련 데이터와 테스트 데이터를 분할하는 데 사용되는 함수인

train_test_split()와 그 파라미터인 stratify에 대하여 알아보겠습니다.

 

train_test_split란?

데이터셋을 train 데이터와 test 데이터로 분할하는 데 사용되는 함수입니다.

주로 모델의 성능을 평가하기 위해 데이터를 나눌 때 사용하며 Python의 머신러닝 라이브러리인 Scikit-Learn에서 제공됩니다.

 

- 주요 목적

• train 데이터는 모델을 학습시키기 위해 사용하고, test 데이터는 학습된 모델이 새로운 데이터에 대해 얼마나 잘 일반화되는지 평가하기 위해 사용합니다.
• 데이터를 훈련과 테스트로 분리하면, train 데이터에만 최적화된 모델이 되는 것을 방지하고 새로운 데이터에서도 좋은 성능을 발휘하는 모델을 만드는 데 도움이 됩니다.

 

사용법

train_test_split의 기본 사용법은 다음과 같습니다:

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=2025)

 

 

주요 파라미터

• X, y: 분할할 입력 데이터(X)와 타겟 레이블(y).
• test_size: 테스트 데이터의 비율을 설정합니다 (예: test_size=0.2는 전체 데이터의 20%를 테스트 세트로, 80%를 훈련 세트로 분할).
• random_state: 난수 시드를 설정해 데이터 분할을 재현 가능하게 합니다. 같은 숫자를 지정하면 매번 동일하게 데이터를 나눕니다.
• stratify: 분류 문제에서 각 클래스의 비율을 훈련 및 테스트 세트에 균일하게 유지하려면 y 값을 설정하여 사용할 수 있습니다.

 

활용 예시

• 분류 문제: 예를 들어, 고양이와 개를 분류하는 모델을 만들 때 train_test_split을 사용해 훈련 세트와 테스트 세트에 각 클래스가 고르게 분포되도록 나눌 수 있습니다.
• 회귀 문제: 연속적인 값을 예측하는 회귀 문제에서도 train_test_split을 사용하여 데이터를 나누어 훈련 및 테스트 성능을 평가할 수 있습니다.