訓練程式碼可封裝為自訂訓練元件,並在管道工作中執行。TensorBoard 記錄會自動串流至 Vertex AI TensorBoard 實驗。您可以使用這項整合功能,在 Vertex AI TensorBoard 記錄檔寫入 Cloud Storage 時,以近乎即時的方式監控訓練作業。
如需初始設定,請參閱「為 Vertex AI TensorBoard 設定」。
訓練指令碼異動
訓練指令碼必須設為將 TensorBoard 記錄寫入 Cloud Storage 值區,Vertex AI 訓練服務會透過預先定義的環境變數 AIP_TENSORBOARD_LOG_DIR
自動提供該位置。
通常只要將 os.environ['AIP_TENSORBOARD_LOG_DIR']
提供給開放原始碼 TensorBoard 記錄寫入 API 做為記錄目錄,即可完成這項操作。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 如何在自訂訓練環境中設定環境變數。
建構及執行管道
以下範例說明如何使用 Kubeflow Pipelines DSL 套件建構及執行管道。如需更多範例和詳細資訊,請參閱 Vertex AI Pipelines 說明文件。
建立訓練元件
將訓練程式碼封裝至自訂元件,確保程式碼已設定為將 TensorBoard 記錄寫入 Cloud Storage 值區。如需更多範例,請參閱「建構自己的管道元件」。
from kfp.v2.dsl import component
@component(
base_image="tensorflow/tensorflow:latest",
packages_to_install=["tensorflow_datasets"],
)
def train_tensorflow_model_with_tensorboard():
import datetime, os
import tensorflow as tf
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
def create_model():
return tf.keras.models.Sequential(
[
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(512, activation="relu"),
]
)
model = create_model()
model.compile(
optimizer="adam",
loss="sparse_categorical_crossentropy",
metrics=["accuracy"]
)
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=5,
validation_data=(x_test, y_test),
callbacks=[tensorboard_callback],
)
建構及編譯管道
在 create_custom_training_job_op_from_component
中指定元件規格,即可根據您建立的元件建立自訂訓練工作。將 tensorboard_resource_name
設為 TensorBoard 例項,並將 staging_bucket
設為 API 呼叫期間的構件暫存位置 (包括 TensorBoard 記錄)。
接著,請建構管道來納入這項工作,並將管道編譯為 JSON 檔案。
from kfp.v2 import compiler
from google_cloud_pipeline_components.v1.custom_job.utils import \
create_custom_training_job_op_from_component
from kfp.v2 import dsl
def create_tensorboard_pipeline_sample(
project, location, staging_bucket, display_name, service_account, experiment, tensorboard_resource_name
):
@dsl.pipeline(
pipeline_root=f"{staging_bucket}/pipeline_root",
name=display_name,
)
def pipeline():
custom_job_op = create_custom_training_job_op_from_component(
component_spec=train_tensorflow_model_with_tensorboard,
tensorboard=tensorboard_resource_name,
base_output_directory=staging_bucket,
service_account=service_account,
)
custom_job_op(project=project, location=location)
compiler.Compiler().compile(
pipeline_func=pipeline, package_path=f"{display_name}.json"
)
提交 Vertex AI 管道
使用 Python 適用的 Vertex AI SDK 提交管道。詳情請參閱「執行管道」。
Python 適用的 Vertex AI SDK
experiment_name
:為實驗命名。pipeline_job_display_name
:管道工作顯示名稱。template_path
:已編譯管道範本的路徑。pipeline_root
:指定管道服務帳戶可存取的 Cloud Storage URI。管道執行作業的構件會儲存在管道根目錄中。parameter_values
:要傳遞至這項執行作業的管道參數。例如,建立dict()
時,可將參數名稱做為字典鍵,參數值做為字典值。project
:您的專案 ID。要執行管道的 Google Cloud 專案。您可以在 Google Cloud 控制台歡迎頁面中找到 ID。
location
:管道執行的位置。這個位置應與您使用的 TensorBoard 執行個體相同。
後續步驟
- 查看結果:查看 Vertex AI Pipelines 的 TensorBoard。
- 瞭解如何使用 Cloud Profiler 最佳化自訂訓練工作效能。