머신러닝 모델을 평가할 때, 데이터를 학습용(train)과 테스트용(test)으로 분할하는 것은 매우 중요합니다. 사이킷런(sklearn) 라이브러리는 이를 쉽게 할 수 있도록 train_test_split 함수를 제공합니다. 이 글에서는 train_test_split 함수의 사용법과 함께 shuffle 및 stratify 옵션에 대해 설명하고, 언제 이러한 옵션을 사용하는 것이 좋은지 알아보겠습니다.
train_test_split 함수란?
train_test_split 함수는 데이터를 학습용과 테스트용으로 나누는 데 사용됩니다. 이를 통해 모델의 성능을 평가하고, 모델이 새로운 데이터에 대해 얼마나 잘 일반화되는지를 확인할 수 있습니다.
주요 파라미터
- arrays (X, y): 분할할 데이터를 나타냅니다. 보통 특징(feature) 데이터 X와 레이블(label) 데이터 y가 전달됩니다.
- test_size: 테스트 데이터셋의 비율 또는 개수를 지정합니다. 예를 들어, test_size=0.2는 전체 데이터의 20%를 테스트 데이터로 분리합니다. 기본값은 0.25입니다.
- train_size: 학습 데이터셋의 비율 또는 개수를 지정합니다. 명시되지 않으면 test_size를 기준으로 자동 설정됩니다.
- random_state: 데이터를 분할하는데 사용되는 난수 시드(seed)를 설정합니다. 동일한 random_state 값을 사용하면 항상 같은 방식으로 데이터를 분할할 수 있습니다.
- shuffle: 데이터를 분할하기 전에 섞을지 여부를 지정합니다. 기본값은 True입니다.
- stratify: 데이터를 분할할 때 층화(stratified) 방식으로 분할하도록 하는 기준입니다. 주로 불균형 데이터셋에서 레이블의 분포를 유지하기 위해 사용됩니다.
예제 코드
다음은 데이터 수를 늘리고 shuffle과 stratify 옵션을 사용한 예제 코드입니다.
from sklearn.model_selection import train_test_split
import numpy as np
# 예제 데이터 생성
# 특징 데이터(X)는 2차원 배열로, 100개의 샘플을 가지며 각각 2개의 특징을 가집니다.
# 레이블 데이터(y)는 0과 1의 두 클래스 레이블을 가집니다.
X = np.random.rand(100, 2)
y = np.random.randint(0, 2, 100)
# 데이터 분할
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.25, random_state=42, shuffle=True, stratify=y)
print("X_train shape:", X_train.shape)
print("X_test shape:", X_test.shape)
print("y_train distribution:", np.bincount(y_train))
print("y_test distribution:", np.bincount(y_test))
shuffle과 stratify 옵션의 중요성
shuffle (섞기)
- 언제 사용하는가?
- 일반적으로 데이터를 분할할 때는 shuffle 옵션을 True로 설정하여 데이터를 랜덤하게 섞는 것이 좋습니다. 이는 특정 순서에 의해 데이터가 편향되는 것을 방지할 수 있습니다. 특히, 데이터가 시간 순서나 어떤 규칙에 따라 정렬되어 있을 때, 이를 섞지 않으면 학습된 모델이 그러한 패턴을 학습하여 과적합(overfitting)될 수 있습니다.
- 예시
- 예를 들어, 날씨 데이터를 예측할 때, 데이터가 시간 순서대로 정렬되어 있다면, shuffle=True로 설정하여 이를 섞어주는 것이 좋습니다.
stratify (층화 분할)
- 언제 사용하는가?
- 레이블의 분포가 불균형한 데이터셋에서 stratify 옵션을 사용하면 원본 데이터의 레이블 비율을 학습용과 테스트용 데이터셋에 동일하게 유지할 수 있습니다. 이는 모델이 소수 클래스에 대해 적절히 학습할 수 있도록 도와줍니다.
- 예시
- 예를 들어, 클래스 0이 90%이고 클래스 1이 10%인 이진 분류 문제에서, stratify=y로 설정하면 학습용과 테스트용 데이터셋 모두에서 클래스 0과 클래스 1의 비율이 원본 데이터셋과 동일하게 유지됩니다.
결론
사이킷런의 train_test_split 함수는 데이터를 학습용과 테스트용으로 나누는 데 매우 유용합니다. shuffle 옵션을 통해 데이터를 랜덤하게 섞어 편향을 방지하고, stratify 옵션을 통해 레이블의 분포를 유지할 수 있습니다. 이러한 옵션들을 적절히 사용하면 모델의 성능을 보다 정확하게 평가하고 일반화 능력을 향상시킬 수 있습니다.
이를 통해 보다 신뢰성 있는 머신러닝 모델을 구축할 수 있습니다. 여러분도 train_test_split 함수를 활용하여 데이터를 분할해보고, 다양한 옵션들을 실험해 보세요!