Prise en main de MLflow, un outil pour tracer les résultats de vos expériences

Dans la tendance de l’industrialisation de la data science, une douleur importante est de garder la trace des expérimentations menées. C’est-à-dire tracer les résultats obtenus avec une version du jeu de données, un type de modèle et un set de méta-paramètres.

Le graal est peut-être de faire comme Facebook : des modèles de modèles. Mais avant d’atteindre cela, encore faut-il collecter les données et s’en servir pour suivre les résultats.

Pour répondre à cet enjeu de collecte de données en lien avec les expérimentations, il existe un projet qui commence à faire parler de lui. Il s’agit du projet MLflow, initié par Databricks et hébergé par la fondation Apache.

Il vise à répondre à trois ambitions :

  • Enregistrer et requêter les expériences : code, données, configuration et résultats (MLflow Tracking),
  • Packager pour rendre reproductible les expériences (MLflow Projects),
  • Proposer un format générique pour envoyer les modèles vers divers outils de déploiements (MLflow Models).

J’expose ci-dessous le résultat de la prise en main de la première ambition.

Prise en main de l’outil

L’installation de l’outil est agréablement simple, il suffit de faire:


pip install mlflow

Pour faciliter la prise en main de l’outil, MLflow propose des exemples très bien faits au format .py ou notebook. En voici un petit résumé.

Pour envoyer des informations vers MLflow, il faut baliser son code avec des “logger”. Il y a deux types : mlflow.log_param pour enregistrer des paramètres et mlflow.log_metric pour enregistrer des résultats.

Enfin pour visualiser les résultats il y a une interface démarrable avec la commande mlflow ui.

Figure 1 : aperçu de l’interface Mlflow

La figure 1 présente 3 entraînements de modèles avec 3 jeux de paramètres différents. Cette vue indique une date, un nom d’utilisateur, l’outil utilisé, les méta-paramètres utilisés et enfin les métriques.

La seule contrainte sur les méta-paramètres et les métriques est qu’il doivent être un objet (string, int, float, ...) de longueur 1. Il est donc possible d’enregistrer beaucoup de choses y compris le nom du jeu de données, par exemple.

Le petit “plus” vraiment appréciable est que l’on peut logger des métriques au fur et à mesure de l'entraînement et ainsi visualiser une évolution.

Figure 2 : visualisation de l’évolution d’une métrique au cours de l'entraînement du modèle.

En synthèse, l’outil est une sorte de logger spécifique avec une interface graphique: c’est simple et efficace.

Pour aller plus loin

Pour aller plus loin et jouer un peu plus avec l’outil, j’ai décidé de créer mon propre modèle. J’ai créé une fake box : c’est un algorithme qui simule le comportement d’un modèle en loggant une métrique qui augmente aléatoirement. L’idée est de voir comment intégrer MLflow à un modèle.

import time
import numpy as np

class fake_box():
    def __init__(self, n_estimator, n_neurons):
        self.n_estimator = n_estimator
        self.n_neurons = n_neurons

    def fit(self, X, y, n_epoch = 1000, n_jobs = -1):
        print("Fitting model with {} estimators, {} neurons on {} jobs".format(self.n_estimator, self.n_neurons, n_jobs))
        perf = 0.5
        for i in range(n_epoch):
            time.sleep(0.1)
            perf += np.random.uniform(high=1./(n_epoch * perf))
            perf = round(min(perf, 1), 4)
            perf_test = round(perf - np.random.uniform(high=perf * 2./100), 4)
            print("Epoch {}: AUC in train: {} AUC in test: {} in 0.1s".format(i, perf, perf_test))

Algorithme 1 : Définition de la fake box

Ajouter le tracking de MLflow est plutôt simple, il s’agit de baliser le code avec les log_param et log_metric :

import mlflow

Algorithme 2 : Ajout de Mlflow à la fake box

Cela permet d'obtenir dans l’interface une courbe de la performance au fil des epochs :

Figure 3 : AUC au fil des epochs

Une petite déception est qu’il ne semble pas possible d’empiler 2 courbes type AUC en train et en test.

Si on relance, les résultats s’accumulent sur le même graphique, même si on change les paramètres.

Figure 4 : Accumulation des résultats au fil des expérimentations.

Cette accumulation est très pratique pour faire des exécutions successives. Mais lors d’un changement de paramètres, il est nécessaire de demander à MLflow de stocker cela dans une autre vue.

Il est temps d’introduire une nouvelle notion celle de run. Il s’agit d’un nouvel entraînement. Il faut déclarer lorsque l’on change d’entraînement pour séparer les résultats et obtenir plusieurs vues comme présentées sur la figure 1.

Pour cela , il faut arrêter le run en cours, puis en lancer un nouveau avec les commandes suivantes :

mlflow.end_run()
mlflow.active_run()

Cette méthode permet de choisir lorsque l’on passe à un nouveau run. Une autre approche est la suivante :

with mlflow.start_run():
    <some code>

Avec cette syntaxe, chaque nouvelle exécution est considérée comme un nouveau run, ce qui empêche de réaliser l’entraînement en plusieurs fois.

En synthèse l’intégration de MLflow à un algorithme est relativement aisée,  grâce à quelques commandes simples il est possible de tracer les résultats de ses expérimentations.

Le collaboratif avec MLflow

Pour travailler à plusieurs sur plusieurs projets, il y a deux notions importantes : le run et l’experiment. La notion de run a déjà été décrite dans la précédente partie : le run correspond à une itération, un set de paramètres.

La documentation décrit l’experiment comme un éventuel moyen d’organiser ses run. Selon les besoins, cette notion pourra être utilisée pour différents projets indépendants, pour le même projet mais avec des approches différentes (ex : réseaux de neurones versus ensembles d’arbres), pour différentes versions du même projet (ex : celle en production versus celle en développement)...

Pour configurer une nouvelle expérience, il faut indiquer un nouveau nom d’expérience:

mlflow.set_experiment("Fake box {}".format(9))

Cela est représenté dans l’ui de la manière suivante :

Figure 5 : Représentation des expériences et des runs

Sur la figure 5, on voit que l’on travaille dans une expérience (entourée en bleu), dans cette expérience on réalise des runs (entourés en orange). On peut également observer que chaque run est associé à un utilisateur ce qui permet de suivre le travail de chacun.

Pour permettre de travailler sur plusieurs projets à la fois, il est possible de reprendre un run et une expérience en spécifiant leur nom :

mlflow.set_experiment(<Insert experiment name>)
mlflow.start_run(<Insert run ID>)

Enfin, pour permettre de travailler ensemble il est possible de démarrer le système sur une machine distante et de partager l’ip / port. La configuration à la machine distante se fait de la manière suivante :

mlflow.set_tracking_uri("http://<Insert ip:port>"):

En synthèse, MLflow permet la collaboration en lançant le service sur une machine distante et en gérant les projets et les essais en utilisant les notions de run et de experiment.

Fonctionnement technique de l’outil

Pour bien comprendre le fonctionnement de cet outil, nous allons regarder l’architecture, la gestion des erreurs et la compatibilité.

Tout d’abord l’architecture comme présenté en figure 6, consiste en 3 briques:

  • L’enregistrement des résultats,
  • Le stockage,
  • La visualisation.

Figure 6 : Architecture macroscopique

L’enregistrement des résultats peut se fait de la façon présentée précédemment en python. Il existe aussi des API pour R, Java, REST.

Le stockage des résultats se fait via des fichiers plats :

mlruns
|    experiment id_1
|    |    run_id_1
|    |    |    artifacts
|    |    |    metrics
|    |    |    |    metric_name_1
|    |    |    params
|    |    |    |    param_name_1
|    |    |    |    param_name_2
|    |    |    meta.yaml
|    |    run_id_2|     | meta.yaml
|    experiment id_2

Figure 7 : Architecture du stockage

En détails:

  • Le dossier mlruns est créé là où MLflow est lancé,
  • Les dossiers experiments contiennent les informations de chaque expérience, ils s’appellent 0, 1, 2, …
  • Les dossier runs contiennent les informations de chaque run, il s'appellent avec un id (par exemple : 0ca26355c1e04326a77118f3d82804a3),
  • Le meta.yaml au niveau de l’experiment contient quelques informations dont le nom de l’expérience,
  • Le meta.yaml au niveau du run contient toutes les informations visible au niveau du run, telles que le nom d’utilisateur, l’outil, le timestamp de début, le timestamp de fin,
  • Les fichiers metric et param ont le nom qui a été envoyé dans le logger. Ce sont des fichiers plats avec une valeur par ligne, correspondant à chaque appel au logger en utilisant leurs noms. Ils sont par la suite parsés et interprétés par l’interface pour faire les graphiques présentés précédemment.

Enfin, l’interface est développée en JavaScript.

Le tout est embarqué dans un package Python, ce qui le rend très facilement installable.

La gestion des erreurs

Le projet est toujours en version bêta, et contient quelques bugs. La chose la moins aboutie selon moi est la gestion des erreurs.

J’ai eu la mauvaise idée d’envoyer un entier en guise de nom pour une expérience. MLflow a alors généré une erreur, mais pas celle que l’on aurait espéré (un message signalant que les entiers ne sont pas acceptés), il ne dit rien. Cependant l’interface ne marche plus et propose un joli écran d'erreur:

Figure 8 : Écran d’erreur

Quand MLflow plante, il fait ça bien ! Il est impossible d’accéder à l’interface et redémarrer l’outil ne change rien. La seule solution que j’ai trouvée est de supprimer le dossier mlruns qui stocke les données. Mais cela fait perdre toutes mes expériences.

Une meilleure solution aurait été d’identifier le nom de l’expérience qui a généré les erreurs et supprimer uniquement celle-ci. Mais encore faut-il le connaître...

En synthèse, la gestion des erreurs n’est vraiment pas aboutie pour l’instant. Une façon de contourner le problème pourrait être la création de backups.

Conclusion

MLflow Tracking répond à un vrai besoin des data scientists de manière très simple et facile à prendre en main. L’outil utilise 4 notions simples : l’experiment, le run, le param et la metric.

La solution n’est pas encore complètement sèche, la gestion des erreurs n’est pas aboutie et les visualisations proposées dans l’interface restent assez basiques.

Le projet est en bêta mais hébergé par la fondation Apache. Il est déjà très utilisé et apprécié par la communauté (près de 3000 étoiles sur Github). Il devrait donc à l’avenir continuer à se développer.

En tout cas, je vais tâcher de l’utiliser dans mes projets de machine learning et je ferai peut-être un article de retour d’expérience dans quelques temps.