Skip to main content

XGBoost with the scikit-learn API

This local-data example uses three partitions. The validation partition selects the boosting round; the test partition remains untouched until the fitted procedure is final.

1. Split before model selection​

Stratification preserves the class ratio in this IID demonstration. Real deployments may require group-aware or time-aware boundaries instead.

from sklearn.datasets import load_breast_cancer
from sklearn.metrics import brier_score_loss, log_loss, roc_auc_score
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

X, y = load_breast_cancer(return_X_y=True)
X_development, X_test, y_development, y_test = train_test_split(
    X,
    y,
    test_size=0.20,
    random_state=42,
    stratify=y,
)
X_train, X_valid, y_train, y_valid = train_test_split(
    X_development,
    y_development,
    test_size=0.20,
    random_state=42,
    stratify=y_development,
)

model = XGBClassifier(
    objective="binary:logistic",
    eval_metric="logloss",
    n_estimators=300,
    learning_rate=0.05,
    max_depth=3,
    min_child_weight=2.0,
    subsample=0.9,
    colsample_bytree=0.9,
    reg_lambda=1.0,
    early_stopping_rounds=15,
    random_state=42,
    n_jobs=1,
    verbosity=0,
)
model.fit(
    X_train,
    y_train,
    eval_set=[(X_valid, y_valid)],
    verbose=False,
)

test_probability = model.predict_proba(X_test)[:, 1]
print("Best iteration:", model.best_iteration)
print("Test ROC AUC:", round(roc_auc_score(y_test, test_probability), 3))
print("Test log loss:", round(log_loss(y_test, test_probability), 3))
print("Test Brier score:", round(brier_score_loss(y_test, test_probability), 3))

assert X_train.shape[0] + X_valid.shape[0] + X_test.shape[0] == X.shape[0]
assert test_probability.shape == y_test.shape
assert 0 <= model.best_iteration < model.n_estimators

2. API details that matter​

In XGBoost 3.4, eval_metric and early_stopping_rounds are estimator parameters. fit receives eval_set; it does not split data for you. The scikit-learn estimator's predict and predict_proba use the selected best_iteration after early stopping.

This example fixes a small, defensible configuration. If these hyperparameters are searched, use training-only folds and keep the test set outside the search. A fixed validation set reused across a large search can itself be overfit; nested cross-validation or a fresh final refit provides a stronger estimate.

3. Evaluate the prediction you deploy​

ROC AUC measures ranking, while log loss and Brier score assess probabilities. They can disagree. Select a hard decision threshold using validation data and real error costs, then report the test confusion matrix at that frozen threshold. Do not choose the threshold from the displayed test metrics.

4. Reproducibility and serving​

Persist the feature order, class mapping, missing-value policy, XGBoost version, and JSON/UBJSON model artifact. Record thread count and hardware for timing comparisons. Recheck probability calibration and subgroup behavior on data representative of deployment.