Beim Trainieren von Machine-Learning-Modellen entsteht schnell eine Vielzahl an Varianten: unterschiedliche Hyperparameter, Trainingsdaten, Modelle und Evaluationsmetriken. Um den Überblick zu behalten, ist Experiment Tracking unerlässlich. Dieser Beitrag zeigt wie es funktioniert mit MLflow und Vertex AI Experiments.
MLflow
Eine einfache leichtgewichtige und sowohl in der Cloud als auch lokal nutzbare Möglichkeit Experimente zu tracken ist das Tool MLflow.
Um MLflow zu nutzen, kann man es einfach mit pip installieren:
pip install mlflowLokale Nutzung
Nach der Installation kann MLflow direkt verwendet werden. Allerdings werden Logs nur lokal im aktuellen Verzeichnis gespeichert. Dies hat einige Nachteile. Die UI funktioniert nur, wenn du mlflow ui im selben Ordner aufrufst. Außerdem sind die Daten nur auf dem lokalen Rechner verfügbar.
Tracking Server
Alternativ wird deshalb empfohlen einen MLflow Trackingserver aufzusetzen. Diese kann von verschiedenen Cloud-Providern gehostet werden, auf Kubernetes laufen oder auch lokal ausgeführt werden. Folgender Befehl startet den Trackingserver lokal:
mlflow server --host 127.0.0.1 --port 8080Zusätzlich muss man im Code die URI entsprechend setzen:
import mlflow
# Allgemein: mlflow.set_tracking_uri(uri="http://<host>:<port>")
mlflow.set_tracking_uri(uri="http://127.0.0.1:8080")Modelltrainings loggen
Für das Beispiel zum Loggen von Modelltrainings verwenden wir ein einfaches XGBoost-Modell, wie in diesem Beitrag gezeigt. Den vollständigen Code zum Modelltraining findet ihr hier.
Für das Logging erstellen wir ein neues MLflow Experiment:
mlflow.set_experiment("MLflow XGBoost Experiment")Das Logging selbst wir dann in einem MLflow run Kontext ausgeführt.
# Start an MLflow run
with mlflow.start_run():Nun kann man mithilfe der entsprechenden mlflow Funktionen Parameter und Metricen loggen. Tags helfen um den jeweiligen Model-Lauf später wiederfinden zu können. Zusätzlich ist es sinnvoll eine Model Signatur zu erstellen und das Modell selbst zu loggen.
import mlflow
from mlflow.models import infer_signature
import mlflow.xgboost
# Set our tracking server uri for logging
mlflow.set_tracking_uri(uri="http://127.0.0.1:8080")
# Create a new MLflow Experiment
mlflow.set_experiment("MLflow XGBoost Experiment")
# Start an MLflow run
with mlflow.start_run():
# Log the hyperparameters
mlflow.log_params(params)
# Log the loss metric
mlflow.log_metric("accuracy", accuracy)
# Set a tag that we can use to remind ourselves what this run was for
mlflow.set_tag("Training Info", "Basic XGBoost for iris data")
# Infer the model signature
signature = infer_signature(dtrain, bst.predict(dtrain))
# Log the model
model_info = mlflow.xgboost.log_model(
xgb_model=bst,
artifact_path="iris_model",
signature=signature,
input_example=X_train,
registered_model_name="xgboost_iris_model",
)Vertex AI Experiments
Ähnlich wie in MLflow können Experimente in der Google Cloud (GCP) mit Vertex AI Experiments getrackt werden. Es ist besonders geeignet für produktive Cloud-Workflows und lässt sich gut mit Pipelines und anderen GCP-Diensten verknüpfen.
Den vollständigen Code finde ihr hier.
Zur Nutzung in einer Localen Umgebung müssen wir zunächst das aiplatform Paket installieren:
pip install google-cloud-aiplatform --upgradeAußerdem wird empfohlen das BigQuery Storage Paket zu installieren:
pip install google-cloud-bigquery-storageDie Daten für das Experiment laden wir wie in diesem Beitrag beschrieben aus BigQuery.
Um ein Experiment in Vertex AI zu initialisieren verwenden wir den Befehl aiplatform.init().
Anschließend können wir einen Testrun mit dem Befehl aiplatform.start_run() starten.
# Init
aiplatform.init(experiment="rf-experiment-01", project=PROJECT_ID, location="us-central1")
aiplatform.start_run(run="run1")Analog zu MLflow nutzen wir die entsprechenden Funktionen um Parameter, Metriken und Artefakte zu loggen:
# Logs
aiplatform.log_params({"n_estimators": 100, "random_state": 42})
aiplatform.log_metrics({"accuracy": accuracy})Artefakte, wie beispielsweise das Modell selbst, müssen als eigene Objekte erstellt werden:
# Modell speichern
with tempfile.TemporaryDirectory() as tmpdir:
model_path = os.path.join(tmpdir, "model.joblib")
joblib.dump(model, model_path)
# Modell als Artifact erstellen
model_artifact = aiplatform.Artifact.create(
schema_title="system.Model", # GCP kennt dieses Schema
uri=model_path, # Lokaler Pfad oder Cloud Storage URI
display_name="Random Forest Model",
description="RandomForest trained on gochx_table",
metadata={
"framework": "sklearn",
"accuracy": accuracy
},
project=PROJECT_ID,
location=REGION,
)Am Ende schließen wir den Run ab mit:
aiplatform.end_run()Im der GCP Console sind die Experimente anschließend unter Vertex AI / Tests abrufbar.
WICHTIG: Die Metadaten liegen separat ab und sind über Vertex AI / Metadaten abrufbar.