Abilitazione del checkpointing - Amazon SageMaker AI

Abilitazione del checkpointing

Dopo aver abilitato il checkpointing, SageMaker AI salva i checkpoint su Amazon S3 e sincronizza il job di addestramento con il bucket S3 del checkpoint. Puoi utilizzare i bucket di directory S3 generici o i bucket di directory S3 per il tuo bucket S3 di checkpoint.

Diagramma architettonico della scrittura dei checkpoint durante l’addestramento.

L’esempio seguente illustra come configurare i percorsi dei checkpoint quando si costruisce uno strumento di stima SageMaker AI. Per abilitare il checkpoint, aggiungi i parametri checkpoint_s3_uri e checkpoint_local_path allo strumento di valutazione.

Il seguente modello di esempio mostra come creare uno strumento di stima SageMaker AI generico e abilitare il checkpoint. È possibile utilizzare questo modello per gli algoritmi supportati specificando il parametro image_uri. Per trovare gli URI delle immagini Docker per gli algoritmi con checkpointing supportati da SageMaker AI, consulta Percorsi del registro Docker e codice di esempio. È inoltre possibile sostituire estimator e Estimator con altre classi principali e classi dello strumento di stima dei framework di SageMaker AI, ad esempio TensorFlow, PyTorch, MXNet, HuggingFace e XGBoost.

import sagemaker
from sagemaker.estimator import Estimator

bucket=sagemaker.Session().default_bucket()
base_job_name="sagemaker-checkpoint-test"
checkpoint_in_bucket="checkpoints"

# The S3 URI to store the checkpoints
checkpoint_s3_bucket="s3://{}/{}/{}".format(bucket, base_job_name, checkpoint_in_bucket)

# The local path where the model will save its checkpoints in the training container
checkpoint_local_path="/opt/ml/checkpoints"

estimator = Estimator(
    ...
    image_uri="<ecr_path>/<algorithm-name>:<tag>" # Specify to use built-in algorithms
    output_path=bucket,
    base_job_name=base_job_name,
    
    # Parameters required to enable checkpointing
    checkpoint_s3_uri=checkpoint_s3_bucket,
    checkpoint_local_path=checkpoint_local_path
)

I due parametri seguenti specificano i percorsi per il checkpoint:

  • checkpoint_local_path— Specifica il percorso locale in cui il modello salva periodicamente i checkpoint in un container di addestramento. Il percorso predefinito è '/opt/ml/checkpoints'. Se stai utilizzando altri framework o state utilizzando un container di addestramento personale, assicurati che la configurazione del checkpoint dello script di addestramento specifichi il percorso verso '/opt/ml/checkpoints'.

    Nota

    È consigliabile specificare i percorsi locali come '/opt/ml/checkpoints', in modo che siano coerenti con le impostazioni predefinite del checkpoint SageMaker AI. Se preferisci specificare il tuo percorso locale, assicurati di abbinare il percorso di salvataggio del checkpoint nello script di addestramento e il parametro checkpoint_local_path degli strumenti di stima SageMaker AI.

  • checkpoint_s3_uri— L'URI di un bucket S3 in cui i checkpoint sono memorizzati in tempo reale. È possibile specificare un bucket S3 generico o di directory per archiviare i checkpoint. Per ulteriori informazioni sui bucket di directory S3, consulta Bucket di directory nella Guida per l’utente di Amazon Simple Storage Service.

Per trovare un elenco completo dei parametri dello strumento di stima SageMaker AI, consulta API Estimator nella documentazione di Amazon SageMaker Python SDK.