Execute trabalhos de treinamento do TensorFlow com o SageMaker Training Compiler

Você pode usar qualquer uma das interfaces do SageMaker AI para executar uma tarefa de treinamento com o SageMaker Training Compiler: Amazon SageMaker Studio Classic, instâncias de caderno do Amazon SageMaker, AWS SDK para Python (Boto3) e AWS Command Line Interface.

Usando o SageMaker Python SDK

Para ativar o SageMaker Training Compiler, adicione o parâmetro compiler_config aos estimadores do TensorFlow ou Hugging Face do SageMaker AI. Importe a classe TrainingCompilerConfig e passe uma instância dela para o parâmetro compiler_config. Os exemplos de código a seguir mostram a estrutura de classes de estimadores do SageMaker AI com o SageMaker Training Compiler ativado.

Para obter informações adequadas ao seu caso de uso, consulte uma das opções a seguir.

A lista a seguir é o conjunto mínimo de parâmetros necessários para executar um trabalho de treinamento do SageMaker com o compilador.

disable_profiler=True,
debugger_hook_config=False

Se o trabalho de treinamento com o compilador for iniciado com êxito, você receberá os seguintes logs durante a fase de inicialização do trabalho:

Usar o SageMaker AI Python SDK e estender os contêineres de deep learning do framework do SageMaker AI

Os Contêineres de Aprendizado Profundo AWS (DLC) para TensorFlow usam versões adaptadas do TensorFlow que incluem alterações no framework de código aberto do TensorFlow. Os contêineres de deep learning do framework do SageMaker AI são otimizados para a infraestrutura subjacente da AWS e para o Amazon SageMaker AI. Com a vantagem de usar os DLCs, a integração do SageMaker Training Compiler adiciona mais melhorias de desempenho em relação ao TensorFlow nativo. Além disso, você pode criar um contêiner de treinamento personalizado estendendo a imagem do DLC.

Para estender e personalizar os DLCs do TensorFlow do SageMaker AI para seu caso de uso, use as instruções a seguir.

Crie um Dockerfile.

Use o modelo do Dockerfile a seguir para estender o DLC do TensorFlow do SageMaker AI. Você deve usar a imagem de DLC do TensorFlow do SageMaker AI como imagem base do seu contêiner do Docker. Para encontrar os URIs de imagem do DLC do TensorFlow do SageMaker AI, consulte Frameworks compatíveis.

# 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
...

Para mais informações, consulte Etapa 2: Como criar e fazer upload dos scripts de treinamento do Dockerfile e do Python.

Considere as seguintes armadilhas ao estender os DLCs do framework do SageMaker AI:

Por exemplo, há um problema conhecido de que as bibliotecas tensorflow/modelos e tensorflow/texto sempre tentam reinstalar o TensorFlow de código aberto. Se precisar instalar essas bibliotecas para escolher uma versão específica para seu caso de uso, recomendamos que você consulte os arquivos Dockerfile do DLC do TensorFlow do SageMaker AI para v2.9 ou posterior. Os caminhos para os Dockerfiles geralmente estão no seguinte formato: tensorflow/training/docker/<tensorflow-version>/py3/<cuda-version>/Dockerfile.gpu. Nos Dockerfiles, você deve encontrar as linhas de código para reinstalar o binário TensorFlow gerenciado da AWS (especificado para a variável de ambiente TF_URL) e outras dependências em ordem. A seção de Reinstalar deve se parecer com o seguinte exemplo:

# 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

Compile e envie para o ECR

Para compilar e enviar seu contêiner do Docker para o Amazon ECR, siga as instruções nos links a seguir.

Execute usando o Estimador do SageMaker Python SDK

Use o estimador de framework TensorFlow do SageMaker AI como de costume. Você deve especificar image_uri para usar o novo contêiner que você hospedou no 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()

Habilitar o SageMaker Training Compiler usando a operação de API CreateTrainingJob do SageMaker AI

As opções de configuração do SageMaker Training Compiler devem ser especificadas por meio dos campos AlgorithmSpecification e HyperParameters na sintaxe da solicitação para a operação CreateTrainingJob da API.

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

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

Para encontrar uma lista completa de URIs de imagens de contêiner de aprendizado profundo que têm o SageMaker Training Compiler implementado, consulte Estruturas compatíveis.