Aktivieren von Checkpointing - Amazon SageMaker AI

Aktivieren von Checkpointing

Nachdem Sie Checkpointing aktiviert haben, speichert SageMaker AI Checkpoints in Amazon S3 und synchronisiert Ihren Trainingsjob mit dem Checkpoint S3-Bucket. Sie können entweder S3-Allzweck-Buckets oder S3-Verzeichnis-Buckets für Ihren Checkpoint-S3-Bucket verwenden.

Architekturdiagramm des Schreibens von Checkpoints während des Trainings.

Im folgenden Beispiel wird gezeigt, wie Checkpoint-Pfade konfiguriert werden, wenn Sie einen SageMaker-AI-Schätzer erstellen. Um Checkpointing zu aktivieren, fügen Sie Ihrem Schätzer die Parameter checkpoint_s3_uri und checkpoint_local_path hinzu.

Das folgende Beispiel zeigt eine Vorlage für die Erstellung eines generischen SageMaker-AI-Schätzers und die Aktivierung von Checkpoints. Sie können diese Vorlage für die unterstützten Algorithmen verwenden, indem Sie den image_uri-Parameter angeben. Docker-Image-URIs für Algorithmen mit Checkpointing, die von SageMaker AI unterstützt werden, finden Sie unter Docker-Registry-Pfade und Beispielcode. Sie können estimator und Estimator auch durch übergeordnete Schätzerklassen und Schätzerklassen anderer SageMaker-AI-Frameworks ersetzen, z. B. TensorFlow, PyTorch, MXNet, HuggingFace und 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
)

Die folgenden beiden Parameter spezifizieren Pfade für Checkpoints:

  • checkpoint_local_path– Geben Sie den lokalen Pfad an, unter dem das Modell die Checkpoints regelmäßig in einem Trainingscontainer speichert. Der Standardpfad ist auf '/opt/ml/checkpoints' gesetzt. Wenn Sie andere Frameworks verwenden oder Ihren eigenen Trainingscontainer mitbringen, stellen Sie sicher, dass die Checkpoint-Konfiguration Ihres Trainingsskripts den Pfad zu '/opt/ml/checkpoints' angibt.

    Anmerkung

    Wir empfehlen, die lokalen Pfade als '/opt/ml/checkpoints' anzugeben, damit sie mit den Standardeinstellungen für SageMaker-AI-Checkpoints konsistent sind. Wenn Sie es vorziehen, Ihren eigenen lokalen Pfad anzugeben, stellen Sie sicher, dass Sie den Checkpoint-Speicherpfad in Ihrem Trainingsskript und den checkpoint_local_path Parametern der SageMaker-AI-Schätzer übereinstimmen.

  • checkpoint_s3_uri – Der URI zu einem S3-Bucket, in dem die Checkpoints in Echtzeit gespeichert werden. Sie können entweder einen S3-Allzweck-Bucket oder einen S3-Verzeichnis-Bucket zum Speichern Ihrer Checkpoints angeben. Weitere Informationen über S3-Verzeichnis-Buckets finden Sie unter Verzeichnis-Buckets im Benutzerhandbuch für Amazon Simple Storage Service.

Eine vollständige Liste der SageMaker-AI-Schätzparameter finden Sie in der Estimator-API in der Dokumentation zum Amazon SageMaker Python SDK.