使用 SageMaker Training Compiler 运行 TensorFlow 训练作业

您可以使用任意 SageMaker AI 界面通过 SageMaker Training Compiler 运行训练作业：Amazon SageMaker Studio Classic、Amazon SageMaker Notebook 实例、适用于 Python (Boto3) 的 Amazon SDK和 Amazon Command Line Interface。

使用 SageMaker Python SDK

要开启 SageMaker Training Compiler，请将 compiler_config 参数添加到 SageMaker AI TensorFlow 或 Hugging Face 估算器中。导入 TrainingCompilerConfig 类，并将它的一个实例传递给 compiler_config 参数。以下代码示例显示了已开启 SageMaker Training Compiler 的 SageMaker AI 估算器类的结构。

有关适合您的使用案例的信息，请参阅下列选项之一。

以下列表是使用编译器运行 SageMaker 训练作业所需的最少参数集。

disable_profiler=True,
debugger_hook_config=False

如果使用编译器成功启动训练作业，则在作业初始化阶段将收到以下日志：

使用 SageMaker AI Python SDK 和扩展 SageMaker AI 框架深度学习容器

适用于 TensorFlow 的 Amazon 深度学习容器 (DLC) 使用 TensorFlow 的更改后版本，其中包括基于开源 TensorFlow 框架的更改。SageMaker AI Framework Deep Learning Containers 针对底层 Amazon 基础设施和 Amazon SageMaker AI 进行了优化。凭借使用 DLC 所获得的优势，SageMaker Training Compiler 集成在原生 TensorFlow 的基础上增加了更多的性能改进。此外，您可以通过扩展 DLC 映像来创建自定义训练容器。

要针对您的使用案例扩展和自定义 SageMaker AI TensorFlow DLC，请按照以下说明进行操作。

创建 Dockerfile

使用以下 Dockerfile 模板扩展 SageMaker AI TensorFlow DLC。您必须使用 SageMaker AI TensorFlow DLC 映像作为 Docker 容器的基本映像。要查找 SageMaker AI TensorFlow DLC 映像 URI，请参阅 Supported Frameworks。

# SageMaker AI TensorFlow Deep Learning Container image
FROM 763104351884.dkr.ecr.<aws-region>.amazonaws.com/tensorflow-training:<image-tag>

ENV PATH="/opt/ml/code:${PATH}"

# This environment variable is used by the SageMaker AI container 
# to determine user code directory.
ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code

# Add more code lines to customize for your use-case
...

有关更多信息，请参阅步骤 2：创建并上传 Dockerfile 和 Python 训练脚本。

在扩展 SageMaker AI Framework DLC 时，请考虑以下缺陷：

例如，有一个已知问题，即 tensorflow/models 和 tensorflow/text 库总是尝试重新安装开源 TensorFlow。如果您需要安装这些库来为您的使用案例选择特定版本，我们建议您考察 2.9 或更高版本的 SageMaker AI TensorFlow DLC Dockerfiles。Dockerfiles 的路径通常采用以下格式：tensorflow/training/docker/<tensorflow-version>/py3/<cuda-version>/Dockerfile.gpu。在 Dockerfiles 中，您应找到代码行来按顺序重新安装 Amazon 管理的 TensorFlow 二进制文件（指定给 TF_URL 环境变量）和其他依赖项。重新安装部分应与以下示例类似：

# tf-models does not respect existing installations of TensorFlow 
# and always installs open source TensorFlow

RUN pip3 install --no-cache-dir -U \
    tf-models-official==x.y.z

RUN pip3 uninstall -y tensorflow tensorflow-gpu \
  ; pip3 install --no-cache-dir -U \
    ${TF_URL} \
    tensorflow-io==x.y.z \
    tensorflow-datasets==x.y.z

构建并推送到 ECR

要构建 Docker 容器并将其推送到 Amazon ECR，请按照以下链接中的说明进行操作：

使用 SageMaker Python SDK 估算器运行

像往常一样使用 SageMaker AI TensorFlow 框架估算器。您必须指定 image_uri 以使用您在 Amazon ECR 中托管的新容器。

import sagemaker, boto3
from sagemaker import get_execution_role
from sagemaker.tensorflow import TensorFlow, TrainingCompilerConfig

account_id = boto3.client('sts').get_caller_identity().get('Account')
ecr_repository = 'tf-custom-container-test'
tag = ':latest'

region = boto3.session.Session().region_name

uri_suffix = 'amazonaws.com'

byoc_image_uri = '{}.dkr.ecr.{}.{}/{}'.format(
    account_id, region, uri_suffix, ecr_repository + tag
)

byoc_image_uri
# This should return something like
# 111122223333.dkr.ecr.us-east-2.amazonaws.com/tf-custom-container-test:latest

estimator = TensorFlow(
    image_uri=image_uri,
    role=get_execution_role(),
    base_job_name='tf-custom-container-test-job',
    instance_count=1,
    instance_type='ml.p3.8xlarge'
    compiler_config=TrainingCompilerConfig(),
    disable_profiler=True,
    debugger_hook_config=False
)

# Start training
estimator.fit()

使用 SageMaker AI CreateTrainingJob API 操作启用 SageMaker AI Training Compiler

必须在 CreateTrainingJobAPI 操作的请求语法中使用 AlgorithmSpecification 和 HyperParameters 字段来指定 SageMaker Training Compiler 配置选项。

"AlgorithmSpecification": {
    "TrainingImage": "<sagemaker-training-compiler-enabled-dlc-image>"
},

"HyperParameters": {
    "sagemaker_training_compiler_enabled": "true",
    "sagemaker_training_compiler_debug_mode": "false"
}

要查找已实施 SageMaker Training Compiler 的深度学习容器映像 URI 的完整列表，请参阅支持的框架。