ONNX Runtime : ton modèle IA 3 fois plus rapide sans changer une ligne de code

Ton modèle scikit-learn est précis, mais dès que tu le mets en production, il rame. Chaque prédiction coûte des millisecondes, et sur des milliers de requêtes par jour, l'addition devient salée. Tu pourrais réécrire tout ton pipeline en C++ ou acheter un GPU. Il existe une troisième voie, plus simple : ONNX Runtime.

ONNX Runtime est un moteur d'inférence open source, porté par Microsoft. Il exécute des modèles exportés au format ONNX (Open Neural Network Exchange), un standard qui fige ton modèle dans un graphe de calcul optimisé. Résultat : la même logique, les mêmes prédictions, mais une inférence souvent deux à cinq fois plus rapide, sur simple CPU. Dans ce tutoriel, on exporte une forêt aléatoire scikit-learn et on mesure le gain. Sur ma machine, x3,4.

Étape 1 : installer le trio

Trois paquets suffisent : scikit-learn pour entraîner, skl2onnx pour convertir, et onnxruntime pour exécuter.

bash

pip install onnxruntime scikit-learn skl2onnx

Étape 2 : entraîner un modèle scikit-learn

On part d'un cas réaliste : une forêt aléatoire qui classe des échantillons sur 20 variables. Les données sont synthétiques, mais le principe vaut pour n'importe quel modèle de ton pipeline.

python

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

# 10 000 lignes, 20 variables
X, y = make_classification(
    n_samples=10_000, n_features=20, n_informative=10, random_state=42
)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
print("Score :", round(model.score(X_test, y_test), 3))

Étape 3 : exporter vers ONNX

skl2onnx transforme ton modèle en graphe ONNX. Il suffit de déclarer le type d'entrée : des flottants, avec un nombre de lignes quelconque.

python

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

# Entrée : des flottants, nombre de lignes non fixé
initial_type = [("float_input", FloatTensorType([None, 20]))]

onx = convert_sklearn(model, initial_types=initial_type)
with open("model.onnx", "wb") as f:
    f.write(onx.SerializeToString())

print("Modèle exporté : model.onnx")

Étape 4 : l'inférence avec ONNX Runtime

On recharge le fichier model.onnx et on fait des prédictions. Attention : pour un classifieur, ONNX Runtime renvoie des probabilités (une ligne par échantillon, une colonne par classe), pas directement l'étiquette. Un simple argmax fait le pont.

python

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name

X_batch = X_test.astype(np.float32)
proba = sess.run(None, {input_name: X_batch})[0]
pred = proba.argmax(axis=1)
print("Prédictions :", pred[:10])

Étape 5 : mesurer le gain, sans se mentir

Un benchmark honnête compare les deux sur le même lot, plusieurs fois, pour lisser le bruit.

python

import time

X_batch = X_test.astype(np.float32)

t0 = time.perf_counter()
for _ in range(50):
    model.predict(X_batch)
sklearn_time = time.perf_counter() - t0

t0 = time.perf_counter()
for _ in range(50):
    sess.run(None, {input_name: X_batch})
onnx_time = time.perf_counter() - t0

print(f"scikit-learn : {sklearn_time:.3f}s")
print(f"ONNX Runtime : {onnx_time:.3f}s")
print(f"Accélération : x{sklearn_time / onnx_time:.1f}")

Le gain dépend du modèle et de la machine : les forêts aléatoires et les petits réseaux y gagnent beaucoup, car ONNX Runtime fusionne les opérations et parallélise les branches de l'arbre. Sur des modèles déjà ultra-optimisés, l'écart se réduit.

Pourquoi c'est plus rapide ?

Un modèle scikit-learn exécute sa prédiction à travers une pile Python : boucles, appels de fonctions, vérifications. ONNX Runtime lit un graphe de calcul statique et le compile en un plan d'exécution optimisé. Les opérations sont fusionnées, les constantes précalculées, et les threads mieux exploités. Tu ne changes ni la logique, ni les poids : tu changes le moteur.

Et après ?

Une fois le fichier .onnx en main, il se déploie partout : dans une API FastAPI, sur un mobile via ONNX Runtime Mobile, ou même dans le navigateur avec onnxruntime-web. Tu peux aussi quantifier le modèle pour le rendre plus léger. Le format ONNX est un passeport : ton modèle n'est plus prisonnier de scikit-learn.

La prochaine fois que ton inférence traîne, ne réécris pas tout. Exporte en ONNX, branche ONNX Runtime, et regarde ton CPU respirer.