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

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

前提条件

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

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

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

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

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

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

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

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

from sagemaker.tensorflow import TensorFlow
from sagemaker.debugger import TensorBoardOutputConfig

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

estimator = TensorFlow(
    entry_point="train.py",
    source_dir="src",
    role=role,
    image_uri=image_uri,
    instance_count=1,
    instance_type="ml.c5.xlarge",
    base_job_name="your-training_job_name",
    tensorboard_output_config=tensorboard_output_config,
    hyperparameters=hyperparameters
)