使用 SageMaker Training Compiler 執行 TensorFlow 訓練任務

您可以使用任何 SageMaker AI 介面，透過 SageMaker Training Compiler 執行訓練任務：Amazon SageMaker Studio Classic、Amazon SageMaker 筆記本執行個體、適用於 Python (Boto3) 的 AWS SDK 和 AWS 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 的 AWS 深度學習容器 (DLC) 使用 TensorFlow 的改編版本，其中包含以開放原始碼 TensorFlow 架構為基礎的變更。SageMaker AI 架構深度學習容器已針對 AWS 基礎架構和 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，請參閱支援的架構。

# 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 架構 DLC 時，請考量下列缺陷：

例如，有一個已知問題，即 tensorflow/models 和 tensorflow/text 程式庫一律嘗試重新安裝開放原始碼 TensorFlow 的已知問題。如果您需要安裝這些程式庫，以為您的使用案例選擇特定版本，我們建議您查看適用於第 2.9 版或更新版本的 SageMaker AI TensorFlow DLC Dockerfile。Dockerfile 路徑通常採用下列格式：tensorflow/training/docker/<tensorflow-version>/py3/<cuda-version>/Dockerfile.gpu。在 Dockerfile 中，您會看到依序重新安裝 AWS 受管 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 Training Compiler

SageMaker Training Compiler 組態選項必須透過 CreateTrainingJob API 作業的請求語法中之 AlgorithmSpecification 和 HyperParameters 欄位來指定。

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

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

如要尋找已實作 SageMaker Training Compiler 的深度學習容器映像 URI 的完整清單，請參閱支援的架構。