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

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

次の例は、SageMaker AI トレーニングオブジェクトを構築するときにチェックポイントパスを設定する方法を示しています。

チェックポイントを有効にするには、 checkpoint_configパラメータを に追加しますModelTrainer。次のサンプルテンプレートは、SageMaker AI を作成しModelTrainer、チェックポイントを有効にする方法を示しています。このテンプレートは、 training_imageパラメータを指定することで、サポートされている任意のアルゴリズムに使用できます。SageMaker AI がサポートするチェックポイントアルゴリズムの Docker イメージ URI を見つけるには、「Docker Registry Paths and Example Code」を参照してください。V3 では、統合ModelTrainerクラスはすべてのフレームワーク固有の推定器クラス (TensorFlow、PyTorch、HuggingFace、XGBoost など) を置き換えます。

from sagemaker.train import ModelTrainer
from sagemaker.train.configs import Compute, CheckpointConfig
from sagemaker.core.helper.session_helper import Session

bucket = 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)

model_trainer = ModelTrainer(
    training_image="<ecr_path>/<algorithm-name>:<tag>",
    role=role,
    compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
    base_job_name=base_job_name,
    checkpoint_config=CheckpointConfig(
        s3_uri=checkpoint_s3_bucket,
        local_path="/opt/ml/checkpoints"
    )
)

checkpoint_config パラメータは、次のフィールドを持つCheckpointConfigオブジェクトを受け入れます。

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