チェックポイントを有効にする

チェックポイントを有効にすると、SageMaker AI はチェックポイントを Amazon S3 に保存し、トレーニングジョブをチェックポイント S3 バケットと同期します。チェックポイント S3 バケットには、S3 汎用バケットまたは S3 ディレクトリバケットを使用できます。

次の例は、SageMaker AI 推定器を構築するときにチェックポイントパスを設定する方法を示しています。チェックポイントを有効にするには、推定器に checkpoint_s3_uri パラメータと checkpoint_local_path パラメータを指定します。

次のサンプルテンプレートは、汎用 SageMaker AI 推定器を作成し、チェックポイントを有効にする方法を示しています。image_uri パラメータを指定すると、このテンプレートをサポート対象のアルゴリズムに使用できます。SageMaker AI でサポートされているチェックポイントを使用するアルゴリズムの Docker イメージ URIs を見つけるには、「Docker Registry Paths and Example Code」を参照してください。estimator と を、、、、 などの他の SageMaker AI フレームワークの推定器の親クラスTensorFlowPyTorchMXNetHuggingFaceと推定器クラスEstimatorに置き換えることもできます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
)

次の 2 つのパラメータは、チェックポイントのパスを指定します。

SageMaker AI 推定器パラメータの完全なリストについては、Amazon SageMaker Python SDK ドキュメントの「推定器 API」を参照してください。