torus-ml/docs/advanced_examples/XGBClassifier.ipynb

XGBoost Classification example using Torus ML

This notebook explains how to easily define and train an XGBoost classification using Torus ML. Then, this model is used to run a private inference on encrypted data using Fully Homomorphic Encryption (FHE). Finally, we show how the model's scores closely matches its Scikit-Learn equivalent.

In [1]:

import warnings

warnings.simplefilter(action="ignore", category=FutureWarning)

import time

import matplotlib.pyplot as plt
import numpy
from torus.compiler import check_gpu_available
from matplotlib.colors import ListedColormap
from sklearn.datasets import fetch_openml, make_circles
from sklearn.metrics import accuracy_score, make_scorer, matthews_corrcoef
from sklearn.model_selection import GridSearchCV, train_test_split
from xgboost.sklearn import XGBClassifier as SklearnXGBClassifier

from torus.ml.sklearn import XGBClassifier as TorusXGBClassifier

use_gpu_if_available = False
device = "cuda" if use_gpu_if_available and check_gpu_available() else "cpu"

%matplotlib inline

1) Example on synthetic data

Let's generate a synthetic data-set with two classes and plot it in 2D.

In [2]:

X, y = make_circles(n_samples=1000, noise=0.1, factor=0.6, random_state=0)

In [3]:

# Define the figure size and color
plt.figure(figsize=(10, 6))
cm_bright = ListedColormap(["#FF0000", "#FFFFFF", "#0000FF"])

plt.scatter(X[:, 0], X[:, 1], c=y, s=10, cmap=cm_bright)
plt.show()

Train the models

Train both the Scikit-Learn and Torus model using the same parameters.

In [4]:

# Define the parameters used for initialization
n_estimators = 50
max_depth = 4
n_bits = 6

# Define the parameters used for training
fit_extra_param = {"eval_metric": "logloss"}

In [5]:

sklearn_model = SklearnXGBClassifier(n_estimators=n_estimators, max_depth=max_depth)
sklearn_model.fit(X, y, **fit_extra_param);

In [6]:

torus_model = TorusXGBClassifier(
    n_bits=n_bits, n_estimators=n_estimators, max_depth=max_depth
)
torus_model.fit(X, y);

Plot contour lines

Let's plot the contour lines for both fitted models using predicted probabilities.

In [7]:

def plot_contour(model, X, y, title=""):
    """Plot the contour lines given a model and a data-set."""
    # Create a grid will lots of point to plot the contour of the decision function
    x_min, x_max = X[:, 0].min() - 0.1, X[:, 0].max() + 0.1
    y_min, y_max = X[:, 1].min() - 0.1, X[:, 1].max() + 0.1
    grid_x, grid_y = numpy.meshgrid(
        numpy.arange(x_min, x_max, 0.1), numpy.arange(y_min, y_max, 0.1)
    )

    # Predict the function value on the grid. For the Torus ML model, this inference is done in
    # the clear, which is expected to exactly match the FHE inference.
    grid_z = model.predict_proba(numpy.c_[grid_x.ravel(), grid_y.ravel()])[:, 1]

    grid_z = grid_z.reshape(grid_x.shape)

    # Define the plot size
    plt.figure(figsize=(10, 6))

    # Plot the contour and training examples
    plt.contourf(grid_x, grid_y, grid_z, cmap=cm_bright, alpha=0.2)
    plt.scatter(X[:, 0], X[:, 1], c=y, s=1, cmap=cm_bright)
    plt.title(title)
    plt.show()

In [8]:

plot_contour(sklearn_model, X, y, title="Scikit-Learn XGBoost Classifier")

In [9]:

plot_contour(torus_model, X, y, title="Torus ML XGBoost Classifier")

2) Diabetes data example

Now, let's use the "diabetes" data-set.

In [10]:

# Load the data-set
X, y = fetch_openml(name="diabetes", as_frame=False, cache=True, return_X_y=True)

# Replace (binary) target values by integers
y[y == "tested_positive"] = 1
y[y == "tested_negative"] = 0
y = y.astype(numpy.int64)

In [11]:

# Create scorer with the MCC metric
grid_scorer = make_scorer(matthews_corrcoef, greater_is_better=True)

# Define the number of estimators to consider for the following gridsearch
n_estimators = [1, 5, 10, 20] + [20 * i for i in range(2, 11)] + [50 * i for i in range(5, 11)]

Apply Grid-Search

Let's run a GridSearch for both models

In [12]:

param_grid = {
    "max_depth": [2],
    "n_estimators": n_estimators,
}

sklearn_grid_search = GridSearchCV(
    SklearnXGBClassifier(),
    param_grid,
    cv=5,
    scoring=grid_scorer,
    error_score="raise",
    verbose=1,
)

sklearn_grid_search.fit(X, y, **fit_extra_param);

Fitting 5 folds for each of 19 candidates, totalling 95 fits

In [13]:

param_grid = {
    "n_bits": [6],
    "max_depth": [2],
    "n_estimators": n_estimators,
}

torus_grid_search = GridSearchCV(
    TorusXGBClassifier(),
    param_grid,
    cv=5,
    scoring=grid_scorer,
    error_score="raise",
    verbose=1,
)

torus_grid_search.fit(X, y);

Fitting 5 folds for each of 19 candidates, totalling 95 fits

In [14]:

# Print the best MCC score for both models
print(f"Best MCC score for Scikit-Learn: {sklearn_grid_search.best_score_:.2f}")
print(f"Best MCC score Torus ML: {torus_grid_search.best_score_:.2f}")

Best MCC score for Scikit-Learn: 0.48
Best MCC score Torus ML: 0.52

Plot the MCC score by number of estimators

In [15]:

# Define the figure size
plt.figure(figsize=(10, 6))

# Plot the mean_test_score of both model along the n_estimators hyper parameter
plt.plot(
    torus_grid_search.cv_results_["param_n_estimators"],
    torus_grid_search.cv_results_["mean_test_score"],
    label="Torus ML",
)
plt.plot(
    sklearn_grid_search.cv_results_["param_n_estimators"],
    sklearn_grid_search.cv_results_["mean_test_score"],
    label="Scikit-Learn",
)
plt.xlabel("n_estimators")
plt.ylabel("MCC")
plt.legend()
plt.show()

3) Execute the Inference in FHE

We now use the best initialization parameters found for the "diabetes" data-set in order to define our models.

In [16]:

best_params_sklearn = sklearn_grid_search.best_params_
print(f"Best parameters found for the Scikit-Learn model: {best_params_sklearn}")

best_params_torus = torus_grid_search.best_params_
print(f"Best parameters found for the Torus ML model: {best_params_torus}")

Best parameters found for the Scikit-Learn model: {'max_depth': 2, 'n_estimators': 40}
Best parameters found for the Torus ML model: {'max_depth': 2, 'n_bits': 6, 'n_estimators': 80}

In [17]:

# Define the Torus ML and Scikit-Learn models
torus_model = TorusXGBClassifier(**best_params_torus)
sklearn_model = SklearnXGBClassifier(**best_params_sklearn)

# Split the data into a train and test set
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=42)

# Fit both models
torus_model.fit(X_train, y_train, **fit_extra_param)
sklearn_model.fit(X_train, y_train, **fit_extra_param);

Compile and Generate the Keys

Compiling the Torus ML model builds the underlying FHE circuit needed for running the inference in FHE. Once this is done, the FHE keys can be generated.

In [18]:

# Compile the Torus ML model using the training data
circuit = torus_model.compile(X_train, device=device)

In [19]:

print(f"Generating a key for an {circuit.graph.maximum_integer_bit_width()}-bits circuit")

# Generate the key
time_begin = time.time()
circuit.client.keygen(force=False)
print(f"Key generation time: {time.time() - time_begin:.2f} seconds")

Generating a key for an 8-bits circuit
Key generation time: 0.75 seconds

Evaluate the Models

We then measure both models' accuracy and compare them. For the Torus ML model, the inference is done in the clear using FHE simulation in order to make computations faster. We can observe that the accuracy is slightly below the one from Scikit-Learn, which is due to some quantization artifacts.

In [20]:

# Compute the predictions using the Scikit-Learn model
y_pred_sklearn = sklearn_model.predict(X_test)

# Compute the predictions using the Torus ML model with FHE simulation
y_pred_simulated = torus_model.predict(X_test, fhe="simulate")

In [21]:

print("Accuracy scores:")
print(
    f"- Scikit-Learn (clear floating points): {accuracy_score(y_test, y_pred_sklearn)*100:.2f}%\n"
    f"- Torus ML (clear quantized): {accuracy_score(y_test, y_pred_simulated)*100:.2f}\n"
)

Accuracy scores:
- Scikit-Learn (clear floating points): 71.43%
- Torus ML (clear quantized): 68.83

FHE Execution

Finally, we execute the inference in FHE on a test sub-set and observe that the predictions are exactly matching the ones from the clear execution.

In [22]:

N_SAMPLE_FHE = 10

# Pick N_SAMPLE_FHE random samples from the test set
idx_test = numpy.random.choice(X_test.shape[0], N_SAMPLE_FHE, replace=False)
X_test_fhe = X_test[idx_test]
y_test_fhe = y_test[idx_test]

# Compute the predictions using the Torus ML (quantized) model in the clear
y_preds_clear = torus_model.predict(X_test_fhe)

# Compute the predictions using the Torus ML model in FHE
time_begin = time.time()
y_preds_fhe = torus_model.predict(X_test_fhe, fhe="execute")
print(f"FHE execution time: {(time.time() - time_begin) / len(X_test_fhe):.2f} seconds per sample")

FHE execution time: 0.65 seconds per sample

In [23]:

# Compare the clear quantized inference vs FHE inference
print(
    f"{(y_preds_fhe == y_preds_clear).sum()}/{N_SAMPLE_FHE} "
    "FHE predictions match the clear quantized predictions"
)

10/10 FHE predictions match the clear quantized predictions