Inbetriebnahme eines scikit-learn-Modells mit ONNX und FastAPI

Dieser Artikel befasst sich mit dem Deployment eines Machine-Learning-Modells, das den Wert eines Hauses in Boston anhand gewisser Merkmale wie der Kriminalitätsrate des Bezirks und der Anzahl der Räume in einer Wohnung bestimmen kann. Im ersten Schritt wird eine lineare Regression mit scikit-learn trainiert. Anschließend wird das Modell mit ONNX und FastAPI für den produktiven Einsatz vorbereitet.

scikit-learn: Übersicht

Übersicht der Schritte

Zunächst wird eine lineare Regression mithilfe von scikit-learn trainiert und in den Modellstandard ONNX konvertiert. Anschließend wird das Modell mit der Inference Engine ONNX Runtime und FastAPI als REST API zur Verfügung gestellt. Bei FastAPI handelt es sich um ein Web-Framework für Python, um Schnittstellen effizient zu entwickeln.

Modell-Training

Als Datensatz wird der im Scikit-Learn-Paket integrierte Boston-Housing-Datensatz verwendet. Es wird eine lineare Regression trainiert. Dazu werden die Daten in einen Trainings- und Testdatensatz aufgeteilt.

Die Konvertierung von scikit-learn nach ONNX wird mit dem Paket sklearn-onnx durchgeführt.

ONNX Runtime

ONNX Runtime ist eine Open Source Inference Engine und wurde von Microsoft veröffentlicht. Der Fokus des Frameworks liegt auf der Inferenz der Modelle. Für das Training ist die Bibliothek nicht geeignet. 

Bevor es die ONNX Runtime gab, mussten die ONNX-Modelle beispielsweise von ONNX nach TensorFlow exportiert werden, um die Inferenz durchzuführen. Dies hat den Nachteil, dass auf dem Weg zur Produktion ein weiterer Schritt benötigt wurde, um die Modelle zu konvertieren. Durch die ONNX Runtime ist dieser Schritt obsolet, da die Modelle direkt ausgeführt werden können.

Für die Verwendung der ONNX Runtime muss lediglich das onnxruntime -Paket installiert werden. 

Rest API mit FastAPI

FastAPI ist ein Web-Framework, das auf den Standard Python Type Hints aufbaut und erst seit Python 3.6+ verfügbar ist. Python Type Hints ist eine neue Syntax, die den Type einer Variable deklariert. Durch die Type Hints ist das Framework in der Lage, aus dem geschriebenen Code automatisch die OpenAPI-Dokumentation zu generieren. Zu dem Framework FastAPI gibt es den hervorragenden Artikel Introducing FastAPI .

Einer der Hauptvorteile des Frameworks ist, dass durch die Python Type Hints der Body der Requests direkt als Klasse definiert wird und dadurch Typ-Überprüfungen, JSON Parsing und vieles weitere direkt unterstützt werden. Dazu muss eine Klasse von BaseModel aus der Bibliothek pydantic erben, um diese Eigenschaften zu besitzen.

Für die Applikation wurden dazu die zwei Klassen HousingFeatures und PredictionResult anlegt. Die HousingFeatures-Klasse repräsentiert die Merkmale und die PredictionResult Klasse den bestimmten Wert.

Anschließend können wir den REST Endpoint /predict implementieren, der als POST Request erreichbar ist. Für die Absicherung des Endpoints wurde ein simpler Sicherheitsmechanismus implementiert, der den Token im Header prüft. Um die Komplexität zu reduzieren, wird hierbei ein festgelegter Token validiert. Der Sicherheitsmechanismus wird über Dependency Injection aufgerufen. 

Abschließend wurde die Anwendung in einen Docker-Container überführt. Der vollständige Code ist unter dem Repository scikit-onnx-fastapi-example verfügbar. Die Anwendung kann direkt mit docker-compose ausgeführt werden und ist unter http://localhost erreichbar. Mit dem Script test.sh wird der /predict Endpoint getestet. Unter http://localhost/docs ist die OpenAPI-Dokumentation verfügbar.

Ablauf POST /predict Request

Fazit

Die ONNX Runtime bietet eine einfache und unkomplizierte Möglichkeit, um Inference direkt mit dem ONNX-Modell auszuführen. Dadurch wird der Arbeitsaufwand deutlich reduziert, da das Modell für das Deployment nicht mehr in ein anderes Format exportiert werden muss. ONNX Runtime und FastAPI bringen Machine-Learning-Modelle schnell in Produktion. Ein weiterer Vorteil von FastAPI ist, dass die Schnittstellen direkt dokumentiert werden und somit die Anwendung schneller in den Unternehmenskontext eingebunden werden kann.