Adapter votre script d'entraînement TensorFlow

Pour commencer à collecter les tenseurs de sortie du modèle et déboguer les problèmes d'entraînement, apportez les modifications suivantes à votre script d'entraînement TensorFlow.

Créez un hook pour les tâches d'entraînement dans SageMaker

import smdebug.tensorflow as smd

hook=smd.get_hook(hook_type="keras", create_if_not_exists=True) 

Cela permet de créer un hook lorsque vous commencez une tâche d'entraînement SageMaker. Lorsque vous lancez une tâche d'entraînement dans Étape 2 : lancer et déboguer des tâches d'entraînement à l'aide du kit SDK Python SageMaker avec l'un des éléments DebuggerHookConfig, TensorBoardConfig ou Rules dans votre estimateur, SageMaker ajoute un fichier de configuration JSON à votre instance d'entraînement qui est récupéré par la méthode smd.get_hook. Notez que si vous n'incluez aucune des API de configuration dans votre estimateur, il n'y aura aucun fichier de configuration susceptible d'être trouvé par le hook, et la fonction renvoie None.

(Facultatif) Créez un hook pour les tâches d'entraînement en dehors de SageMaker

Si vous exécutez des tâches d'entraînement en mode local, directement sur des instances de bloc-notes SageMaker, des instances Amazon EC2 ou vos propres appareils locaux, utilisez la classe smd.Hook pour créer un hook. Toutefois, cette approche peut stocker uniquement les collections de tenseurs utilisables pour la visualisation de TensorBoard. Les règles intégrées de SageMaker Debugger ne fonctionnent pas avec le mode local. La méthode smd.get_hook renvoie également None dans ce cas.

Si vous souhaitez créer un hook manuel, utilisez l'extrait de code suivant avec la logique permettant de vérifier si le hook renvoie None, et créez un hook manuel à l'aide de la classe smd.Hook.

import smdebug.tensorflow as smd

hook=smd.get_hook(hook_type="keras", create_if_not_exists=True) 

if hook is None:
    hook=smd.KerasHook(
        out_dir='/path/to/your/local/output/',
        export_tensorboard=True
    )

Après avoir ajouté le code de création du hook, passez à la rubrique suivante pour TensorFlow Keras.

Note

SageMaker Debugger prend actuellement uniquement en charge TensorFlow Keras.

Enregistrer le hook dans votre script d'entraînement TensorFlow Keras

La procédure suivante explique comment utiliser le hook et ses méthodes pour collecter des scalaires et des tenseurs de sortie à partir de votre modèle et de votre optimiseur.

1. Enveloppez votre modèle Keras et votre optimiseur avec les méthodes de classe du hook.

   La méthode hook.register_model() prend votre modèle et itère sur chaque couche, recherchant tous les tenseurs qui correspondent aux expressions régulières que vous fournirez via la configuration dans Étape 2 : lancer et déboguer des tâches d'entraînement à l'aide du kit SDK Python SageMaker. Les tenseurs collectables via cette méthode de hook sont des poids, des biais et des activations.

   model=tf.keras.Model(...)
   hook.register_model(model)

2. Enveloppez l'optimiseur avec la méthode hook.wrap_optimizer().

   optimizer=tf.keras.optimizers.Adam(...)
   optimizer=hook.wrap_optimizer(optimizer)

3. Compilez le modèle en mode Eager dans TensorFlow.

   Pour collecter des tenseurs à partir du modèle, tels que les tenseurs d'entrée et de sortie de chaque couche, vous devez exécuter l'entraînement en mode Eager. Sinon, SageMaker Debugger ne sera pas en mesure de collecter les tenseurs. Cependant, d'autres tenseurs, tels que les poids, les biais et les pertes du modèle, peuvent être collectés sans exécuter explicitement le mode Eager.

   model.compile(
       loss="categorical_crossentropy", 
       optimizer=optimizer, 
       metrics=["accuracy"],
       # Required for collecting tensors of each layer
       run_eagerly=True
   )

4. Enregistrez le hook avec la méthode tf.keras.Model.fit().

   Pour collecter les tenseurs des hooks que vous avez enregistrés, ajoutez callbacks=[hook] à la méthode de classe model.fit() Keras. Le hook sagemaker-debugger sera alors transmis en tant que rappel Keras.

   model.fit(
       X_train, Y_train,
       batch_size=batch_size,
       epochs=epoch,
       validation_data=(X_valid, Y_valid),
       shuffle=True, 
       callbacks=[hook]
   )

5. TensorFlow 2.x fournit uniquement des variables de gradient symboliques qui ne donnent pas accès à leurs valeurs. Pour collecter des gradients, enveloppez tf.GradientTape avec la méthode hook.wrap_tape(), ce qui vous oblige à écrire votre propre étape d'entraînement comme suit.

   def training_step(model, dataset):
       with hook.wrap_tape(tf.GradientTape()) as tape:
           pred=model(data)
           loss_value=loss_fn(labels, pred)
       grads=tape.gradient(loss_value, model.trainable_variables)
       optimizer.apply_gradients(zip(grads, model.trainable_variables))

   En enveloppant la bande, le hook sagemaker-debugger peut identifier les tenseurs de sortie tels que les gradients, les paramètres et les pertes. Le fait d'envelopper la bande garantit que la méthode hook.wrap_tape() basée sur les fonctions de l'objet de bande, tels que push_tape(), pop_tape() et gradient(), configurera les agents d'écriture de SageMaker Debugger et enregistrera les tenseurs qui sont fournis en entrée de gradient() (variables entraînables et pertes) et en sortie de gradient() (gradients).

   Note

   Pour collecter avec une boucle d'entraînement personnalisée, assurez-vous d'utiliser le mode Eager. Sinon, SageMaker Debugger ne pourra pas collecter des tenseurs.

Pour obtenir la liste complète des actions proposées par les API du hook sagemaker-debugger pour créer des hooks et enregistrer des tenseurs, consultez Hook Methods dans la documentation du SDK sagemaker-debugger Python.

Une fois que vous avez terminé d'adapter votre scénario d'entraînement, passez à Étape 2 : lancer et déboguer des tâches d'entraînement à l'aide du kit SDK Python SageMaker.