TensorFlow - 亚马逊 SageMaker AI
TensorFlow 模型带 Hugging Face Transformers 的 TensorFlow 模型
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

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

TensorFlow

将您自己的 TensorFlow 模型引入 SageMaker AI，然后使用 SageMaker Training Compiler 运行训练作业。

TensorFlow 模型

SageMaker Training Compiler 会自动优化基于原生 TensorFlow API 或高级 Keras API 构建的模型训练工作负载。

提示

要预处理输入数据集，请确保使用静态输入形状。动态输入形状可以启动对模型的重新编译，并且可能会增加总训练时间。

使用 Keras（推荐）

为了获得最佳的编译器加速，我们建议使用属于 TensorFlow Keras 子类的模型 (tf.keras.Model)。

对于单个 GPU 训练

您无需在训练脚本中进行额外的更改。

没有 Keras 的情况下

SageMaker Training Compiler 不支持在 TensorFlow 中紧急执行。相应地，您应使用 TensorFlow 函数装饰器 (@tf.function) 包装模型和训练循环，以利用编译器加速。

SageMaker Training Compiler 执行图表级优化，并使用装饰器来确保将 TensorFlow 函数设置为在图表模式下运行。

对于单个 GPU 训练

默认情况下，TensorFlow 2.0 或更高版本具有紧急执行功能，您应在用于构建 TensorFlow 模型的每个函数的前面添加 @tf.function 装饰器。

带 Hugging Face Transformers 的 TensorFlow 模型

带 Hugging Face Transformers 的 TensorFlow 模型基于 TensorFlow 的 tf.keras.Model API。Hugging Face Transformers 还为 TensorFlow 提供预训练的模型类，以便减少配置自然语言处理 (NLP) 模型所需的工作量。在使用 Transformers 库创建自己的训练脚本后，您可以通过将 SageMaker AI HuggingFace 估算器与 SageMaker Training Compiler 配置类结合使用来运行训练脚本，如使用 SageMaker Training Compiler 运行 TensorFlow 训练作业中的上一主题所示。

SageMaker Training Compiler 会自动优化基于原生 TensorFlow API 或高级 Keras API（例如 TensorFlow 转换器模型）构建的模型训练工作负载。

提示

在训练脚本中使用 Transformers 为 NLP 模型创建分词器时，请务必通过指定 padding='max_length' 来使用静态输入张量形状。请勿使用 padding='longest'，因为填充批处理中的最长序列可能会改变每个训练批处理的张量形状。动态输入形状可以启动对模型的重新编译，并可能会增加总训练时间。有关 Transformers 分词器的填充选项的更多信息，请参阅 Hugging Face Transformers 文档中的填充和截断。

使用 Keras

为了获得最佳的编译器加速，我们建议使用属于 TensorFlow Keras 子类的模型 (tf.keras.Model)。正如 Hugging Face Transformers 文档中的快速导览页面上所述，您可以将这些模型用作常规 TensorFlow Keras 模型。

对于单个 GPU 训练

您无需在训练脚本中进行额外的更改。

对于分布式训练

在 tf.distribute.Strategy.scope() 调用范围内使用 Keras API 构建和训练模型时，SageMaker Training Compiler 加速可以透明方式用于多 GPU 工作负载。

  1. 选择合适的分布式训练策略。

    1. 对于单节点多 GPU，请使用 tf.distribute.MirroredStrategy 设置策略。

      strategy = tf.distribute.MirroredStrategy()

    2. 对于多节点多 GPU，请在创建策略之前，先添加以下代码来正确设置 TensorFlow 分布式训练配置。

      def set_sm_dist_config():
    DEFAULT_PORT = '8890'
    DEFAULT_CONFIG_FILE = '/opt/ml/input/config/resourceconfig.json'
    with open(DEFAULT_CONFIG_FILE) as f:
        config = json.loads(f.read())
        current_host = config['current_host']
    tf_config = {
        'cluster': {
            'worker': []
        },
        'task': {'type': 'worker', 'index': -1}
    }
    for i, host in enumerate(config['hosts']):
        tf_config['cluster']['worker'].append("%s:%s" % (host, DEFAULT_PORT))
        if current_host == host:
            tf_config['task']['index'] = i
    os.environ['TF_CONFIG'] = json.dumps(tf_config)

set_sm_dist_config()

      使用 tf.distribute.MultiWorkerMirroredStrategy 设置策略。

      strategy = tf.distribute.MultiWorkerMirroredStrategy()

  2. 使用所选策略来包装模型。

    with strategy.scope():
    # create a model and do fit

没有 Keras 的情况下

如果您要使用不带 Keras 的 TensorFlow 来引入带自定义训练循环的自定义模型，则应使用 TensorFlow 函数装饰器 (@tf.function) 来包装模型和训练循环，以便利用编译器加速。

SageMaker Training Compiler 执行图表级优化，并使用装饰器来确保将 TensorFlow 函数设置为在图表模式下运行。

对于单个 GPU 训练

默认情况下，TensorFlow 2.0 或更高版本具有紧急执行功能，您应在用于构建 TensorFlow 模型的每个函数的前面添加 @tf.function 装饰器。

对于分布式训练

除了进行为使用 Keras 进行分布式训练所需的更改外，您还需确保使用 @tf.function 为每个 GPU 上运行的函数添加注释，同时不为跨 GPU 通信函数添加注释。示例训练代码应类似于以下内容：

@tf.function()
def compiled_step(inputs, outputs):
    with tf.GradientTape() as tape:
        pred=model(inputs, training=True)
        total_loss=loss_object(outputs, pred)/args.batch_size
    gradients=tape.gradient(total_loss, model.trainable_variables)
    return total_loss, pred, gradients

def train_step(inputs, outputs):
    total_loss, pred, gradients=compiled_step(inputs, outputs)
    if args.weight_decay > 0.:
        gradients=[g+v*args.weight_decay for g,v in zip(gradients, model.trainable_variables)]

    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss.update_state(total_loss)
    train_accuracy.update_state(outputs, pred)

@tf.function()
def train_step_dist(inputs, outputs):
    strategy.run(train_step, args= (inputs, outputs))

请注意，此指令可用于单节点多 GPU 和多节点多 GPU。