调整您的 TensorFlow 训练脚本 - Amazon SageMaker

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

调整您的 TensorFlow 训练脚本

要开始收集模型输出张量并调试训练问题,请对 TensorFlow 训练脚本进行以下修改。

在其中创建用于训练作业的挂钩 SageMaker

import smdebug.tensorflow as smd hook=smd.get_hook(hook_type="keras", create_if_not_exists=True)

这会在你开始 SageMaker 训练作业时产生一个挂钩。当您在估算器中使用 Python 使用调试器启动训练作业 SageMaker SDK使用DebuggerHookConfigTensorBoardConfig、或Rules中的任何一个启动训练作业时, SageMaker 会向您的训练实例添加一个由该smd.get_hook方法获取的JSON配置文件。请注意,如果您在估算器APIs中不包含任何配置,则不会有配置文件可供钩子查找,并且函数会返回。None

(可选)为外部训练作业创建挂钩 SageMaker

如果您在本地模式下直接在 SageMaker Notebook 实例、Amazon EC2 实例或您自己的本地设备上运行训练作业,请使用smd.Hook类来创建挂钩。但是,这种方法只能存储张量集合并可用于可 TensorBoard视化。 SageMaker 调试器的内置规则不适用于本地模式。在这种情况下,smd.get_hook 方法也会返回 None

如果您要创建手动钩子,请使用以下带有逻辑的代码片段来检查钩子是否返回 None,并使用 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 )

添加挂钩创建代码后,继续阅读以下 TensorFlow Keras 主题。

注意

SageMaker 调试器目前仅支持 TensorFlow Keras。

在你的 TensorFlow Keras 训练脚本中注册挂钩

以下过程介绍如何使用钩子及其方法,从模型和优化器中收集输出标量和张量。

  1. 用钩子的类方法包装您的 Keras 模型和优化器。

    hook.register_model() 方法获取您的模型并遍历每一层,寻找与您通过 使用 Python 使用调试器启动训练作业 SageMaker SDK 中配置提供的正则表达式匹配的任何张量。通过这种钩子方法可以收集到的张量包括权重、偏差和激活。

    model=tf.keras.Model(...) hook.register_model(model)
  2. hook.wrap_optimizer() 方法包装优化器。

    optimizer=tf.keras.optimizers.Adam(...) optimizer=hook.wrap_optimizer(optimizer)
  3. 在急切模式下编译模型 TensorFlow。

    要从模型中收集张量,例如每层的输入和输出张量,必须在急切模式下运行训练。否则, SageMaker调试器将无法收集张量。但是,模型权重、偏差和损失等其他张量,无需在急切模式下运行即可收集。

    model.compile( loss="categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"], # Required for collecting tensors of each layer run_eagerly=True )
  4. 将钩子注册到 tf.keras.Model.fit() 方法。

    要从您注册的钩子中收集张量,请将 callbacks=[hook] 添加到 Keras model.fit() 类方法中。这会将 sagemaker-debugger 钩子作为 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 仅提供不提供对其值的访问的符号渐变变量。要收集梯度,请使用 hook.wrap_tape() 方法包装 tf.GradientTape,这要求您如下所示编写自己的训练步骤。

    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))

    通过对磁带进行包装,sagemaker-debugger 钩子可以识别输出张量,例如梯度、参数和损失。封装磁带可确保围绕磁带对象函数hook.wrap_tape()的方法(例如、、)将设置 D SageMaker ebugger 的编写器并保存作为输入(可训练变量和损失)和gradient()(梯度)的输出gradient()(渐变)提供的张量。push_tape() pop_tape() gradient()

    注意

    要使用自定义训练循环进行收集,请确保使用急切模式。否则, SageMaker 调试器无法收集任何张量。

有关钩子APIs提供的用于构造sagemaker-debugger钩子和保存张量的操作的完整列表,请参阅 sagemaker-debuggerPython SDK 文档中的钩子方法

调整完训练脚本后,继续到 使用 Python 使用调试器启动训练作业 SageMaker SDK