MLflow – Tracking von Parametern und Metriken

Was ist MLflow?

MLflow ist ein Open Source Tool, um typische Workflows von Machine Learning (ML) zu verwalten. Durch MLflow werden verschiedene Bereiche der Teamzusammenarbeit in ML-Projekten standardisiert und damit erheblich vereinfacht. Es stellt die folgenden Komponenten bereit:

Tracking von Parametern und Metriken für Experimente
Reproduzierbarkeit von ML-Projekten
Modellversionierung und Modellregistry
Model Serving

Mit diesen Funktionalitäten gibt MLflow Lösungen auf typische Problemstellungen in einem ML-Projekt: Wie können verschiedene Trainingsdurchläufe unter unterschiedlichen Bedingungen reproduzierbar gemacht werden? Wie teilt man trainierte Modelle für andere Teammitglieder? Wie kann man schnell eine zentral verwaltete Datenbank für Metadaten wie Parameter und Metriken für ein Team schaffen? Wie kann ich mein ML-Modell schnell als alleinstehenden Service ausrollen?

Weitere Aspekte von MLOps sind in meinem vorherigen Blogpost zu finden. In diesem Beitrag widmen wir uns nur dem Thema Tracking von Parametern und Metriken, da dieses Gebiet oft vernachlässigt wird, aber gleichzeitig unheimlich wichtig ist. Mit MLflow kann man diese Aufgabe sehr schnell und ohne großen Aufwand umsetzen, wie wir im nächsten Kapitel sehen werden.

Durch die ständige Weiterentwicklung der aktiven Community von MLflow kommen immer mehr Funktionen hinzu. Zudem bietet es viele Integrationen mit Deep Learning Frameworks wie TensorFlow, Keras oder PyTorch, mit Spark, Kubernetes, Docker und Cloudanbietern wie Amazon, Microsoft und Google. Das macht es zu einem sehr flexiblen Werkzeug. Des Weiteren kann es in bereits bestehenden Code integriert werden.

Komponenten

Für das Tracking von Experimenten benötigt MLflow zwei Komponenten. Das Backend-Store ist eine Datenbank, in der Metadaten zu Experimenten wie Parameter und Metriken sowie Modellversionen persistent gespeichert werden. Es können verschiedene Datenbanken wie SQLite, MySQL oder PostgreSQL als Backend verwendet werden. Im Artifact-Store liegen Artefakte wie beispielsweise Modelle oder vorverarbeitete Daten. Dafür können unterschiedliche Möglichkeiten wie ein lokales Dateisystem, Amazon S3 Storage, Azure Blob Storage, Google Cloud Storage und noch weitere genutzt werden.

Quickstart

Für das einfachste Setup von MLflow verwenden wir eine SQLite Datenbank als Backend und den lokalen Ordner „mlruns“ als Artifact-Store. Im zugehörigen Repository zu diesem Blogpost ist noch eine weitere Variante durch ein Docker-basiertes MySQL Backend gegeben. Beispielsweise wäre eine solche Konfiguration mit einem Cloud-Storage als Artifact-Store in einer produktiven Umgebung eher geeignet.

Trackingserver

Als erstes muss ein MLflow Trackingserver gestartet werden. Dieser bringt ein fertiges Dashboard mit, in dem wir unsere Experimente und alle Metadaten vergleichen und visualisieren können. Um den Trackingserver zu starten, muss vorher MLflow über pip installiert werden:

pip install mlflow

1	pip install mlflow

Jetzt können wir den Trackingserver mit einem Befehl starten:

mlflow server --backend-store-uri sqlite:///mlflow.db --default-artifact-root ./mlruns --host 0.0.0.0

1	mlflow server --backend-store-uri sqlite:///mlflow.db --default-artifact-root ./mlruns --host 0.0.0.0

Als Parameter geben wir die beiden URIs (Uniform Resource Identifier) für das SQLite Backend-Store und das lokale Artifact-Store sowie die lokale IP-Adresse an.

MLflow Python API

Als weitere Voraussetzung muss in unserem Python-Script die MLflow-Bibliothek importiert und die URI für den Trackingserver und das Backend konfiguriert werden:

import mlflow

registry_uri = 'sqlite:///mlflow.db'
tracking_uri = 'http://localhost:5000'

mlflow.tracking.set_registry_uri(registry_uri)
mlflow.tracking.set_tracking_uri(tracking_uri)

import mlflow

registry_uri = 'sqlite:///mlflow.db'

tracking_uri = 'http://localhost:5000'

mlflow.tracking.set_registry_uri(registry_uri)

mlflow.tracking.set_tracking_uri(tracking_uri)

Alle Voraussetzungen sind nun geschaffen, um Experimente mit MLflow zu dokumentieren.

Im folgenden Codeausschnitt starten wir den Run eines Experiments und speichern darunter den Parameter „batch_size“ mit dem Wert 32 und die Metrik „mse“ mit dem Wert 0.1 ab:

with mlflow.start_run() as run:
    mlflow.log_param('batch_size', 32)
    mlflow.log_metric('mse', 0.1)

with mlflow.start_run() as run:

mlflow.log_param('batch_size', 32)

mlflow.log_metric('mse', 0.1)

MLflow kümmert sich darum, wie die Daten in der Datenbank abgelegt werden. Als nächstes betrachten wir ein etwas ausführlicheres Beispiel: die Vorhersage einer Zeitreihe mit einem LSTM-basierten Modell.

Tracking von Metadaten am Beispiel einer Zeitreihenvorhersage mit einem LSTM-basierten RNN

Der vollständige Code für dieses Beispiel ist in diesem Repository abgelegt. Die Funktionsweise von Recurrent Neural Networks wird hier nicht erläutert, sondern kann z.B. in diesem Blogpost nachgelesen werden. Als Datengrundlage bedienen wir uns einem bekannten öffentlichen Datensatz, der den Stromverbrauch von Haushalten enthält. Dieses Merkmal benutzen wir als Zielvariable, die von unserem ML-Modell vorhergesagt wird. Als Eingabe dient ein festgelegtes Zeitfenster mit 7 Spalten bzw. Merkmalen. Die detaillierten Beschreibungen der Merkmale können der Originalseite entnommen werden. Ausführliche Kommentare sind sowohl im Jupyter Notebook als auch im Pythonscript des Repositorys zu finden.

In diesem Beispiel werden wir eine RNN-Architektur mit 2 LSTM-Schichten unter verschiedenen Hyperparametern mit TensorFlow 2 trainieren und die Metadaten dokumentieren. Für den Anfang starten wir mit 4 Konfigurationen: 20 und 40 Neuronen für die LSTM-Layer und die Werte 0 bzw. 0.2 für die Dropoutrate des letzten Fully-Connected-Layers. Wenn das Pythonscript einmal ausgeführt wurde, sind im MLflow Dashboard die 4 „Child Runs“ zu sehen:

MLflow Dashboard: Übersicht der Runs

Die Parameter enthalten die Batch-Größe, Dropoutrate, wie viele Zeitschritte in die Zukunft vorhergesagt werden (future_length), Neuronenanzahl (hidden_layer_size), die Anzahl der vorhergesagten Merkmale (n_output_features) und die Fensterlänge des Inputs (window_length). Bei den Metriken werden während des Trainings der Mean Absolute Error (MAE) und der Mean Squared Error (MSE) als oft verwendete Messgrößen bei einer Regression geloggt. Dafür kann ein simpler benutzerdefinierter TensorFlow Callback verwendet werden, der alle vorhandenen Metriken im Training innerhalb der „model.fit()“-Methode trackt:

class MlflowLogging(tf.keras.callbacks.Callback):
    def __init__(self, **kwargs):
        super().__init__()  # handles base args (e.g., dtype)

    def on_epoch_end(self, epoch, logs=None):
        keys = list(logs.keys())
        for key in keys:
            mlflow.log_metric(str(key), logs.get(key), step=epoch)

class MlflowLogging(tf.keras.callbacks.Callback):

def __init__(self, **kwargs):

super().__init__() # handles base args (e.g., dtype)

def on_epoch_end(self, epoch, logs=None):

keys = list(logs.keys())

for key in keys:

mlflow.log_metric(str(key), logs.get(key), step=epoch)

Nach dem Training werden die Modelle in der Modelregistry registriert und versioniert. Alle Metadaten befinden sich in der SQLite Datenbank und die Modelle liegen in dem lokalen Artifact-Store unter dem Ordner „mlruns“ mit der folgenden Struktur:

MLflow Artifact-Store Struktur

Der Ordner „0“ steht für das Experiment mit der ID 0. Darunter wird für jeden Run eine Hash-ID generiert und ein entsprechender Ordner erstellt. In diesem werden alle zugehörigen Artefakte gespeichert. In unserem Fall wird ein TensorFlow Modell im sogenannten SavedModel Format mit den Gewichten und dem Metagraph dort abgelegt. Zusätzlich existieren noch weitere Metadateien wie eine conda.yaml, die durch MLflow automatisch erstellt werden. Sie halten die Versionen der benötigten Bibliotheken fest, sodass mit der MLflow Model Serving Funktionalität dieses Modell als alleinstehender Service beispielsweise in einem Docker-Container ausgerollt werden kann.

Des Weiteren stellt das Dashboard auch die Möglichkeit bereit, Runs oder Modelle zu vergleichen:

MLflow Dashboard: Vergleich verschiedener Modelle

Wenn sich Werte in einer Zeile wie die Neuronenanzahl und die Dropoutrate unterscheiden, werden sie farblich markiert. Als Schlussfolgerung sehen wir, dass das Modell Version 3 mit 40 Neuronen und einer Dropoutrate von 0 den geringsten MSE und MAE aufweist.

Fazit

In diesem Beispiel haben wir gesehen, dass mit sehr wenig Aufwand das Tracking der Metadaten von verschiedenen Trainingsdurchläufen realisiert wurde. Zudem stellt MLflow ein fertiges Dashboard zur Visualisierung zur Verfügung.

Wir sehen uns in einem neuen Blogpost der MLOps-Serie über ML-Pipelines.