0Pricing
Machine Learning Academy · 강의

train_test_split: 비율, 시드, 계층화

sklearn의 train_test_split을 다양한 테스트 크기로 사용하고, 재현성을 위해 난수 시드를 설정하며, 균형 잡힌 분할에 계층화를 적용합니다.

train_test_split: 비율, 시드, 계층화은(는) CoddyKit의 무료 Machine Learning Academy 강의입니다. 이것은 4개 중 2번째 강의입니다. 아래에서 전체 강의를 무료로 읽을 수 있으며, 내장 코드 에디터와 24/7 AI 튜터와 함께 브라우저에서 직접 실습할 수 있습니다. 이 강의는 Machine Learning Academy 학습 경로의 일부이며, 진행 상황이 웹과 CoddyKit 앱에 동기화됩니다. Machine Learning Academy 강의에는 총 4개의 강의가 포함되어 있습니다.

이 강의의 일부는 아직 번역되지 않았으며 영어로 표시됩니다.

The train_test_split Function

Scikit-learn's train_test_split() function is the standard tool for partitioning datasets into training and test subsets. It randomly shuffles the data and allocates the specified proportion to the test set, returning four arrays: X_train, X_test, y_train, y_test.

Understanding its parameters fully is important because poor splitting choices can compromise the validity of all subsequent evaluation. A too-small test set gives noisy estimates; a too-large test set wastes training data. The right choice depends on your dataset size and the variance you can afford in performance estimates.

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_breast_cancer
import numpy as np

X, y = load_breast_cancer(return_X_y=True)

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.20,     # 20% of data goes to test set
    random_state=42,    # reproducible random shuffle
    shuffle=True,       # default: shuffle before splitting
    stratify=y          # preserve class proportions
)

print(f'Training size: {X_train.shape}')
print(f'Test size:     {X_test.shape}')

Choosing the Test Size Ratio

The test_size parameter accepts either a float (proportion) or an integer (absolute count):

  • Large datasets (>100k examples): use 10-15% test (there is plenty of training data; a large test set gives very stable estimates).
  • Medium datasets (1k–100k): 20-30% test is standard (80/20 or 70/30 splits).
  • Small datasets (<1k): consider cross-validation instead of a single split, because any fixed split may be unlucky.

There is no universally correct ratio. The goal is enough test examples for a statistically stable performance estimate, while keeping enough training data for the model to learn well.

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_breast_cancer
from sklearn.linear_model import LogisticRegression
import numpy as np

X, y = load_breast_cancer(return_X_y=True)

# Compare different test sizes
for test_size in [0.1, 0.2, 0.3, 0.4]:
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=42)
    model = LogisticRegression(max_iter=1000)
    model.fit(X_train, y_train)
    acc = model.score(X_test, y_test)
    print(f'test_size={test_size}: train={len(X_train)}, test={len(X_test)}, acc={acc:.3f}')

Random Seeds for Reproducibility

The random_state parameter seeds the random number generator that shuffles the data before splitting. With the same seed, the split is identical every time. Without a seed, the split changes on every run.

Reproducibility matters for: comparing models fairly (same split for all), sharing code with colleagues, debugging (same data each run), and scientific publication (results can be verified). Use a fixed seed for all experiments, but try two or three different seeds to confirm your model is not accidentally benefiting or suffering from one lucky split.

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_breast_cancer

X, y = load_breast_cancer(return_X_y=True)

# Test variance across random seeds
results = []
for seed in [0, 1, 42, 123, 999]:
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=seed
    )
    model = LogisticRegression(max_iter=1000)
    model.fit(X_train, y_train)
    acc = model.score(X_test, y_test)
    results.append(acc)
    print(f'seed={seed}: accuracy={acc:.4f}')

import numpy as np
print(f'\nMean: {np.mean(results):.4f}  Std: {np.std(results):.4f}')

The Problem Without Stratification

Without stratification, random splitting may produce a test set where the class proportions differ significantly from the overall dataset. If your dataset has 90% class 0 and 10% class 1, a random split might give a test set with only 3% class 1 by chance.

This matters because: performance metrics computed on an atypical class distribution are not representative of production performance. A model might score perfectly on a test set that contains almost no minority-class examples, giving false confidence in a model that actually fails on the positive class.

from sklearn.model_selection import train_test_split
import numpy as np

# Imbalanced dataset: 90% class 0, 10% class 1
np.random.seed(0)
n = 1000
y_imb = np.array([0]*900 + [1]*100)

results = []
for seed in range(10):
    _, _, y_train, y_test = train_test_split(y_imb, y_imb, test_size=0.2, random_state=seed)
    results.append(y_test.mean())
    print(f'seed={seed}: test positive rate = {y_test.mean():.3f} (expected: 0.100)')

print(f'\nMin: {min(results):.3f}  Max: {max(results):.3f}')

Stratification: Preserving Class Proportions

Setting stratify=y ensures that both the training and test sets contain the same proportion of each class as the full dataset. This is essential for imbalanced classification problems and should be used by default whenever you have a classification target.

Stratification works by splitting each class separately and then combining: if your dataset has 90% class 0 and 10% class 1, both the training set and test set will also have approximately 90%/10% proportions. This makes performance estimates on the test set much more stable and representative.

from sklearn.model_selection import train_test_split
import numpy as np

np.random.seed(0)
n = 1000
y_imb = np.array([0]*900 + [1]*100)
X_imb = np.random.randn(n, 5)

# Without stratification
_, _, y_train_no, y_test_no = train_test_split(X_imb, y_imb, test_size=0.2, random_state=42)
print(f'Without stratify: test positive rate = {y_test_no.mean():.3f}')

# With stratification
_, _, y_train_s, y_test_s = train_test_split(X_imb, y_imb, test_size=0.2, random_state=42, stratify=y_imb)
print(f'With stratify:    test positive rate = {y_test_s.mean():.3f}  (exactly 0.100)')

Stratification in Multi-Class Problems

Stratification works equally well for multi-class targets. When you pass stratify=y with a multi-class y, scikit-learn ensures all classes are represented in the same proportions in both splits.

This is particularly important for rare classes. If class C appears in only 2% of examples and you have a 500-example dataset, a random split might put all 10 class-C examples in training and none in the test set — making the test set useless for evaluating performance on class C. Stratification prevents this.

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
import numpy as np

X, y = load_iris(return_X_y=True)  # 3 balanced classes

# Stratified split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42, stratify=y
)

# Verify class proportions are preserved
for cls in [0, 1, 2]:
    full_ratio = (y == cls).mean()
    train_ratio = (y_train == cls).mean()
    test_ratio = (y_test == cls).mean()
    print(f'Class {cls}: full={full_ratio:.2f} train={train_ratio:.2f} test={test_ratio:.2f}')

Shuffling: Why Order Matters

By default, shuffle=True randomly reorders the data before splitting. This is critical because real datasets are often not in random order — they may be sorted by date, user ID, or label. Without shuffling, the test set would contain only the last examples in the dataset, which may be systematically different from earlier examples.

The exception is time-series data, where you should NOT shuffle. For time-ordered data, the test set must come from the future relative to training data. Setting shuffle=False preserves time order. Scikit-learn's TimeSeriesSplit provides more sophisticated time-ordered cross-validation.

from sklearn.model_selection import train_test_split
import numpy as np

# Time-series data: ordered by date
dates = np.arange(1000)           # day 0 to 999
sales = dates * 0.5 + np.random.randn(1000) * 10

# WRONG for time series: shuffling leaks future data into training
X_train_wrong, X_test_wrong, _, _ = train_test_split(dates.reshape(-1,1), sales, shuffle=True)
print(f'Wrong: test date range = {X_test_wrong.min():.0f} to {X_test_wrong.max():.0f}')

# CORRECT for time series: last 20% of time is the test set
X_train_right, X_test_right, _, _ = train_test_split(dates.reshape(-1,1), sales, shuffle=False)
print(f'Right: test date range = {X_test_right.min():.0f} to {X_test_right.max():.0f}')

Splitting Multiple Arrays Consistently

A common need is to split multiple arrays together — for example, X features, y labels, and a separate ID array — while keeping the same indices aligned. train_test_split accepts any number of arrays as positional arguments and applies the same shuffle and split to all of them.

This is also how you split pre-processed feature matrices alongside original DataFrames if you want to track which original rows ended up in which split for error analysis.

from sklearn.model_selection import train_test_split
import numpy as np

# Multiple arrays with same sample dimension
X = np.random.randn(100, 5)
y = np.random.randint(0, 2, 100)
sample_ids = np.arange(100)  # e.g., customer IDs

# Split all three together
X_train, X_test, y_train, y_test, ids_train, ids_test = train_test_split(
    X, y, sample_ids,
    test_size=0.2,
    random_state=42
)

print('Train IDs sample:', ids_train[:5])
print('Test IDs sample:', ids_test[:5])
print('All arrays remain aligned.')

When to Use Cross-Validation Instead

A single train-test split has high variance — the exact split can produce performance estimates that vary by several percentage points. Cross-validation averages performance across multiple splits, giving a more stable and reliable estimate.

Use cross-validation (not a single split) when:

  • Your dataset is small (fewer than a few thousand examples).
  • You want a reliable estimate with low variance.
  • You are comparing models or tuning hyperparameters.

Use a single split when you have very large datasets where cross-validation is computationally prohibitive.

from sklearn.model_selection import cross_val_score
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_breast_cancer
import numpy as np

X, y = load_breast_cancer(return_X_y=True)

# 5-fold cross-validation (more reliable than single split)
scores = cross_val_score(
    LogisticRegression(max_iter=1000),
    X, y,
    cv=5,
    scoring='accuracy'
)

print('CV scores:', scores.round(4))
print(f'Mean: {scores.mean():.4f}  Std: {scores.std():.4f}')
# Much more stable estimate than any single split

Complete Evaluation Workflow

The complete evaluation workflow that avoids all the pitfalls discussed in this lesson:

  1. Split off 20% as a final test set (stratified, random seed set).
  2. Use the remaining 80% for all development (training, cross-validation, hyperparameter tuning).
  3. Once development is complete and you have selected your final model, fit it on the entire 80% training data.
  4. Evaluate on the 20% test set exactly once and report that score.

Following this workflow gives you a trustworthy estimate of how your model will perform on future unseen data.

from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_breast_cancer
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline

X, y = load_breast_cancer(return_X_y=True)

# Step 1: Reserve test set FIRST
X_dev, X_test, y_dev, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# Step 2: Cross-validate on development set
pipeline = Pipeline([('scaler', StandardScaler()), ('clf', LogisticRegression(max_iter=1000))])
cv_scores = cross_val_score(pipeline, X_dev, y_dev, cv=5, scoring='accuracy')
print(f'Dev CV accuracy: {cv_scores.mean():.4f} (+/- {cv_scores.std():.4f})')

# Step 3: Fit final model on all dev data
pipeline.fit(X_dev, y_dev)

# Step 4: One-time final evaluation
print(f'Final test accuracy: {pipeline.score(X_test, y_test):.4f}')

Quick Check

Test your understanding of Machine Learning with Python concepts from this lesson.

Lesson Recap

In this lesson you learned: test_size controls the proportion of data held out for final evaluation, random_state seeds the shuffle for reproducibility and fair model comparison, and stratify=y preserves class proportions in both splits, which is critical for imbalanced datasets. Next up we study the bias-variance trade-off and learn to diagnose underfitting versus overfitting by plotting training and validation error curves.

자주 묻는 질문

“train_test_split: 비율, 시드, 계층화” 강의는 무료인가요?

네 — “train_test_split: 비율, 시드, 계층화” 전체 내용을 이 웹사이트에서 무료로 읽을 수 있습니다. 인터랙티브하게 실습하려면(내장 코드 에디터와 24/7 AI 튜터), CoddyKit PRO로 업그레이드하면 Machine Learning Academy 강의 전체를 잠금 해제할 수 있습니다. Machine Learning Academy 강의에는 총 4개의 강의가 포함되어 있습니다.

“train_test_split: 비율, 시드, 계층화”에서 뭘 배우나요?

sklearn의 train_test_split을 다양한 테스트 크기로 사용하고, 재현성을 위해 난수 시드를 설정하며, 균형 잡힌 분할에 계층화를 적용합니다. 브라우저에서 직접 실행하는 실습 코드로 Machine Learning Academy을(를) 배우며, 24/7 AI 튜터가 강의를 진행하면서 질문에 답변해줍니다.

Machine Learning Academy을(를) 시작하는 데 경험이 필요한가요?

사전 경험은 필요하지 않습니다. CoddyKit의 Machine Learning Academy은(는) 초급자부터 고급 학습자까지를 위해 구성되어 있으므로, 여기서 시작하거나 처음부터 시작할 수 있으며 자신의 속도대로 진행할 수 있습니다. 이것은 4개 중 2번째 강의입니다.

“train_test_split: 비율, 시드, 계층화” 강의는 얼마나 걸리나요?

대부분의 CoddyKit 강의는 약 5~10분이 소요됩니다. 각 강의는 간결하고 인터랙티브하여 꾸준한 진행이 가능하며, 웹과 앱에서 중단한 부분부터 바로 시작할 수 있습니다.

이 Machine Learning Academy 강의에서 코드를 작성하고 실행할 수 있나요?

네. 모든 Machine Learning Academy 강의에는 내장 코드 에디터가 포함되어 있으므로, 브라우저에서 바로 실제 코드를 작성하고 실행한 후 즉시 AI 피드백을 받을 수 있습니다 — 로컬 설정이 필요 없습니다.

이 강의의 모든 강의

  1. 훈련 데이터로 평가하면 안 되는 이유
  2. train_test_split: 비율, 시드, 계층화
  3. 편향-분산 상충 관계: 과소적합과 과적합
  4. 기준선 모델: DummyClassifier를 항상 능가하기
← Machine Learning Academy(으)로 돌아가기