Exporting cuml.accel Models to ONNX#

Models trained under cuml.accel can be exported to the ONNX format using sklearn-onnx. Because cuml.accel proxy objects are recognized by sklearn-onnx as scikit-learn estimators, you can pass them directly to convert_sklearn() without any conversion step.

The resulting .onnx file can then be loaded with ONNX Runtime for inference on both CPU and GPU, with no cuML dependency at inference time.

Supported Estimators#

The following cuml.accel estimators have been tested with sklearn-onnx:

  • Classifiers: RandomForestClassifier, KNeighborsClassifier, LinearSVC

  • Regressors: RandomForestRegressor, KNeighborsRegressor, LinearSVR

  • Transformers/Clusterers: PCA, TruncatedSVD, KMeans

Estimators not listed here (e.g. SVC, SVR, DBSCAN, TSNE, HDBSCAN, UMAP) are not currently supported for ONNX export.

[1]:

%load_ext cuml.accel

Training a Model#

Train a RandomForestClassifier on the Iris dataset. Under cuml.accel, this runs on the GPU.

[2]:

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

# ONNX requires float32 input
X, y = load_iris(return_X_y=True)
X = X.astype(np.float32)

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

clf = RandomForestClassifier(n_estimators=10, max_depth=3, random_state=0)
clf.fit(X_train, y_train)

accel_predictions = clf.predict(X_test)
print(f"Training accuracy: {clf.score(X_train, y_train):.4f}")
print(f"Test accuracy:     {clf.score(X_test, y_test):.4f}")

/opt/conda/envs/docs/lib/python3.13/site-packages/cuml/ensemble/randomforestclassifier.py:237: UserWarning: The number of bins, `n_bins` is greater than the number of samples used for training. Changing `n_bins` to number of training samples.
  return self._fit_forest(X_m, y_m)

Training accuracy: 0.9583
Test accuracy:     1.0000

Exporting to ONNX#

Convert the fitted model to ONNX using skl2onnx.convert_sklearn(). The proxy object is passed directly — no as_sklearn() call needed.

[3]:

from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# zipmap=False returns class probabilities as a 2D array instead of a
# list of dicts. This option only applies to classifiers with predict_proba.
initial_type = [("float_input", FloatTensorType([None, X_train.shape[1]]))]
onnx_model = convert_sklearn(
    clf, initial_types=initial_type, options={"zipmap": False}
)

onnx_path = "./rf_classifier.onnx"
with open(onnx_path, "wb") as f:
    f.write(onnx_model.SerializeToString())

print(f"ONNX model saved to {onnx_path} ({len(onnx_model.SerializeToString())} bytes)")

ONNX model saved to ./rf_classifier.onnx (4945 bytes)

Inference with ONNX Runtime#

The saved .onnx file can be loaded on any machine with onnxruntime installed — no cuML or sklearn-onnx needed. Use onnxruntime for CPU inference, or onnxruntime-gpu for GPU inference.

[4]:

import onnxruntime as ort

sess = ort.InferenceSession(onnx_path)
input_name = sess.get_inputs()[0].name
onnx_results = sess.run(None, {input_name: X_test})

onnx_predictions = onnx_results[0]
onnx_probabilities = onnx_results[1]

print(f"ONNX predictions shape:    {onnx_predictions.shape}")
print(f"ONNX probabilities shape:  {onnx_probabilities.shape}")

ONNX predictions shape:    (30,)
ONNX probabilities shape:  (30, 3)