ONNX – Portabilität von Deep-Learning-Modellen

In den vergangenen Jahren ist die Anzahl an Frameworks für Deep Learning explodiert. Unternehmen wie Google, Facebook und Amazon haben ihre Deep Learning Frameworks TensorFlow , PyTorch und MXNet quelloffen zur Verfügung gestellt oder entwickeln aktiv dabei mit. Jedes dieser Frameworks bringt unterschiedliche Vor- und Nachteile mit, die verschiedene Folgen auf Entwicklung und Inbetriebnahme haben. In diesem Artikel wird Open Neural Network Exchange (ONNX)  vorgestellt. ONNX ist ein Standard, der es ermöglicht, Modelle zwischen den Frameworks auszutauschen und die Vorteile eines Frameworks je nach Situation besser zu nutzen.

Hintergrund

In den letzten Jahren wurde stark auf das Framework theano gesetzt. Heutzutage ist das Framework schon fast ausgestorben und wird nicht mehr weiterentwickelt. Derzeit ist es ungewiss, welche Frameworks sich etablieren und welche wieder verschwinden. Dies liegt under anderem auch daran, dass die Frameworks aus unterschiedlichen Situation heraus entwickelt wurden und spezifischer in der Forschung oder Produktion verwendet werden. Neben den Deep-Learning-Bibliotheken gibt es noch weitere numerische Frameworks, die ebenfalls eine Auswirkung auf die Entwicklung und Laufzeit der Modelle haben. Diese Frameworks werden von Hardware-Herstellern wie bspw. NVIDIA und INTEL entwickelt, und eigenen sich in bestimmten Anwendungsfeldern wie GPU- oder CPU-basiertes Rechnen. Unternehmen, die Deep Learning in das Tagesgeschäft implementieren wollen, sind von der Bandbreite an Möglichkeiten überfordert. Dabei kann die Auswahl eines Frameworks gravierende Auswirkungen auf verschiedene Unternehmensbereiche haben. Die Innovationsgeschwindigkeit kann deutliche Einbußen vernehmen, da sich eventuell nach der Modellentwicklung die  Inbetriebnahme eines Modell verzögert. Ein Grund dafür kann sein, dass das gewählte Framework stärker für die Entwicklung als für die Produktion ausgelegt ist.

Die untenstehende Grafik zeigt eine kleine Auswahl des Deep Learning Frameworks Zoo und dessen technischen Möglichkeiten. Ein Problem, das es generell unter den Frameworks gibt, ist die Portabilität der Modelle zu einem anderen Framework. Dadurch könnten je nach Phase, ob Entwicklung oder Inbetriebnahme, die Vorteile der unterschiedlichen Frameworks genutzt werden. Beispielsweise eignet sich PyTorch hervorragend für die prototypische Entwicklung und Experimentation der Modelle, während beispielsweise die TensorFlow-Serving -Komponente unkompliziertes Deployment der Modelle mit TensorFlow ermöglicht.

ONNX

2017 haben sich Microsoft, Facebook und Amazon vereint, um die Herausforderung der Portabilität der Modelle zu lösen. Entstanden dabei ist der neue Standard Open Neural Network Exchange (ONNX) . Die Vision hinter ONNX: ein Modell, das mit Framework A entwickelt wurde, zu exportieren und problemlos in ein Framework B zu importieren. Bis heute wird an dieser Vision festgehalten und eine Vielzahl an Features, die dies garantieren können, sind bereits implementiert. Eine Liste der derzeit unterstützten Frameworks und Tools kann hier gefunden werden. Hierbei sollte auch erwähnt sein, dass nur die Eigenschaften und Funktionen von einem Framework übertragen werden können, wenn diese bereits in einem Ziel-Framework existieren und über den ONNX-Standard ausgetauscht werden.

Bei der abstrakten Betrachtung der Arten an Deep-Learning-Bibliotheken ist der Datenfluss einer der wesentlichen Hauptunterschiede. Während bei TensorFlow und Caffe2 die Daten durch einen statischen Graphen fließen, wird bei PyTorch auf einen dynamischen Graphen gesetzt. In der Programmierung sowie zur Laufzeit kann dies zu einigen Unterschiede führen. Jedoch ist das unproblematisch für den ONNX-Standard. Durch die Interfaces der Bibliotheken können die relevanten Informationen wie Struktur und Gewichte extrahiert und transformiert werden. Die ONNX-Spezifikation besteht aus diesen drei wesentlichen Komponenten, die Import und Export ermöglichen:

1. Ein erweiterbaren Berechnungsgraphen
2. Festgelegte Operatoren und Funktionen
3. Festgelegte Standard Datentypen

Die exakte Definition kann in der Spezifikation auf Github onnx/onnx nnachgelesen werden und wird nicht weiter in diesem Artikel behandelt.

MNIST Beispiel

Um ONNX etwas näher kennenzulernen, schauen wir uns ein praktisches Beispiel mit PyTorch und TensorFlow an. Mit PyTorch wird ein Modell trainiert, das nach ONNX exportiert wird. Anschließend wird das ONNX transformierte Modell in TensorFlow geladen und eine Vorhersage getätigt. Als Datensatz wird MNIST verwendet. Um das Tutorial selbst durchzuführen, werden python3 und pip3 benötigt. Die benötigten Pakete können über pip3 installiert werden:

pip3 install torch torchvision tensorflow onnx onnx-tf

Zunächst wird das zu trainierende Netz in PyTorch definiert. Es wird ein Convolutional Neural Network verwendet, das im wesentlichen aus zwei Convolutional Layer und zwei Full Connected Layer besteht, die über die Aktivierungsfunktion ReLU und Max Pooling Layer verbunden sind. Als Input wird ein Bild mit nur einen Farbkanal erwartet. Anschließend werden die Trainings- und Testfunktion definiert. In der main()-Funktion werden die wesentliche Teile zusammengesteckt. Wichtig dabei ist, dass nach dem Training die Gewichte des Modells mit torch.save(model.state_dict(),  file) abgespeichert werden. Die Trainings-, Test- und main()-Funktionen können im Repository nachgelesen werden.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

Für den Export nach ONNX wird zunächst das trainierte Modell wieder in PyTorch eingelesen. Dazu werden die Gewichte geladen. Anschließend wird noch ein dummy_input angelegt, um die Eingangsvektoren des Modells festzulegen. Grund dafür ist, dass PyTorch durch den dynamischen Input keine Information für die exakte Größe des Eingangsvektor hat. Dies wird dennoch für die Repräsentation in ONNX benötigt.

from torch.autograd import Variable
trained_model = Net()
trained_model.load_state_dict(torch.load('output/mnist_cnn.pt'))
# Shape of input to the model
dummy_input = Variable(torch.randn(1, 1, 28, 28)) 
# Export the trained model to ONNX
torch.onnx.export(trained_model, dummy_input, "output/mnist.onnx")

Durch onnx.load(file) kann das Modell eingelesen werden. Über die prepare(model)-Methode des onnx/onnx-tensorflow -Pakets werden die Gewichte an den statischen Graphen gebunden.

import onnx
from onnx_tf.backend import prepare
model = onnx.load('output/mnist.onnx')
tf_rep = prepare(model)

Anschließend kann das Modell in der Laufzeitumgebung von TensorFlow für neue Daten verwendet werden. In diesem Fall muss das Bild noch auf 28×28 Pixel skaliert und nach Greyscale konvertiert werden. Danach kann die Umwandlung des Datentyps zu Float32 sowie die Transformation der Darstellung der Daten innerhalb des Arrays für das neuronale Netz erfolgen.

import numpy as np
from IPython.display import display
from PIL import Image
img = Image.open('images/five.png').resize((28, 28)).convert('L')
display(img)
output = tf_rep.run(np.asarray(img, dtype=np.float32)[np.newaxis, np.newaxis, :, :])
print('The digit is classified as ', np.argmax(output))

Grenzen von ONNX

Auf den ersten Blick ist der ONNX-Standard eine leicht zu bedienende Möglichkeit, die Portabilität der Modelle zu gewährleisten. Die Verwendung von ONNX ist sehr einfach, wenn die unterstützen Datentypen und Operationen der Spezifikation genutzt werden und es keine große Eigenentwicklung bei der Architekturgestaltung des neuronalen Netzes gab. Weiterhin sollte auch geprüft werden, welche der Operationen und Funktionen bereits in den Backends für den Export und Import der Frameworks implementiert sind. Hierbei muss angemerkt werden, dass beim ONNX-Projekt mit einer rasanten Geschwindigkeit entwickelt wird und ständig neue Releases veröffentlicht werden, die die Kompatibilität zwischen den Frameworks verstärken. Wird ein Projekt in diesem Rahmen durchgeführt, ist der Einsatz von ONNX völlig unproblematisch. Ist aber eine dieser Bedingungen nicht gegeben, müssen die Funktionalität in den ONNX Backends selbst implementiert werden, um diese zu verwenden. Dies kann sich als sehr zeitintensiv und mühsam herausstellen.

Zusammenfassung

Der Bedarf an der Modell-Portabilität ist größer denn je. Es gibt immer mehr und mehr Deep Learning Frameworks am Markt und durch die Portabilität können die Vorteile der einzelnen Frameworks besser genutzt werden. ONNX ist ein leicht zu verwendendes Framework, das viel Potenzial hat um der Standard für den Austausch der Modelle zwischen den Bibliotheken zu sein. Dadurch wird gewährleistet, entwickelte Modelle langfristig und flexible einzusetzen. Weiterhin können die Ergebnisse der Forschung schneller in Produktion gelangen, solange die unterstützten Datentypen und Operationen von ONNX genutzt werden. Andernfalls müssen diese in ONNX nachimplementiert werden.