Deep Java Learning - NDManager & NDArray

Wer es nicht abwarten kann, findet im Github Repository von DJL bereits eine Menge Komplettbeispiele, sowohl als Java Projekte als auch als interaktive Jupyter Notebooks.

Wir werden allerdings etwas tiefer einsteigen und zunächst mit den zwei essentiellsten Interfaces der DJL API beginnen: ai.djl.ndarray.NDManager, ai.djl.ndarray.NDArray. Beides sind Interfaces, die zur Laufzeit durch eine der zugrundeliegenden Engines implementiert werden. Meist wird es sich hierbei zur Zeit noch um Apache MXNet handeln, aber Implementierungen auf Basis von TensorFlow und PyTorch sind bereits in Arbeit.

Der Einstieg in die API: einen NDManager erzeugen

Der NDManager kümmert sich darum, Daten auf einem Device - oft der GPU zu verwalten. Zugriff auf diese Daten erhält man in Form von NDArray-Instanzen. Trainiert man ein neues DJL-Modell oder nutzt ein bestehendes, so wird der NDManager von den entsprechenden Hilfsklassen erzeugt. Möchte man für Testzwecke oder “nicht-Deep-Learning” Anwendungen direkt darauf zugreifen, so kann man ihn einfach wie folgt erzeugen:

NDManager manager = NDManager.newBaseManager();
    NDManager managerOnCPU = NDManager.newBaseManager(Device.cpu());

Bei der ersten Variante wählt DJL ein sogenanntes Device aus, auf dem die Operationen ausgeführt werden – in der Regel die erste verfügbare GPU oder ansonsten die CPU, falls keine GPUs nutzbar sind. Möchte man ein ganz spezielles Device manuell auswählen, nutzt man die zweite Variante.

Die wichtigste Klasse der DJL API: NDArray

Möchte man nun Berechnungen ausführen, muss man die Werte, mit denen man rechnen möchte, in NDArrays packen. Um einen neuen NDArray zu erzeugen, braucht man einen NDManager. Dieser legt die Daten dann auf seinem Device außerhalb des Java-Heaps an und verwaltet den dafür benötigten Speicher:

NDArray pi        = manager.create((float)Math.PI);
    NDArray e         = manager.create(Math.E);
    NDArray one       = manager.create((byte)1);
    NDArray theAnswer = manager.create(42);
    NDArray big       = manager.create(Long.MAX_VALUE);
    NDArray isTrue    = manager.create(true);

Die einfachste Methode, einen NDArray zu erzeugen, ist einen einzelnen Wert in einem NDArray zu wrappen. Dies kann ein Java-Primitiv sein, oder eine Klasse, die Number implementiert, wie Integer oder Float. Im Gegensatz zu z.B. java.util.List ist NDArray nicht generisch, wir können also nicht am Typ erkennen, was für Daten gespeichert sind. Während man also eine List<Float> erzeugen kann, gibt es keinen NDArray<Float>. Um herauszubekommen, was der Typ im NDArray gespeicherten Daten ist, gibt es die Methode NDArray.getDataType(). Das sind die Datentypen, die die zuvor erzeugten NDArrays angeben:

System.out.println(pi.getDataType());        //float32
    System.out.println(e.getDataType());         //float64
    System.out.println(one.getDataType());       //int8
    System.out.println(theAnswer.getDataType()); //int32
    System.out.println(big.getDataType());       //int64
    System.out.println(isTrue.getDataType());    //boolean

Die möglichen Datentypen eines NDArray findet man im enum ai.djl.ndarray.types.DataType. Die meisten NDArray-Datentypen entsprechen 1:1 einem Java Primitiv:

• float → DataType.FLOAT32
• double → DataType.FLOAT64
• byte → DataType.INT8
• int → DataType.INT32
• long → DataType.INT64
• boolean → DataType.BOOLEAN

Der Datentyp des erzeugten NDArray hängt also vom Java-Datentyp ab, der an die create-Methode übergeben wurde. Es gibt allerdings auch zwei Datentypen, die keine Java-Entsprechung haben: UINT8 (ein unsigned byte) und FLOAT16 (ein float Wert mit geringerer Präzision; weniger genau, spart aber dafür Speicher, der auf Grafikkarten auch mal knapp sein kann). Um NDArrays dieses Typs zu erzeugen, muss man erst einen Array eines anderen Typs erzeugen und dann den Datentyp manuell konvertieren:

NDArray pi16 = pi.toType(DataType.FLOAT16, true);

Der zweite Parameter, copy gibt dabei an ob der bestehende NDArray modifiziert wird oder ob man eine neue Kopie erhält und der alte NDArray erhalten bleibt.

Weitere Möglichkeiten zur Erzeugung von NDArrays

Es gibt eine Reihe weiterer Möglichkeiten, einen NDArray zu erzeugen. Praktisch alle sind Memberfunktionen des NDManagers. Die wichtigste Methode ist – wie oben – die create Methode. Allerdings akzeptiert diese natürlich nicht nur einzelne Werte, sondern auch Arrays von Java Primitiven und Number Instanzen. Sehr häufig wird man also NDArrays aus ein- oder zweidimensionalen float[] oder int[] Arrays erzeugen.

Dazu gibt es noch die Methoden NDManager.arange und NDManager.linspace, mit denen man Sequenzen von Zahlen als NDArray erzeugen kann, z.B. 0, 1, 2, 3 oder 0.0, -0.1, -0.2, -0.3. Dabei sind Startwert, Endwert und Schrittweite einstellbar. Dies ist sehr nützlich, um schnell einige Testdaten zu erstellen, aber auch zum Beispiel bei sehr kleinteiligen Berechnungsschritten in einem neuronalen Netz Offsets für Input Daten zu erzeugen.

Mit NDManager.ones und NDManager.zeros kann man NDArrays beliebiger Größe erzeugen, die mit Einsen oder Nullen befüllt sind. In der Praxis sehr wichtig sind schließlich die Methoden, mit denen man mit Zufallszahlen gefüllte NDArrays erzeugt. Mit NDManager.randomNormal, NDManager.randomUniform und NDManager.randomMultinomial kann man Zufallszahlen mit den entsprechenden Wahrscheinlichkeitsverteilungen erzeugen. Dies ist besonders wichtig für neuronale Netze, denn diese müssen erst zufällig initialisiert werden, bevor sie trainiert werden können.

Berechnungen auf NDArrays

Nun, da wir Daten so verpackt haben, das DJL damit arbeiten kann, können wir auch mathematische Operationen ausführen:

System.out.println(pi.sin().getFloat()); //-8.742278E-8

Sämtliche Berechnungen werden nun nativ auf dem Device des zugrundeliegenden NDManagers ausgeführt. Bei der Berechnung eines einzelnen Wertes ist dies natürlich weder spannend noch nützlich. Das Hin und Her zwischen GPU und JVM ist langsamer und aufwändiger, als alles einfach in Java zu berechnen. Spannend wird es, wenn wir sehr viel auf einmal zu berechnen haben. Zum Testen erzeugen wir 100 Millionen Zufallszahlen:

float[] random = new float[1000 * 1000 * 100];
    Random rand = new Random();
    for (int i = 0; i < random.length; ++i) {
        random[i] = rand.nextFloat();
    }

Nun berechnen wir den Sinus jeder dieser Zahlen in Java:

float[] sines1 = new float[random.length];
    for (int i = 0; i < random.length; ++i) {
        sines1[i] = (float)Math.sin(random[i]);
    }

Auf einem unserer Arbeitslaptops dauert dies ca 3s. Nun führen wir die gleiche Berechnung mittels DJL auf der GPU durch:

NDArray randOnGpu = manager.create(random);
    float[] sines2 = randOnGpu.sin().toFloatArray();

Dies dauert ca. 500ms, ist also sechsmal so schnell. In der Regel sind Berechnungen mit DJL sogar noch um einen wesentlich größeren Faktor schneller als in Plain Java. Der Hauptzeitfresser in unserem Beispiel ist der Transfer von und zur Grafikkarte. Bleibt man auf der GPU und führt viele Operationen in Folge durch, wird der relative Zeitgewinn gegenüber einer unbeschleunigten Lösung immer größer.

Die Form von NDArrays – Shape

Wichtig bei der Nutzung von NDArrays gegenüber normaler Arrays ist aber nicht nur die höhere Geschwindigkeit, sondern auch der wesentlich lesbarere Code. Alle Operationen werden „vektorisiert” ausgeführt, dass heißt mit allen Elementen auf einmal. Bei einer Operation wie sin() kann man sich dies noch leicht vorstellen, da man für einen Sinus nur eine Eingabe braucht – die Operation wird einfach auf jedem Element des Arrays wiederholt.

Spannend wird es bei Operationen, bei denen NDArrays kombiniert werden, z.B. bei einer einfachen Addition (das Ergebnis ist immer im Kommentar über dem Aufruf angegeben):

// I. 4
    manager.create(2).add(manager.create(2));
    // II. [10, 12, 14, 16, 18, 20, 22, 24]
    manager.arange(0, 8).add(manager.arange(10, 18));
    // III. [ 2,  3,  4,  5,  6,  7,  8,  9]
    manager.arange(0, 8).add(manager.create(2));
    // IV.
    // [[ 100, 1001],
    //  [ 102, 1003],
    //  [ 104, 1005],
    //  [ 106, 1007],
    // ]
    manager.arange(0, 8).reshape(4, 2)
        .add(manager.create(new int[]{100, 1000}));

Das erste Beispiel ist wenig überraschend: 2 + 2 = 4. Das zweite ist schon interessanter: Man kann zwei Arrays einfach mit einem Aufruf addieren, es werden jeweils die Elemente miteinander addiert (dies entspricht einer Vektoraddition). Das dritte Beispiel wird noch interessanter: Es zeigt, dass die NDArrays nicht unbedingt die gleiche Größe haben müssen. Addiert man einen einzelnen Wert, wird er zu allen Element des ersten NDArray addiert. Richtig spannend wird es dann in Beispiel IV. Hier sehen wir eine neue, wichtige Operation auf Arrays, reshape. Lässt man sie in diesem Beispiel weg, crasht der Code. Aber was macht reshape und wie kommt das Ergebnis zustande?

Bisher haben wir gelernt, dass ein NDArray einen Datentyp hat (z.B. FLOAT32) und eine Größe (die Anzahl der Elemente im Array). Aber ein NDArray hat auch eine Form. Die Form bestimmt, wie Rechenoperationen, die Arrays miteinander kombinieren, mit dem Array umgehen müssen. Im obigen Beispiel, wird die Zahlenreihe [0, 1, … , 7] mittels reshape in eine neue Form gebracht. Es ist nicht mehr eine Zahlenreihe (ein Vektor), sondern eine Reihe von Reihen (eine Matrix). Der Aufruf reshape(4, 2) bedeutet, dass die existierende Reihe in vier Stücke der Länge zwei zerlegt werden soll. Damit das klappt, muss die resultierende Form die gleiche Größe wie die ursprüngliche haben. Da 2 * 4 = 8, ist das hier kein Problem. Da wir aber nun am „Ende“ des NDArray immer nur Reihen der Länge zwei haben, kann nun eine andere Reihe der Länge zwei addiert werden. Zwar gibt es nur eine davon, aber die wird dann einfach jedes mal benutzt. Dieses Verhalten nennt sich broadcasting und ist ein essentielles Feature aller Deep Learning Frameworks.

Wenn man einmal nicht weiß, welche Form ein NDArray hat, kann man dies jederzeit mit getShape herausfinden. Das Umformen von NDArrays und das richtige miteinander verknüpfen von NDArrays unterschiedlicher Formen ist eine der wichtigsten und kniffligsten Aufgaben bei der Programmierung von Deep Learning System. Für die angehende Java-Deep-Learning-Expertin ist es wichtig, das Broadcasting Verhalten für eine Reihe von wichtigen Operationen wie add, sub, mul, dot, matMul etc. pp. zu kennen, um so effektiv und elegant Formeln und Pseudocode in eine Verkettung von NDArray Operationen zu übertragen.

Die Speicherverwaltung von NDArrays

Nun, da wir wissen, wie man NDArrays erzeugt und benutzt, bleibt nur noch, nach getaner Arbeit wieder hinter uns aufzuräumen. Wie bereits erwähnt, sind NDArrays Platzhalter für Daten auf einem vom NDManager genutzten Device. Der Speicher dieses Devices kann allerdings nicht vom Memory Manager der JVM verwaltet werden, daher müssen wir uns hier selber darum kümmern. Jeder initiale NDArray muss wieder geschlossen werden (wie bei Streams mit .close()) , damit der zugrunde liegende native Speicher z.B. auf der GPU wieder verfügbar wird.

Dies könnte man nun natürlich mit jedem Array einzeln machen, am besten mit try…finally Blöcken. Allerdings erzeugen Operationen auf NDArrays ebenfalls neue NDArrays im nativen Speicher des Device. Addieren wir zwei Arrays, entsteht ein neuer für das Ergebnis. (Ausgenommen sind hier einige spezielle Varianten der Operationen, die Veränderungen im NDArray vornehmen, wie addi. Diese haben das Suffix -i für „in place“) . All diese Zwischenergebnisse zu schließen, kann schnell mühselig werden. Aber dafür gibt es zum Glück eine einfache Lösung: Alle diese NDArrays sind mit einem NDManager verknüpft, über den sie direkt oder indirekt erzeugt wurden. NDManager selbst implementiert aber auch AutoClosable! Das Schließen des Manager bewirkt nun wiederum das Schließen aller von ihm „abstammenden“ NDArrays, so dass man bequem mit einer Operation alle Speicher aufräumen kann.

Was aber, wenn man nicht alle NDArrays schließen will, sondern nur die, die z.B. bei einer Zwischenrechnung angefallen sind? Das ist ebenfalls ganz einfach: Man kann mit NDManager.newSubManager() einen „Submanager“ erzeugen, der sich wie der ursprüngliche Manager verhält, aber nicht dessen NDArrays „erbt“. Mit diesem Submanager kann man nun Berechnungen durchführen, und dann nur den Submanager schließen. Der ursprüngliche Manager und dessen Arrays bleiben dann erhalten.

Fazit

In dieser Einleitung haben wir gesehen, wie man die grundlegendsten Klassen der DJL API benutzt: NDManager und NDArray. Im nächsten Post wollen wir dann den nächsten Schritt Richtung Deep Learning machen und Daten für unser erstes Beispiel so laden, dass sie für DJL zum Training nutzbar sind. Dafür werden wir NDArrays erzeugen, befüllen und umformen müssen, sowie die ersten Berechnungen anstellen, damit unsere Daten für ein neuronales Netz auch zu „verdauen“ sind.