Utiliser AWS Batch et DeepLabCut pour automatiser l’apprentissage de modèles d’estimation de pose 3D dans le Cloud

L’accélération des protocoles expérimentaux à travers l’accès rapide et automatisé aux technologies cloud-natives est un défi auquel sont confrontés les équipes de recherche scientifique. L’estimation de pose 3D (article en anglais) est un exemple de tels protocoles, appliqué au domain de la neuroscience. Cette méthode consiste à enregistrer des flux vidéo pour un animal en laboratoire, pour ensuite les annoter par le scientifique, avant d’être migrés vers le cloud pour entrainer des modèles d’analyse scientifique. L’application de ces modèles par un scientifique permet de mesurer quantitativement, de manière rapide et efficiente, le comportement d’un animal afin de comprendre la réaction du système nerveux face à plusieurs stimuli. DeepLabCut est un logiciel sous licence libre, utilisé par nos clients, qui implémente cette méthode.

Cet article présente une solution pour automatiser la phase d’apprentissage de l’outil DeepLabCut dans le cloud, en utilisant les technologies natives AWS. L’architecture hybride décrite dans cet article permet à un scientifique d’enregistrer et d’annoter les flux vidéo dans le laboratoire de recherche, ensuite d’instancier automatiquement un environnement de calcul préconfiguré sur le cloud AWS. La maitrise des coûts se fait à travers l’instanciation évènementielle de l’environnement de calcul, et la possibilité de le supprimer instantanément une fois le modèle d’apprentissage est établi.

Présentation de la solution

L’architecture hybride comporte trois phases :

• Etape 1 – Définition de la tâche d’apprentissage
• Etape 2 – Préparation des données d’analyse
• Etape 3 – Initialisation de la tâche et export des paramètres du modèle

Le diagramme suivant illustre l’architecture proposée.

Dans la suite de l’article, nous décrivons plus en détails chacune de ces étapes.

Etape 1 – Définition de la tâche d’apprentissage

La définition de la tâche d’apprentissage utilise le service AWS Batch qui permet à des scientifiques d’exécuter des tâches de traitement et de calcul d’apprentissage par lots tout en optimisant les ressources de calcul. Il convient d’abord de créer un environnement de calcul dans AWS Batch en utilisant Amazon EC2 comme support, et en choisissant une instance de calcul accéléré, équipée d’unité de traitement graphique GPU, telle que par exemple la famille G4dn disponible en région AWS (Paris). Pour plus d’information sur comment créer un environnement de calcul utilisant AWS Batch, veuillez vous référer au guide Démarrer avec AWS Batch. Le choix de l’environnement de calcul Amazon EC2 est recommandé par AWS pour exécuter avec AWS Batch des tâches qui requièrent une unité de traitement graphique GPU, ce qui est le cas pour DeepLabCut.

Le fonctionnement nominal de DeepLabCut, disponible en tant qu’image docker, prévoit le chargement manuel des données d’analyse dans le conteneur (e.g., à travers bind mount), et un lancement de l’apprentissage par le scientifique en ligne de commande.

Ici, nous automatisons ces étapes en construisant une nouvelle image à partir de l’image existante. La nouvelle image récupère à son lancement, depuis un bucket dans Amazon Simple Storage Service (S3), les données du projet d’analyse préchargées par le scientifique dans l’étape 2 de cet article. Elle déclenche ensuite l’apprentissage du modèle, en invoquant les fonctions natives DeepLabCut. A la fin de l’exécution, elle exporte les paramètres de sortie du modèle dans un bucket S3.

Créer la nouvelle image conteneur et l’enregistrer dans AWS ECR

Nous utilisons le fichier Dockerfile suivant pour construire la nouvelle image conteneur.

FROM deeplabcut/deeplabcut:latest

ENV LOCAL_PROJECT_PATH=/usr/local/deeplabcut/DLC
ENV config_path=/usr/local/deeplabcut/DLC/config.yaml

RUN pip install boto3
RUN mkdir -p /usr/local/deeplabcut/DLC
WORKDIR /usr/local
ADD dlc-train.py .
CMD ["dlc-train.py"]
ENTRYPOINT ["python"]

Ce fichier extrait l’image de base publique du conteneur deeplabcut et crée les variables d’environnement utilisées pour la phase d’apprentissage. Ensuite, il invoque en entrée le code Python dlc-train.py qui permet de charger les données du projet d’analyse depuis S3, et lancer la phase d’apprentissage. Ce code, disponible dans GitHub, contient trois fonctions :

• import_dlc_zip_project : Importe un projet d’analyse depuis S3 et le déploie dans le conteneur,
• train_dlc_model : lance une tâche d’analyse en invoquant les fonctions natives de l’outil DeepLabCut,
• export_dlc_model_s3 : Exporte le modèle ainsi généré dans S3, où il sera accessible au scientifique.

Utiliser la commande suivante pour construire la nouvelle image Docker :

docker build -t dlc-aws-batch .

Ensuite exporter l’image dans un répertoire privé Amazon Elastic Container Registry (ECR). On suppose vous avez déjà créé votre répertoire « dlc-aws-batch » dans Amazon ECR. Pour plus d’informations sur comment créer un répertoire privé avec Amazon ECR, veuillez vous référer au guide “ Créer un référentiel privé”.

Utiliser la séquence suivante pour importer l’image vers Amazon ECR (12345678910 est à remplacer par l’identifiant du compte AWS utilisé) :

aws ecr get-login --region eu-west-3
docker tag dlc-aws-batch:latest 12345678910.dkr.ecr.eu-west-3.amazonaws.com/dlc-aws-batch:latest
docker push 12345678910.dkr.ecr.eu-west-3.amazonaws.com/dlc-aws-batch:latest

Créer la tâche d’apprentissage dans AWS Batch

Suivre les étapes suivantes pour créer une nouvelle définition de tâche d’apprentissage depuis la console AWS Batch :

• Utiliser l’environnement de calcul Amazon EC2.
• Choisir l’URI de l’image enregistrée dans Amazon ECR dans l’étape précédente.
• Laisser vide le champs Commande
• Dans le champ Rôle d’exécution, renseigner un rôle qui autorise l’accès pour votre tâche à Amazon S3. Pour plus d’information sur les politiques de contrôle d’accès à Amazon S3, veuillez vous référer au guide “Stratégies et autorisations dans Amazon S3”,
• Dans la rubrique Configuration de l’environnement, choisir 1 vCPU, 10Go de mémoire, et 1 GPU, suivant les recommandations de l’outil DeepLabCut,
• Dans la rubrique Variables d’environnement, créer les variables suivantes :
  • TGT_S3_BUCKET: < Bucket dans Amazon S3 qui héberge les données du projet d’analyse importé par le scientifique >
  • DLC_PROJECT_PATH : < Clé du projet deepLabCut dans le bucket Amazon S3 >
  • OUTPUT_PATH : < Préfixe dans le bucket Amazon S3 où sera exporté le modèle d’apprentissage en sortie de la tâche >
• Garder les autres valeurs par défaut,
• Choisir « Créer définition de tâche ».

La capture d’écran suivante montre la configuration de la tâche de calcul dans la console AWS Batch.

Etape 2 – Préparation et chargement des données d’analyse

Cette étape se déroule dans le laboratoire de recherche, et utilise les fonctions native DeepLabCut. Un projet d’analyse initié par l’outil deepLabCut comporte l’arborescence de fichiers suivante :

~ % ipython

In [1] : import deeplabcut
In [2] : deeplabcut.create_new_project(‘demo’, ‘test-dlc’, [‘/path-to-videos/demo.avi’])
Created “/home/dlc/projects/videos”
Created “/home/dlc/projects/labeled-data”
Created “/home/dlc/projects/training-datasets”
Created “/home/dlc/projects/dlc-models”
Generated “/home/dlc/projects/config.yaml”

In [3] : path_config = ‘/home/dlc/projects/config.yaml’ 
In [4] : deeplabcut.extract_frames(path_config, ‘automatic’, ‘kmeans’)
In [5] : deeplabcut.label_frames(path_config)
In [6] : deeplabcut.create_training_dataset(path_config)
In [7] : exit()

• Le répertoire « videos » contient des liens symboliques vers les enregistrements vidéo renseignés en paramètre d’entrée à la fonction create_new_project,
• Le répertoire « labeled-data » contient des extraits d’images annotées manuellement par le scientifique à travers une interface graphique de l’outil,
• Le répertoire « training-datasets » contient les données d’entrainement utilisées en entrée du modèle d’apprentissage,
• Le fichier « config.yaml » introduit les paramètres du modèle d’entrainement, tel que les labels, et la répartition des données d’entrainement et de test,
• Le répertoire « dlc-models » est vide à l’instant. Il est destiné à contenir les modèles d’apprentissage construits sur le cloud AWS, selon la tâche de calcul décrite dans l’étape 1.

La fonction deeplabcut.extract_frames en ligne 4 extrait, depuis les enregistrements vidéo, des images à annoter par le scientifique. La fonction deeplabcut.label_frames en ligne 5 expose une interface graphique où sont affichées les images, et où le scientifique positionne les points de pose relatif à l’animal.

Une fois le projet d’analyse créé et les données d’entrainement labellisées en laboratoire, générer un dossier zip du répertoire local et l’exporter dans le bucket Amazon S3 TGT_S3_BUCKET, en utilisant la clé correspondante à la variable d’environnement DLC_PROJECT_PATH définie dans l’étape 1.

zip -r dossier.zip /home/dlc/projects

Etape 3 – Automatisation de la tâche d’apprentissage et export du modèle

Cette étape consiste à automatiser le lancement de la tâche d’apprentissage décrite dans l’étape 1, suite à la soumission du projet d’analyse tel que décrit dans l’étape 2. Nous utilisons ici le service Amazon EventBridge.

Dans la console Amazon EventBridge, suivre les étapes suivantes :

• Choisir Créer une règle, et désigner un nom pour cette règle (ex., « dlc-batch-trigger-rule »),
• Choisir Règle avec un modèle d’évènement en tant que Type de règle,
• Choisir Événements AWS ou événements partenaires EventBridge en tant que Source d’évènement,
• Choisir la méthode « Utiliser le formulaire d’évènement », et éditer le document JSON comme suit :

{
  "source": ["aws.s3"],
  "detail-type": ["Object Created"],
  "detail": {
    "bucket": {  "name": ["testbucket987654321"]},
    "object": { "key": [{ "prefix": "testDLC/dossier.zip" }] }}
}

• Dans l’onglet « Types de cibles », choisir « Service AWS », et renseigner les champs selon l’exemple illustré dans la figure ci-dessous

• Choisir Créer une règle

La création d’un objet « testDLC/dossier.zip » dans le bucket Amazon S3 TGT_S3_BUCKET déclenche automatiquement la règle EventBridge ainsi créée, qui par la suite invoque la tâche de calcul AWS Batch décrite dans l’étape 1.

Conclusion

En utilisant des services tel que AWS Batch, Amazon ECR, Amazon S3, et Amazon EventBridge, vous pouvez créer une solution à base de technologie et services entièrement gérés par AWS, avec une architecture évènementielle permettant aux scientifiques et aux équipes de recherche d’étendre leurs protocoles expérimentaux dans le Cloud.

Dans cet article, nous avons illustré un exemple d’implémentation, permettant d’automatiser l’apprentissage dans le Cloud AWS de modèles d’estimation de pose 3D en utilisant l’outil DeepLabCut. L’architecture hybride décrite dans cet article offre un environnement intégré pour les équipes de recherche, permettant de concilier d’un côté les contraintes locales liées à la collecte et le pré-traitement des données d’analyse en laboratoire, et de l’autre côté l’usage automatisé et à la demande des ressources de calcul dans le Cloud, dans le but d’accélérer la réalisation des protocoles expérimentaux.

Article écrit par Nizar Kheir, Solutions Architect dans les équipes AWS France.