TensorFlow - Amazon SageMaker
TensorFlow 模型TensorFlow 带有 Hugging Face 变形金刚的模型

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

TensorFlow

将自己的 TensorFlow 模型带到训练编译器中 SageMaker，然后使用训练编译器运行 SageMaker 训练作业。

TensorFlow 模型

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

提示

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

使用 Keras（推荐）

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

对于单个 GPU 训练

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

没有 Keras 的情况下

SageMaker 训练编译器不支持在中即时执行 TensorFlow。因此，您应该使用 TensorFlow 函数装饰器 (@tf.function) 包装模型和训练循环，以利用编译器加速。

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

对于单个 GPU 训练

TensorFlow 2.0 或更高版本默认开启了 Eager Execution，因此你应该在用于构造 TensorFlow 模型的每个函数前面添加@tf.function装饰器。

TensorFlow 带有 Hugging Face 变形金刚的模型

TensorFlow 带有 Hu g ging Fac e Transformers 的模型基 TensorFlow于 tf.keras.Model API。Hugging Face Transformers 还提供了预训练 TensorFlow 的模型类，以帮助减少配置自然语言处理 (NLP) 模型的工作量。使用 Transformers 库创建自己的训练脚本后，您可以使用带有训练编译器配置类的 SageMaker HuggingFace估算器运行 SageMaker 训练脚本，如上一主题所示。使用 TensorFlow 训练编译器运行 SageMaker 训练作业

SageMaker Trainging 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 Transform ers 文档的快速浏览页面所述，你可以将这些模型用作 TensorFlow 常规 Keras 模型。

对于单个 GPU 训练

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

对于分布式训练

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

  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 或更高版本默认开启了 Eager Execution，因此你应该在用于构造 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。