カスタムトレーニングで Vertex AI TensorBoard を使用する

カスタムトレーニングを使用してモデルをトレーニングする際に、Vertex AI TensorBoard のログを Vertex AI TensorBoard に自動的にアップロードするようにトレーニングジョブを設定できます。

Vertex AI TensorBoard では、Cloud Storage に Vertex AI TensorBoard ログが書き込まれる際にそれをストリーミングするため、この統合を使用することで、ほぼリアルタイムにトレーニングをモニタリングできます。

初期設定については、Vertex AI TensorBoard を設定するをご覧ください。

トレーニングスクリプトの変更点

トレーニングスクリプトは、TensorBoard ログを Cloud Storage バケットに書き込むように構成する必要があります。Cloud Strage バケットの場所は、Vertex AI Training サービスにより事前定義の環境変数 AIP_TENSORBOARD_LOG_DIR で自動的に利用可能にされます。

通常これは、オープンソースの TensorBoard ログ書き込み API に、ログディレクトリとして os.environ['AIP_TENSORBOARD_LOG_DIR'] を指定することで実現できます。AIP_TENSORBOARD_LOG_DIR の場所は、staging_bucket 変数で設定します。

TensorFlow 2.x でトレーニングスクリプトを構成するには、TensorBoard コールバックを作成して、log_dir 変数を os.environ['AIP_TENSORBOARD_LOG_DIR'] に設定します。これにより、TensorBoard コールバックが、TensorFlow model.fit コールバックのリストに含まれます。

  tensorboard_callback = tf.keras.callbacks.TensorBoard(
       log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],
       histogram_freq=1
  )

  model.fit(
       x=x_train,
       y=y_train,
       epochs=epochs,
       validation_data=(x_test, y_test),
       callbacks=[tensorboard_callback],
  )

詳しくは、Vertex AI がカスタムトレーニング環境で環境変数を設定する方法をご覧ください。

カスタムトレーニングジョブを作成する

次の例は、独自のカスタムトレーニングジョブを作成する方法を示しています。

カスタムトレーニングジョブを作成する方法のサンプルについては、Hello カスタムトレーニングをご覧ください。カスタムトレーニングコンテナを構築する手順については、トレーニング用のカスタムコンテナイメージの作成をご覧ください。

カスタムトレーニングジョブを作成するには、Vertex AI SDK for Python または REST を使用します。

Python

def create_training_pipeline_custom_job_sample(
    project: str,
    location: str,
    staging_bucket: str,
    display_name: str,
    script_path: str,
    container_uri: str,
    model_serving_container_image_uri: str,
    dataset_id: Optional[str] = None,
    model_display_name: Optional[str] = None,
    args: Optional[List[Union[str, float, int]]] = None,
    replica_count: int = 0,
    machine_type: str = "n1-standard-4",
    accelerator_type: str = "ACCELERATOR_TYPE_UNSPECIFIED",
    accelerator_count: int = 0,
    training_fraction_split: float = 0.8,
    validation_fraction_split: float = 0.1,
    test_fraction_split: float = 0.1,
    sync: bool = True,
    tensorboard_resource_name: Optional[str] = None,
    service_account: Optional[str] = None,
):
    aiplatform.init(project=project, location=location, staging_bucket=staging_bucket)

    job = aiplatform.CustomTrainingJob(
        display_name=display_name,
        script_path=script_path,
        container_uri=container_uri,
        model_serving_container_image_uri=model_serving_container_image_uri,
    )

    # This example uses an ImageDataset, but you can use another type
    dataset = aiplatform.ImageDataset(dataset_id) if dataset_id else None

    model = job.run(
        dataset=dataset,
        model_display_name=model_display_name,
        args=args,
        replica_count=replica_count,
        machine_type=machine_type,
        accelerator_type=accelerator_type,
        accelerator_count=accelerator_count,
        training_fraction_split=training_fraction_split,
        validation_fraction_split=validation_fraction_split,
        test_fraction_split=test_fraction_split,
        sync=sync,
        tensorboard=tensorboard_resource_name,
        service_account=service_account,
    )

    model.wait()

    print(model.display_name)
    print(model.resource_name)
    print(model.uri)
    return model

project: 実際のプロジェクト ID。これらの ID は、Google Cloud コンソールの [ようこそ] ページで確認できます。
location: CustomJob を実行するリージョン。指定した TensorBoard インスタンスと同じリージョンにする必要があります。
staging_bucket: API 呼び出し中にアーティファクトをステージングするための Cloud Storage バケット（TensorBoard ログなど）。
display_name: カスタムトレーニングジョブの表示名。
script_path: ローカルファイルシステム上のワーキングディレクトリから、トレーニングコードのエントリポイントとなるスクリプトへの相対パス。
container_uri: トレーニングコンテナイメージの URI は、Vertex AI のビルド済みトレーニングコンテナ、またはカスタムコンテナです。
model_serving_container_image_uri: トレーニングスクリプトによって生成されたモデルを提供するのに適したモデルサービングコンテナの URI。
dataset_id: トレーニングに使用するデータセットの ID 番号。
model_display_name: トレーニング済みモデルの表示名。
args: Python スクリプトに渡すコマンドライン引数。
replica_count: 使用するワーカーレプリカの数。ほとんどの場合、最初のワーカープールでは 1 に設定します。
machine_type: 使用する VM のタイプ。サポートされている VM の一覧については、マシンタイプをご覧ください。
accelerator_type:（省略可）リソースプール内の各 VM にアタッチする GPU のタイプ。サポートされている GPU の一覧については、GPU をご覧ください。
accelerator_count:（省略可）リソースプール内の各 VM にアタッチする GPU の数。デフォルト値は 1 です。
training_fraction_split: モデルのトレーニングに使用するデータセットの割合。
validation_fraction_split: モデルの検証に使用するデータセットの割合。
test_fraction_split: モデルの評価に使用するデータセットの割合。
sync: このメソッドを同期的に実行するかどうか。
tensorboard_resource_name: CustomJob が TensorBoard ログをアップロードする先の Vertex TensorBoard インスタンスのリソース名。
service_account: TensorBoard で実行する場合は必須です。必要な権限を持つサービスアカウントを作成するをご覧ください。

REST

リクエストのデータを使用する前に、次のように置き換えます。

LOCATION_ID: 使用するリージョン。
PROJECT_ID: 実際のプロジェクト ID。
TENSORBOARD_INSTANCE_NAME:（必須）Vertex AI TensorBoard のログを保存する既存の Vertex AI TensorBoard インスタンスのフルネーム。
projects/PROJECT_ID/locations/LOCATION_ID/tensorboards/TENSORBOARD_INSTANCE_ID
注: これが既存の TensorBoard インスタンスでない場合、customJobs を作成すると 404 がスローされます。
GCS_BUCKET_NAME:「${PROJECT_ID}-tensorboard-logs-${LOCATION}」
USER_SA_EMAIL:（必須）前の手順で作成したサービスアカウント、または独自のサービスアカウント。「USER_SA_NAME@${PROJECT_ID}.iam.gserviceaccount.com」
TRAINING_CONTAINER: TRAINING_CONTAINER。
INVOCATION_TIMESTAMP:「$(date +'%Y%m%d-%H%M%S')」
JOB_NAME:「tensorboard-example-job-${INVOCATION_TIMESTAMP}」
BASE_OUTPUT_DIR:（必須）トレーニングのすべての出力が書き込まれる Google Cloud パス。「gs://$GCS_BUCKET_NAME/$JOB_NAME」

HTTP メソッドと URL:

POST https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/customJobs

リクエストの本文（JSON）:

{
"displayName": JOB_NAME,
"jobSpec":{
"workerPoolSpecs":[
  {
    "replicaCount": "1",
     "machineSpec": {
        "machineType": "n1-standard-8",
      },
      "containerSpec": {
        "imageUri": TRAINING_CONTAINER,
      }
    }
  ],

  "base_output_directory": {
  "output_uri_prefix": BASE_OUTPUT_DIR,
   },
  "serviceAccount": USER_SA_EMAIL,
  "tensorboard": TENSORBOARD_INSTANCE_NAME,
  }
}

リクエストを送信するには、次のいずれかのオプションを展開します。

curl（Linux、macOS、Cloud Shell）

注: 次のコマンドは、gcloud init または gcloud auth login を実行して、ユーザーアカウントで gcloud CLI にログインしているか、Cloud Shell を使用して自動的に gcloud CLI にログインしていることを前提としています。gcloud auth list を実行すると、現在アクティブなアカウントを確認できます。

リクエスト本文を request.json という名前のファイルに保存します。ターミナルで次のコマンドを実行して、このファイルを現在のディレクトリに作成または上書きします。

cat > request.json << 'EOF'
{
"displayName": JOB_NAME,
"jobSpec":{
"workerPoolSpecs":[
  {
    "replicaCount": "1",
     "machineSpec": {
        "machineType": "n1-standard-8",
      },
      "containerSpec": {
        "imageUri": TRAINING_CONTAINER,
      }
    }
  ],

  "base_output_directory": {
  "output_uri_prefix": BASE_OUTPUT_DIR,
   },
  "serviceAccount": USER_SA_EMAIL,
  "tensorboard": TENSORBOARD_INSTANCE_NAME,
  }
}
EOF

その後、次のコマンドを実行して REST リクエストを送信します。

curl -X POST \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json; charset=utf-8" \
    -d @request.json \
    "https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/customJobs"

PowerShell（Windows）

@'
{
"displayName": JOB_NAME,
"jobSpec":{
"workerPoolSpecs":[
  {
    "replicaCount": "1",
     "machineSpec": {
        "machineType": "n1-standard-8",
      },
      "containerSpec": {
        "imageUri": TRAINING_CONTAINER,
      }
    }
  ],

  "base_output_directory": {
  "output_uri_prefix": BASE_OUTPUT_DIR,
   },
  "serviceAccount": USER_SA_EMAIL,
  "tensorboard": TENSORBOARD_INSTANCE_NAME,
  }
}
'@  | Out-File -FilePath request.json -Encoding utf8

その後、次のコマンドを実行して REST リクエストを送信します。

$cred = gcloud auth print-access-token
$headers = @{ "Authorization" = "Bearer $cred" }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/customJobs" | Select-Object -Expand Content

次のような JSON レスポンスが返されます。

{
  "name": "projects/PROJECT_ID/locations/LOCATION_ID/customJobs/CUSTOM_JOB_ID",
  "displayName": "DISPLAY_NAME",
  "jobSpec": {
    "workerPoolSpecs": [
      {
        "machineSpec": {
          "machineType": "n1-standard-8"
        },
        "replicaCount": "1",
        "diskSpec": {
          "bootDiskType": "pd-ssd",
          "bootDiskSizeGb": 100
        },
        "containerSpec": {
          "imageUri": "IMAGE_URI"
        }
      }
    ],
    "serviceAccount": "SERVICE_ACCOUNT",
    "baseOutputDirectory": {
      "outputUriPrefix": "OUTPUT_URI_PREFIX"
    },
    "tensorboard": "projects//locations/LOCATION_ID/tensorboards/tensorboard-id"
  },
  "state": "JOB_STATE_PENDING",
  "createTime": "CREATE-TIME",
  "updateTime": "UPDATE-TIME"
}

次のステップ

Vertex AI TensorBoard を表示するを参照する。
Vertex AI TensorBoard Profiler を使用して、カスタムトレーニングジョブのパフォーマンスを最適化する方法を学習する。

カスタム トレーニングで Vertex AI TensorBoard を使用する

トレーニング スクリプトの変更点

カスタム トレーニング ジョブを作成する

Python

Python

REST

curl（Linux、macOS、Cloud Shell）

PowerShell（Windows）

次のステップ

カスタムトレーニングで Vertex AI TensorBoard を使用する

トレーニングスクリプトの変更点

カスタムトレーニングジョブを作成する