XGBoost

Train gradient-boosted tree models with XGBoost in your browser: regression, classification, early stopping, feature importance and missing values.

Python
# Training a regression model using XGBoost

import xgboost as xgb
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# Generate a synthetic dataset
np.random.seed(42)
num_samples = 1000
X = pd.DataFrame({
    'num_rooms': np.random.randint(1, 10, num_samples),
    'square_footage': np.random.randint(500, 5000, num_samples),
    'age': np.random.randint(1, 100, num_samples)
})
y = X['num_rooms'] * 50000 + X['square_footage'] * 100 + X['age'] * -200 + np.random.normal(0, 10000, num_samples)

# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Create DMatrix for XGBoost
dtrain = xgb.DMatrix(X_train, label=y_train)
dtest = xgb.DMatrix(X_test, label=y_test)

# Define the parameters for the XGBoost model
params = {
    'objective': 'reg:squarederror',
    'max_depth': 3,
    'eta': 0.1,
    'eval_metric': 'rmse'
}

# Train the model
num_rounds = 10
model = xgb.train(params, dtrain, num_rounds)

# Make predictions
y_pred = model.predict(dtest)

# Evaluate the model
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
print(f'Root Mean Squared Error: {rmse}')

# Plot the true vs predicted values
import matplotlib.pyplot as plt

plt.scatter(y_test, y_pred, alpha=0.3)
plt.xlabel('True')
plt.ylabel('Predicted')
plt.title('True vs Predicted')
plt.tight_layout()
plt.show()
Root Mean Squared Error: 82159.70786651419

XGBoost trains gradient-boosted decision trees: it adds small trees one at a time, and each new tree corrects the errors the earlier ones left. People use it on tables of numbers and categories, to predict prices, score risk or flag customers likely to leave. On this page it runs in your browser, so you can try it without installing anything. Run the example first, then paste any snippet below into a new cell to try it.

What the example does

The code makes up 1,000 houses with a number of rooms, a floor area and an age, and a price that follows a fixed formula plus noise with a standard deviation of 10,000. train_test_split() from scikit-learn keeps 200 houses for testing, and xgb.DMatrix packs the data into XGBoost's own format. In params, reg:squarederror makes it a regression on squared error, max_depth: 3 limits each tree to three levels, and eta: 0.1 is the learning rate, which scales down each tree's correction. xgb.train() builds 10 trees, and the chart plots predicted against true prices.

Ten rounds at a learning rate of 0.1 are too few for this data. The RMSE is about 81,700: predictions run only from about 313,000 to 716,000, while true test prices run from about 92,000 to 914,000, so the points lie on a line flatter than the diagonal. With num_rounds = 100, the RMSE falls to about 12,300.

Classify with XGBClassifier

XGBClassifier and XGBRegressor work like scikit-learn models, with fit(), predict() and predict_proba(). This one tells malignant from benign tumors:

import xgboost as xgb
from sklearn.datasets import load_breast_cancer
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0, stratify=y)

clf = xgb.XGBClassifier(n_estimators=100, max_depth=3, learning_rate=0.1)
clf.fit(X_train, y_train)

print("accuracy:", round(accuracy_score(y_test, clf.predict(X_test)), 3))
print(clf.predict_proba(X_test[:3]).round(3))   # columns: malignant, benign

Stop training at the best round

Set n_estimators high and pass early_stopping_rounds. Training stops once the score on eval_set has not improved for that many rounds, and predict() then uses the trees up to the best round:

import xgboost as xgb
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split

X, y = make_regression(n_samples=1000, n_features=5, noise=10, random_state=0)
X_train, X_valid, y_train, y_valid = train_test_split(X, y, random_state=0)

model = xgb.XGBRegressor(n_estimators=1000, learning_rate=0.1, max_depth=3,
                         early_stopping_rounds=20)
model.fit(X_train, y_train, eval_set=[(X_valid, y_valid)], verbose=False)

print("best round:", model.best_iteration)
print("validation RMSE:", round(model.best_score, 2))

Rounds count from 0, so the best model has 461 trees.

See which features matter

plot_importance() draws one bar per feature. By default it counts how often a feature is used in a split. importance_type="gain" shows how much its splits reduce the error on average:

import matplotlib.pyplot as plt
import pandas as pd
import xgboost as xgb
from sklearn.datasets import load_diabetes

data = load_diabetes()
X = pd.DataFrame(data.data, columns=data.feature_names)
model = xgb.XGBRegressor(n_estimators=100, max_depth=3, learning_rate=0.1)
model.fit(X, data.target)

gain = model.get_booster().get_score(importance_type="gain")
print(sorted(gain, key=gain.get, reverse=True)[:3])

xgb.plot_importance(model, importance_type="gain", show_values=False)
plt.tight_layout()
plt.show()

s5, a blood serum measurement, ranks first, ahead of bmi and bp.

Train with missing values and categories

XGBoost reads NaN as missing: each split learns which branch missing values take, so you don't need to fill them in. A pandas category column works too, with enable_categorical=True:

import numpy as np
import pandas as pd
import xgboost as xgb
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import train_test_split

rng = np.random.default_rng(0)
df = pd.DataFrame({
    "size": rng.uniform(50, 200, 500),
    "city": pd.Categorical(rng.choice(["Austin", "Boston", "Denver"], 500)),
})
price = 3 * df["size"] + df["city"].map({"Austin": 100, "Boston": 300, "Denver": 200}).astype(float)
df.loc[rng.random(500) < 0.1, "size"] = np.nan   # about 10% missing

X_train, X_test, y_train, y_test = train_test_split(df, price, random_state=0)
model = xgb.XGBRegressor(n_estimators=200, max_depth=3, enable_categorical=True)
model.fit(X_train, y_train)
print("MAE:", round(mean_absolute_error(y_test, model.predict(X_test)), 1))

Most of the error comes from the 10 test rows without a size.

Good to know

  • xgb.train() takes eta and num_boost_round (default 10). The scikit-learn classes take learning_rate and n_estimators (default 100).
  • With early stopping, xgb.train() keeps every tree it built and predict() uses them all. Pass iteration_range=(0, model.best_iteration + 1) to use the best round only.
  • In this Python build, XGBRegressor.fit() on a text or category column without enable_categorical=True crashes Python with a Pyodide fatal error that try/except cannot catch. Convert text columns with .astype("category") and set the flag.