トレーニングジョブを準備して TensorBoard 出力データ設定を収集する

SageMaker AI での機械学習の一般的なトレーニングジョブは、トレーニングスクリプトの準備と、SageMaker AI Python SDK の SageMaker AI ModelTrainer (以前の PySDK v2 の推定器) オブジェクトの設定という 2 つの主要なステップで構成されます。このセクションでは、SageMaker トレーニングジョブから TensorBoard 互換データを収集するために必要な変更を説明します。

前提条件

次のリストは、SageMaker AI with TensorBoard の使用を開始するための前提条件を示しています。

ステップ 1: オープンソースの TensorBoard ヘルパーツールを使用してトレーニングスクリプトを変更する

収集する出力テンソルとスカラーを決定し、TensorBoardX、TensorFlow サマリーライター、PyTorch サマリーライター、または SageMaker Debugger のいずれかのツールを使用してトレーニングスクリプトのコード行を変更してください。

また、トレーニングコンテナのコールバックのログディレクトリ (log_dir) として TensorBoard データ出力パスを指定していることを確認してください。

フレームワークごとのコールバックについて詳しくは、次のリソースを参照してください。

ステップ 2: TensorBoard 出力設定を使用して SageMaker トレーニング ModelTrainer オブジェクトを作成する

SageMaker AI ModelTrainer を設定するsagemaker.debugger.TensorBoardOutputConfigときは、 を使用します。この設定 API は、TensorBoard データを保存するために指定した S3 バケットをトレーニングコンテナ (/opt/ml/output/tensorboard) のローカルパスにマッピングします。モジュールの オブジェクトを ModelTrainer クラスの tensorboard_output_configパラメータに渡します。次のコードスニペットは、TensorBoard 出力設定パラメータを使用して TensorFlow ModelTrainer を準備する例を示しています。 TensorBoard

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

import os
from sagemaker.train import ModelTrainer
from sagemaker.train.configs import SourceCode
from sagemaker.core.debugger import TensorBoardOutputConfig
from sagemaker.train.configs import Compute

# 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
)

model_trainer = ModelTrainer(
    source_code=SourceCode(entry_script="train.py", source_dir="src"),
    role=role,
    training_image=image_uri,
    compute=Compute(instance_count=1, instance_type="ml.c5.xlarge"),
    base_job_name="your-training_job_name",
    tensorboard_output_config=tensorboard_output_config,
    hyperparameters=hyperparameters
)