调整您的 TensorFlow 训练脚本 - Amazon SageMaker
在你的 TensorFlow Keras 训练脚本中注册挂钩
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

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

调整您的 TensorFlow 训练脚本

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

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

import smdebug.tensorflow as smd

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

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

（可选）为外部训练作业创建挂钩 SageMaker

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

    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 调试器将无法收集任何张量。

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

调整完训练脚本后，继续到 第 2 步：使用 SageMaker Python SDK 启动和调试训练作业。