Skalierbare Bildklassifizierung mit ONNX und AWS Lambda

In meinem Blogartikel ONNX – Portabilität von Deep-Learning-Modellen haben wir bereits ONNX kennengelernt und was es damit auf sich hat. Zur Erinnerung: ONNX ist ein Open Source geführter Standard, mit dem Modelle zwischen verschiedenen Deep-Learning-Bibliotheken serialisiert werden können. Wie sieht es jetzt aber aus, nachdem wir ein Modell trainiert haben und dieses skalierbar verwenden wollen? Eine Möglichkeit ist, das ONNX Modell in einem AWS Lambda zu deployen.

Die Situation

Wir stellen uns vor, dass wir als Anbieter eines Online-Versandhauses unsere Artikel anhand der dazugehörigen Bilder klassifizieren wollen. Damit wollen wir die Artikel mithilfe der Klassifikation mit weiteren Stichwörtern anreichern, so dass Nutzer anhand der Stichwörter diese Artikel auch finden können. Im ersten Schritt benötigen wir ein Modell, das in der Lage ist, Bilder zu klassifizieren. Im zweiten Schritt müssen wir uns Gedanken machen, wie wir das Modell in die Produktion bekommen und skalierbar verwenden können. Die Produktionsumgebung muss in der Lage sein, die bereits vorhandenen Artikel sowie neu erstellte Artikel zu klassifizieren. Dies bedeutet, dass die genaue Auslastung der Server nicht vorhergesagt werden kann.

Das Modell

Für die Entwicklung des Modells gibt es zwei Ansätze. Der erste Ansatz ist ein Modell von Grund auf mit den bereits vorhandenen Daten zu trainieren. Einige der Artikel verfügen bereits über verschiedene Tags, die als Label dienen können. Allerdings müssen diese Tags vor dem Training verifiziert werden, um eine gewisse Datenqualität zu erhalten und Inkonsistenzen auszuschließen. Sowohl diese Vorgehensweise als auch das ressourcenintensive Training des Modells ist mit sehr hohem Zeitaufwand verbunden. Da wir noch nicht wissen, wie stark dieses Feature genutzt wird, könnte dies auch darauf hinauslaufen, dass zu viel Zeit investiert wird.

Der zweite Ansatz besteht darin, ein vorhandenes Modell zu verwenden, das bereits auf einer Vielzahl von Klassen trainiert wurde. Dies hat den Vorteil, dass zunächst keine Daten erforderlich wären und die Funktionalität des Modells direkt begutachtet werden kann. Ein mögliches Modell wäre ResNet, das 2015 auf dem jährlich stattfindenden Software Wettbewerb ILSVRC u.a. in der Kategorie Bildklassifizierung gewonnen hat. Es gibt bereits eine Variante von ResNet, die auf der ImageNet Datenbank trainiert wurde. Diese umfasst Bilder aus über 20.000 Kategorien. Im Framework PyTorch ist bereits ein trainiertes ResNet-Modell vorhanden, das auf 1.000 Klassen trainiert wurde.

Das Modell kann nach ONNX konvertiert werden, um es für die produktiven Einsatz zu nutzen. Dazu muss die Forward-Funktion aufgrund der von ONNX unterstützten Operationen leicht verändert werden. Die Änderungen wurden in der FixResNet50 Klasse implementiert. Anschließend kann das Modell nach ONNX konvertiert werden.

model = FixResNet50(models.resnet.Bottleneck, [3, 4, 6, 3])
from torch.utils.model_zoo import load_url as load_state_dict_from_url
state_dict = load_state_dict_from_url(model_url,
                                    progress=True)
model.load_state_dict(state_dict)
from torch.autograd import Variable
dummy_input = Variable(torch.randn(1, 3, 224, 224))
torch.onnx.export(model, dummy_input, "./model/resnet.onnx")

Inferenz mit ONNX.js

Nun stellt sich die Frage, wie wir effizient mit dem Modell Vorhersagen tätigen können. Dazu verwenden wir ONNX.js – eine Javascript Library, um ONNX Modelle in Node.js und im Browser ausführen zu lassen. Das Framework lässt sich sowohl auf der CPU mit WebAssembly und der GPU mit WebGL ausführen. Ermöglicht wird dies durch die Implementierung des ONNX Standards. Durch ONNX.js haben wir eine Runtime, um die Modelle auszuführen.

const session = new onnx.InferenceSession();
await session.loadModel("./model/resnet.onnx");
const prediction = await session.run([inputTensor]);

AWS Lambda in der Produktion

AWS Lambda ist ein Service, der von Amazon Web Services (AWS) angeboten wird, um Code beim Eintreten von festgelegten Events wie bspw. einem REST Request auszuführen. Der Entwickler muss lediglich den Code schreiben, während die Inbetriebnahme und Skalierung von AWS übernommen wird. Dies wird heutzutage als serverless bezeichnet.

Serverless bedeutet, dass sich die Verwaltung der Server, auf denen die Anwendungen laufen, an einen Anbieter wie bspw. AWS abgegeben ist. Im Klartext heißt dies für den Entwickler, dass die technischen Aufgaben – wie bspw. das Konfigurieren der Firewalls, Aufsetzen von virtuellen Maschinen und Patchen des Betriebssystem – nicht verwaltet werden müssen. Neben der Ausführung wird von Lambda Service die Skalierung und damit die Verfügbarkeit der Funktion übernommen. Mein Kollege Dennis Traub hat einen ausführlichen Artikel über AWS Lambda verfasst.

Im Vergleich zur herkömmlichen Skalierung von Server hat AWS Lambda den Vorteil, dass die Kosten nur für die Nutzung entstehen, nicht aber für die Idle-Zeit der Server. In unserem Anwendungsfall ist AWS Lambda hervorragend geeignet, da für die anfängliche Batch-Klassifizierung die Lambdas nach Bedarf individuell skalieren. Nach der Batch-Klassifizierung fallen lediglich Kosten an, sobald wir neue Artikel aufnehmen.

Lambdas mit dem Serverless Framework

Ein Lambda kann in verschiedenen Runtime-Umgebungen deployed werden. In unserem Fall entwickeln wir mit Node.js um ONNX.js zu verwenden. Für die Programmierung sowie das Deployment verwenden wir die Third Party Library Serverless . Das Serverless Framework ist eines der bekanntesten Toolkits, um Serverless-Anwendungen zu entwickeln, und es vereinfacht durch die Automatisierung den Deploymentprozess des Lambdas.

Die Bibliothek wird über npm installiert. Durch den Command serverless create –template aws-nodejs –path image_classifier wird ein Blueprint für das Projekt erstellt. Im Verzeichnis image_classifier werden die zwei Dateien serverless.yml und handler.js erstellt. In der YML-Datei werden die Lambda-Funktionen registriert sowie die Konfigurationen hinterlegt.

Die Anwendung

Nachdem die Tools zur Realisierung des  Bildklassifizierungs-Lambda vorgestellt wurden, widmen wir uns in diesem Teil der Implementierung. Um den Anwendungsfalls etwas einfacher zu gestalten, wird das Feature als ein POST Request implementiert, der im Body das Bild als Base64 encoded enthält. Die Response beinhaltet die Top-5-Klassen mit der höchsten Wahrscheinlichkeit und dem dazugehörigen realen Label. Folgende Schritte müssen implementiert werden:

1. Base64 Decoding des Bildes
2. Resizing des Bildes für das Modell
3. Normalisierung der Pixel
4. Klassifizierung
5. Erstellung der Response

Die Schritte 1 und 2 sind hier implementiert. Das Ergebnis nach Schritt 2 ist ein ImageData Objekt. Die Daten sind in einem eindimensionalen Array nach RGBA angeordnet. Für die Normalisierung der Bilder transformieren wir die Daten in ein mehrdimensionales Array mithilfe der ndarray Bibliothek.

Zunächst wird aus ImageData ein ndarry erstellt. Neben der Weite und Höhe müssen wir zusätzlich noch den Stride von 4 angeben, der die Anzahl an Kanäle RGBA beschreibt. Die Pixelwerte werden zwischen 0 und 1 normalisiert, indem durch 255 geteilt wird. Zusätzlich werden die RGB Kanäle mit den Mittelwerten [0.485, 0.456, 0.406] und den Standardabweichungen [0.229, 0.224, 0.225] normalisiert. Dazu wird der Mittelwert von dem jeweiligen Farbkanal subtrahiert und durch die Standardabweichung geteilt.

const dataFromImage = ndarray(new Float32Array(data), [width, height, 4]);
ops.divseq(dataFromImage, 255);
ops.subseq(dataFromImage.pick(0, null, null), 0.485);
ops.divseq(dataFromImage.pick(0, null, null), 0.229);
ops.subseq(dataFromImage.pick(1, null, null), 0.456);
ops.divseq(dataFromImage.pick(1, null, null), 0.224);
ops.subseq(dataFromImage.pick(2, null, null), 0.406);
ops.divseq(dataFromImage.pick(2, null, null), 0.225);

Nach diesen Schritten verfügen wir über ein Array mit der Form [224,224,4]. Dieses muss nach [1,3,224,224] für das Modell transformiert werden.

const dataProcessed = ndarray(new Float32Array(width * height * 3), [1, 3, height, width]);
ops.assign(dataProcessed.pick(0, 0, null, null), dataFromImage.pick(null, null, 0));
ops.assign(dataProcessed.pick(0, 1, null, null), dataFromImage.pick(null, null, 1));
ops.assign(dataProcessed.pick(0, 2, null, null), dataFromImage.pick(null, null, 2));

Anschließend wird das Modell mit ONNX.js in den Speicher geladen. Für die Berechnung werden die Daten in einen Tensor transformiert. Nach der Klassifizierung erhalten wir ein Array mit den Wahrscheinlichkeiten aller möglichen Klassen. Über eine Hilfsfunktion werden die Wahrscheinlichkeiten anschließend in die realen Klassennamen überführt. Im handler.js werden die Codestücke zusammengeführt.

const input = require("./lib/input");
const process = require("./lib/process");
const inference = require('./lib/inference');

module.exports.predict = async (event) => {
 let body = JSON.parse(event.body);
 let base64String = body.image;
 let img = await input(base64String);
 let preprocessedData = await process(img.data, 224, 224);
 let predictions = await inference.predict(preprocessedData);
 let output = await inference.formatOutput(predictions);
 return {
   statusCode: 200,
   body: JSON.stringify({output}, null, 2),
 };
};

Die Lambda-Funktion wird anschließend über die Shell durch den Befehl serverless deploy in AWS deployed. Die vollständige serverless.yml ist hier einzusehen. Weiterhin verwenden wir einen Docker Container, um die Dependencies in einer Lambda ähnlichen Umgebung zu installieren. Nach dem erfolgreichen Deployment wird in der Konsole eine URL ausgegeben, unter der die Lambda-Funktion erreichbar ist. Das test_lambda.sh Script zeigt exemplarisch wie die Funktion aufgerufen werden kann. Der vollständige Code kann im GitHub Repository nachgelesen werden.

Fazit

In diesem Artikel haben wir uns die Inbetriebnahme eines ONNX Modells auf einem AWS Lambda angeschaut. Dazu wurde im ersten Schritt ein trainiertes ResNet von PyTorch nach ONNX konvertiert. Das konvertierte Modell wird anschließend mit ONNX.js verwendet, um die Inferenz auszuführen. Die Berechnungen erfolgen über WebAssembly oder WebGL. Das Serverless Framework hilft uns, die Node.js-Anwendung reibungslos als AWS Lambda zu deployen. Durch AWS Lambda skaliert die Anwendung nach Bedarf und wir müssen uns keine Gedanken um den Betrieb der Applikation machen.

AWS Lambda eignet sich je nach Use Case als eine hervorragende Alternative zu EC2 Instanzen, um Inferenz auszuführen. Dazu muss beachtet werden, dass durch das RAM Limit (<3GB) von AWS Lambda nicht jedes Modell in Betrieb genommen werden kann. Eine weitere Einschränkung ist der ONNX Standard, der nicht unbedingt jede Operation und Datentyp unterstützt. Dadurch ist es möglich, dass einige Modelle nicht konvertiert werden können.