Usar o TensorBoard da Vertex AI com o Vertex AI Pipelines

O código de treinamento pode ser empacotado em um componente de treinamento personalizado e executado em um job de pipeline. Os registros do TensorBoard são transmitidos automaticamente para o experimento do Vertex AI TensorBoard. Essa integração é necessária para monitorar o treinamento quase em tempo real, já que o TensorBoard da Vertex AI fará streaming nos registros do TensorBoard Vertex AI à medida que eles forem gravados no Cloud Storage.

Para a configuração inicial, consulte Configuração para o TensorBoard da Vertex AI.

Alterações no script de treinamento

Seu script de treinamento precisa ser configurado para gravar registros do TensorBoard no bucket do Cloud Storage, o local do serviço de treinamento do Vertex AI que será disponibilizado automaticamente por meio de uma variável de ambiente predefinida AIP_TENSORBOARD_LOG_DIR.

Geralmente, isso pode ser feito ao fornecer os.environ['AIP_TENSORBOARD_LOG_DIR'] como o diretório de registros para as APIs de gravação de registros de TensorBoard de código aberto. O local do AIP_TENSORBOARD_LOG_DIR é normalmente definido com a variável staging_bucket.

Para configurar seu script de treinamento no TensorFlow 2.x, crie um callback do TensorBoard e defina oslog_dir comoos.environ['AIP_TENSORBOARD_LOG_DIR'] O callback do TensorBoard é incluído nomodel.fit de callbacks.

  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],
  )
  

Saiba mais sobre como a Vertex AI define variáveis de ambiente no ambiente de treinamento personalizado.

Criar e executar um pipeline

O exemplo a seguir mostra como criar e executar um pipeline usando o pacote DSL do Kubeflow Pipelines. Para mais exemplos e detalhes, consulte a documentação do Vertex AI Pipelines.

Criar um componente de treinamento

Empacote seu código de treinamento em um componente personalizado, certificando-se de que o código esteja configurado para gravar registros do TensorBoard em um bucket do Cloud Storage. Para mais exemplos, consulte Criar seus próprios componentes de pipeline.

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],
    )

Crie e compile um pipeline.

Crie um job de treinamento personalizado a partir do componente criado. Basta especificar a especificação do componente em create_custom_training_job_op_from_component. Defina tensorboard_resource_name como sua instância do TensorBoard e staging_bucket como o local para preparar os artefatos durante chamadas de API (incluindo registros do TensorBoard).

Em seguida, crie um pipeline para incluir esse job e compile-o em um arquivo JSON.

Para mais exemplos e informações, consulte Componentes de job personalizados e Criar um pipeline.

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

Enviar um pipeline da Vertex AI

enviar seu pipeline usando o SDK da Vertex AI para Python. Para mais informações, consulte Executar um pipeline.

from typing import Any, Dict, Optional

from google.cloud import aiplatform


def log_pipeline_job_to_experiment_sample(
    experiment_name: str,
    pipeline_job_display_name: str,
    template_path: str,
    pipeline_root: str,
    project: str,
    location: str,
    parameter_values: Optional[Dict[str, Any]] = None,
):
    aiplatform.init(project=project, location=location)

    pipeline_job = aiplatform.PipelineJob(
        display_name=pipeline_job_display_name,
        template_path=template_path,
        pipeline_root=pipeline_root,
        parameter_values=parameter_values,
    )

    pipeline_job.submit(experiment=experiment_name)

A seguir

• Veja os resultados: Veja o TensorBoard para o Vertex AI Pipelines.
• Saiba como otimizar o desempenho dos jobs de treinamento personalizado usando Cloud Profiler.