Ejecutar trabajos de entrenamiento de TensorFlow con el Compilador de entrenamiento de SageMaker

Puede utilizar cualquiera de las interfaces de SageMaker AI para ejecutar un trabajo de entrenamiento con el Compilador de entrenamiento de SageMaker: Amazon SageMaker Studio Classic, instancias de cuadernos de Amazon SageMaker, AWS SDK para Python (Boto3) y AWS Command Line Interface.

Uso de SageMaker Python SDK

Para activar el Compilador de entrenamiento de SageMaker, añada el parámetro compiler_config al estimador TensorFlow o Hugging Face de SageMaker AI. Importe la clase TrainingCompilerConfig y pase una instancia de ella al parámetro compiler_config. Los siguientes ejemplos de código muestran la estructura de las clases de estimador de SageMaker AI con el Compilador de entrenamiento de SageMaker activado.

Para obtener información que se adapte a su caso de uso, consulte una de las siguientes opciones.

La siguiente lista es el conjunto mínimo de parámetros necesarios para ejecutar un trabajo de entrenamiento de SageMaker con el compilador.

disable_profiler=True,
debugger_hook_config=False

Si el trabajo de entrenamiento con el compilador se ha iniciado correctamente, recibirá los siguientes registros durante la fase de inicialización del trabajo:

Uso del SageMaker AI Python SDK y ampliación de los contenedores de aprendizaje profundo de marco de SageMaker AI

AWSLos Contenedores de aprendizaje profundo (DLC) para TensorFlow utilizan versiones adaptadas de TensorFlow que incluyen cambios en la estructura de código abierto de TensorFlow. Los Contenedores de aprendizaje profundo de marco de SageMaker AI están optimizados para la infraestructura de AWS subyacente y para Amazon SageMaker AI. Con la ventaja de utilizar los DLC, la integración del Compilador de entrenamiento de SageMaker añade más mejoras de rendimiento en comparación con el TensorFlow nativo. Además, puede crear un contenedor de entrenamiento personalizado ampliando la imagen del DLC.

Para ampliar y personalizar los DLC TensorFlow de SageMaker AI para su caso de uso, siga las instrucciones siguientes.

Cree un archivo Dockerfile.

Utilice la siguiente plantilla de Dockerfile para ampliar el DLC TensorFlow de SageMaker AI. Debe usar la imagen DLC TensorFlow de SageMaker AI como imagen base del contenedor de Docker. Para encontrar los URI de imagen del DLC TensorFlow de SageMaker AI, consulte Marcos compatibles.

# 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 obtener más información, consulte Paso 2: Crear y cargar los scripts de entrenamiento de Dockerfile y Python.

Tenga en cuenta las siguientes dificultades a la hora de ampliar los DLC de marco de SageMaker AI:

Por ejemplo, existe un problema conocido por el que las bibliotecas tensorflow/models y tensorflow/text siempre intentan reinstalar TensorFlow de código abierto. Si necesita instalar estas bibliotecas para elegir una versión específica para su caso de uso, le recomendamos que consulte los Dockerfiles de DLC TensorFlow de SageMaker para v2.9 o posterior. Las rutas a los Dockerfiles suelen tener el siguiente formato: tensorflow/training/docker/<tensorflow-version>/py3/<cuda-version>/Dockerfile.gpu. En los Dockerfiles, debería encontrar las líneas de código para reinstalar el binario de TensorFlow administrado por AWS (especificado en la variable de entorno TF_URL) y otras dependencias en orden. La sección de reinstalación debe ser similar a la siguiente especificación de ejemplo:

# 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

Compilar e insertar en ECR

Para compilar e insertar el contenedor de Docker en Amazon ECR, siga las instrucciones de los siguientes enlaces:

Ejecute con el estimador de SageMaker Python SDK

Utilice el estimador de marco TensorFlow de SageMaker AI de la forma habitual. Debe especificar image_uri para usar el nuevo contenedor alojado en 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 el Compilador de entrenamiento de SageMaker mediante la operación de la API CreateTrainingJob de SageMaker AI

Las opciones de configuración del Compilador de entrenamiento de SageMaker deben especificarse a través del campo AlgorithmSpecification y HyperParameters en la sintaxis de solicitud de la operación de API CreateTrainingJob.

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

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

Para encontrar una lista completa de los URI de imágenes de contenedores de aprendizaje profundo que tienen implementado el Compilador de entrenamiento de Sagemaker, consulte Marcos admitidos.