调整您的 TensorFlow 训练脚本 - 亚马逊 SageMaker AI
在你的 TensorFlow Keras 训练脚本中注册挂钩

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

调整您的 TensorFlow 训练脚本

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

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

import smdebug.tensorflow as smd

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

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

（可选）创建用于在 SageMaker AI 之外训练作业的挂钩

如果您在本地模式下直接在 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() 方法获取您的模型并遍历每一层，寻找与您通过 使用 SageMaker Python 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 AI 调试器将无法收集张量。但是，模型权重、偏差和损失等其他张量，无需在急切模式下运行即可收集。

    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 文档中的钩子方法。

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