Preparación de un trabajo de entrenamiento para recopilar los datos de salida de TensorBoard

Un trabajo de entrenamiento típico para machine learning en SageMaker consta de dos pasos principales: la preparación de un script de entrenamiento y la configuración de un objeto estimator de SageMaker del SageMaker Python SDK. En esta sección, conocerá los cambios necesarios para recopilar datos compatibles con TensorBoard de los trabajos de entrenamiento de SageMaker.

Requisitos previos

En la siguiente lista se muestran los requisitos previos para empezar a utilizar SageMaker con TensorBoard.

Paso 1: modificación del script de entrenamiento con las herramientas auxiliares de TensorBoard de código abierto

Asegúrese de determinar qué escalares y tensores de salida desea recopilar, y modifique las líneas de código del script de entrenamiento con cualquiera de las siguientes herramientas: TensorBoardX, TensorFlow Summary Writer, PyTorch Summary Writer o SageMaker Debugger.

Asegúrese también de especificar la ruta de salida de datos de TensorBoard como el directorio de registro (log_dir) para la devolución de llamadas en el contenedor de entrenamiento.

Para obtener más información sobre las devoluciones de llamadas por marco, consulte los siguientes recursos.

Paso 2: creación de un objeto de estimador de entrenamiento de SageMaker con la configuración de salida de TensorBoard

Use sagemaker.debugger.TensorBoardOutputConfig al configurar un estimador del marco de SageMaker. Esta API de configuración asigna el bucket de S3 que especifica para guardar los datos de TensorBoard con la ruta local en el contenedor de entrenamiento (/opt/ml/output/tensorboard). Transfiere el objeto del módulo al parámetro tensorboard_output_config de la clase estimator. En el siguiente fragmento de código, se muestra un ejemplo de cómo preparar un estimador de TensorFlow con el parámetro de configuración de salida de TensorBoard.

"TensorBoardOutputConfig": { 
  "LocalPath": "/opt/ml/output/tensorboard",
  "S3OutputPath": "s3_output_bucket"
}

from sagemaker.tensorflow import TensorFlow
from sagemaker.debugger import TensorBoardOutputConfig

# Set variables for training job information, 
# such as s3_out_bucket and other unique tags.
... 

LOG_DIR="/opt/ml/output/tensorboard"

output_path = os.path.join(
    "s3_output_bucket", "sagemaker-output", "date_str", "your-training_job_name"
)

tensorboard_output_config = TensorBoardOutputConfig(
    s3_output_path=os.path.join(output_path, 'tensorboard'),
    container_local_output_path=LOG_DIR
)

estimator = TensorFlow(
    entry_point="train.py",
    source_dir="src",
    role=role,
    image_uri=image_uri,
    instance_count=1,
    instance_type="ml.c5.xlarge",
    base_job_name="your-training_job_name",
    tensorboard_output_config=tensorboard_output_config,
    hyperparameters=hyperparameters
)