[Python] train_test_split()

○ 개요
train_test_split은 데이터를 학습용(train)과 평가용(test)으로 무작위 분할하는 scikit-learn의 핵심 함수입니다. 지도학습에서 과적합을 방지하고, 모델의 일반화 성능을 공정하게 추정하기 위한 가장 기본 단계예요.

 

 

○ 핵심 정의

• 위치: sklearn.model_selection.train_test_split
• 역할: 전달받은 하나 이상의 배열/시퀀스(예: X, y)를 동일한 인덱스 기준으로 무작위 분할하여 (X_train, X_test, y_train, y_test, ...) 형태로 반환합니다.

 

 

○ 반환 순서

• 인자로 들어온 배열의 순서를 그대로 유지하여 train 먼저, test 다음 순서로 각 배열이 한 쌍씩 반환됩니다.
  예) train_test_split(X, y) → (X_train, X_test, y_train, y_test)

 

 

○ 주요 파라미터(요점 정리)

• test_size: 테스트 셋 비율 또는 개수. (예: 0.2면 20%, 100이면 샘플 100개)
• train_size: 학습 셋 비율 또는 개수. 일반적으로 test_size만 설정하고 나머지는 자동 계산.
• random_state: 난수 시드(재현성 보장). 같은 시드를 쓰면 매번 같은 분할 결과.
• shuffle: 분할 전 섞을지 여부(기본 True). 시계열 등 순서가 중요한 데이터는 False 권장.
• stratify: 층화 분할에 사용할 레이블 배열(보통 y). 분류(classification)에서 클래스 비율을 train/test에 동일하게 유지합니다. (단, shuffle=False와 함께 쓸 수 없음)
• 가변 인자: *arrays 형태로 여러 배열을 동시에 입력 가능(특성 X, 타깃 y, 샘플 가중치 등).

 

○ 사용 예시(필수 패턴)

 

1. 가장 기본적인 분할

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=42
)

2. 분류에서 클래스 비율 보존(층화 분할)

X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.25, random_state=0, stratify=y
)

3. 다중 배열 동시 분할(특성, 타깃, 샘플 가중치)

X_tr, X_te, y_tr, y_te, w_tr, w_te = train_test_split(
    X, y, sample_weight, test_size=0.2, random_state=1
)

4. 시계열(순서 보존) 데이터

# 순서가 의미 있으므로 섞지 않음
X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, shuffle=False
)
# 더 적합한 방법: TimeSeriesSplit 사용(아래 관련 개념 참고)

 

 

○ 왜 중요한가(의의)

• 일반화 추정: 테스트 셋은 모델이 ‘보지 않은’ 데이터여야 함.
• 누수 방지: 전처리·스케일링은 train에서 학습 → test에 적용(fit은 금지).
• 데이터 편향 제어: stratify로 클래스 불균형에 대응.

 

 

○ 자주 겪는 실수 & 예방 팁

• 데이터 누수: 분할 전에 스케일링/인코딩/특성 선택을 fit하면 누수 발생 → Pipeline 사용.
• 클래스 불균형 무시: 분류 문제에서 stratify=y 미설정 → 결과가 클래스 비율에 따라 왜곡.
• 시계열 셔플: 순서를 섞으면 미래 정보가 과거로 새어들 수 있음 → shuffle=False 또는 전용 분할기 사용.
• 작은 데이터셋의 과도한 test 비율: 표본 수가 너무 적으면 분산 커짐 → KFold/StratifiedKFold로 교차검증 권장.
• 재현성 미보장: random_state 미설정 → 실험 반복 시 결과가 달라짐.

 

 

○ 파라미터 심화 설명

• test_size vs train_size: 둘 다 줄 수도 있지만 보통 하나만 지정. 합이 1이 넘지 않도록 주의.
• random_state: int 또는 np.random.RandomState/Generator 호환 객체. 팀 내 고정값 공유 권장.
• shuffle: True가 기본. 단, stratify를 쓰려면 shuffle=True여야 하며, shuffle=False일 때 stratify는 사용 불가.
• stratify: 타깃의 범주형 레이블 분포를 유지. 연속 타깃(회귀)에는 부적합 → 회귀는 대신 분위수 기반 binning으로 가짜 레이블을 만들어 층화할 수 있음(고급 기법).

 

 

○ 실무 체크리스트

• 분할 전 데이터 샘플 수, 클래스 분포 확인했는가?
• 분류면 stratify=y를 적용했는가?
• 시계열/순서 데이터면 shuffle=False 또는 전용 분할기를 썼는가?
• 전처리 파이프라인을 구성해 누수를 차단했는가?
• 실험 로그에 random_state를 기록했는가?

 

○ 파이프라인과 함께 쓰기(누수 방지 정석)

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

pipe = make_pipeline(StandardScaler(), LogisticRegression())
pipe.fit(X_tr, y_tr)             # 스케일러는 train에서 학습
score = pipe.score(X_te, y_te)   # test에는 transform만 거쳐 평가

 

 

○ 관련/대체 도구

• KFold, StratifiedKFold: 더 안정적인 성능 추정을 위한 교차검증 분할기.
• GroupKFold, GroupShuffleSplit: 동일 그룹(예: 사용자, 세션) 간 정보 누수 방지.
• TimeSeriesSplit: 시계열 전용 점진적 분할.
• ShuffleSplit, StratifiedShuffleSplit: 반복적 랜덤 분할.
• train_test_split + validation 셋: 하이퍼파라미터 튜닝을 위해 train/valid/test 3-way 분할도 일반적.

 

○ 작은 데이터셋의 3분할 예시

# 1) 먼저 test 분리
X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
# 2) train을 다시 train/valid로
X_tr, X_va, y_tr, y_va = train_test_split(
    X_tr, y_tr, test_size=0.25, random_state=42, stratify=y_tr
)  # 전체 기준 0.2(valid), 0.2(test)

 

 

○ 성능/재현성 팁

• 동일한 데이터/코드라면 시드 고정이 가장 큰 재현성 확보 수단.
• 분할 직후 np.bincount(y_train)/y_test 등으로 클래스 비율을 바로 확인.
• 실험 기록에 분할 방식, 시드, 비율, 층화 여부를 함께 남기기.

 

○ 마무리 점검 체크리스트

• test_size가 과하지 않은가?
• 분류면 stratify=y로 클래스 분포를 유지했는가?
• 전처리·특성 선택이 train에서만 fit되도록 파이프라인을 썼는가?
• 시계열/그룹 데이터 특성을 고려한 분할기를 골랐는가?
• random_state를 문서화했는가?