Exécution de tâches d’entraînement TensorFlow avec SageMaker Training Compiler

Vous pouvez utiliser n’importe quelle interface SageMaker AI pour exécuter une tâche d’entraînement avec SageMaker Training Compiler : Amazon SageMaker Studio Classic, instances de bloc-notes Amazon SageMaker, AWS SDK pour Python (Boto3) ou AWS Command Line Interface.

Utilisation du kit SDK SageMaker Python

Pour activer SageMaker Training Compiler, ajoutez le paramètre compiler_config à l’estimateur SageMaker AI TensorFlow ou Hugging Face. Importez la classe TrainingCompilerConfig et transmettez-en une instance au paramètre compiler_config. Les exemples de code suivants montrent la structure des classes d’estimateur SageMaker AI avec SageMaker Training Compiler activé.

Pour obtenir des informations adaptées à votre cas d’utilisation, consultez l’une des options suivantes.

La liste suivante est l’ensemble minimal de paramètres requis pour exécuter une tâche d’entraînement SageMaker avec le compilateur.

disable_profiler=True,
debugger_hook_config=False

Si la tâche d’entraînement avec le compilateur est lancée avec succès, vous recevez les journaux suivants lors de la phase d’initialisation de la tâche :

Utilisation du kit SageMaker AI Python SDK et extension des conteneurs de deep learning de framework SageMaker AI

Les conteneurs AWS de deep learning (DLC, Deep Learning Containers) pour TensorFlow utilisent des versions adaptées de TensorFlow qui incluent des modifications en plus du framework TensorFlow open source. Les conteneurs de deep learning de framework SageMaker AI sont optimisés pour l’infrastructure AWS sous-jacente et Amazon SageMaker AI. Grâce aux conteneurs de deep learning, l’intégration de SageMaker Training Compiler améliore encore les performances par rapport au TensorFlow natif. En outre, vous pouvez créer un conteneur d’entraînement personnalisé en étendant l’image DLC.

Pour étendre et personnaliser les conteneurs DLC de SageMaker AI TensorFlow pour votre cas d’utilisation, utilisez les instructions suivantes.

Création d’un fichier Dockerfile

Utilisez le modèle Dockerfile suivant pour étendre le conteneur DLC de SageMaker AI TensorFlow. Vous devez utiliser l’image de conteneur DLC de SageMaker AI TensorFlow comme image de base de votre conteneur Docker. Pour trouver les URI d’image de conteneur DLC de SageMaker AI TensorFlow, consultez Frameworks pris en charge.

# SageMaker AI TensorFlow Deep Learning Container image
FROM 763104351884.dkr.ecr.<aws-region>.amazonaws.com/tensorflow-training:<image-tag>

ENV PATH="/opt/ml/code:${PATH}"

# This environment variable is used by the SageMaker AI container 
# to determine user code directory.
ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code

# Add more code lines to customize for your use-case
...

Pour plus d’informations, consultez Étape 2 : créer et télécharger le fichier Dockerfile et les scripts d’entraînement Python.

Tenez compte des pièges suivants lors de l’extension des conteneurs DLC de framework SageMaker AI :

Par exemple, un problème connu est que les bibliothèques tensorflow/models et tensorflow/text tentent toujours de réinstaller TensorFlow open source. Si vous devez installer ces bibliothèques pour choisir une version spécifique pour votre cas d’utilisation, nous vous recommandons de consulter les fichiers Dockerfile de conteneur DLC de SageMaker AI TensorFlow pour v2.9 ou version ultérieure. Les chemins d’accès aux fichiers Dockerfile sont généralement au format suivant : tensorflow/training/docker/<tensorflow-version>/py3/<cuda-version>/Dockerfile.gpu. Dans les fichiers Dockerfile, vous devez trouver les lignes de code permettant de réinstaller le fichier binaire TensorFlow géré par AWS (spécifié dans la variable d’environnement TF_URL) et d’autres dépendances dans l’ordre. La section de réinstallation doit ressembler à l’exemple suivant :

# tf-models does not respect existing installations of TensorFlow 
# and always installs open source TensorFlow

RUN pip3 install --no-cache-dir -U \
    tf-models-official==x.y.z

RUN pip3 uninstall -y tensorflow tensorflow-gpu \
  ; pip3 install --no-cache-dir -U \
    ${TF_URL} \
    tensorflow-io==x.y.z \
    tensorflow-datasets==x.y.z

Génération et envoi (push) vers ECR

Pour générer et envoyer (push) votre conteneur Docker vers Amazon ECR, suivez les instructions des liens suivants :

Exécution en utilisant l’estimateur du kit SageMaker Python SDK

Utilisez l’estimateur de framework SageMaker AI TensorFlow comme d’habitude. Vous devez spécifier image_uri pour utiliser le nouveau conteneur que vous avez hébergé dans Amazon ECR.

import sagemaker, boto3
from sagemaker import get_execution_role
from sagemaker.tensorflow import TensorFlow, TrainingCompilerConfig

account_id = boto3.client('sts').get_caller_identity().get('Account')
ecr_repository = 'tf-custom-container-test'
tag = ':latest'

region = boto3.session.Session().region_name

uri_suffix = 'amazonaws.com'

byoc_image_uri = '{}.dkr.ecr.{}.{}/{}'.format(
    account_id, region, uri_suffix, ecr_repository + tag
)

byoc_image_uri
# This should return something like
# 111122223333.dkr.ecr.us-east-2.amazonaws.com/tf-custom-container-test:latest

estimator = TensorFlow(
    image_uri=image_uri,
    role=get_execution_role(),
    base_job_name='tf-custom-container-test-job',
    instance_count=1,
    instance_type='ml.p3.8xlarge'
    compiler_config=TrainingCompilerConfig(),
    disable_profiler=True,
    debugger_hook_config=False
)

# Start training
estimator.fit()

Activation de SageMaker Training Compiler à l’aide de l’opération d’API CreateTrainingJob de SageMaker AI

Les options de configuration de SageMaker Training Compiler doivent être spécifiées via les champs AlgorithmSpecification et HyperParameters dans la syntaxe de demande pour l’opération d’API CreateTrainingJob.

"AlgorithmSpecification": {
    "TrainingImage": "<sagemaker-training-compiler-enabled-dlc-image>"
},

"HyperParameters": {
    "sagemaker_training_compiler_enabled": "true",
    "sagemaker_training_compiler_debug_mode": "false"
}

Pour trouver une liste complète des URI d’image de conteneur Deep Learning Containers qui ont SageMaker Training Compiler implémenté, consultez Cadres pris en charge.