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.

SDK da Vertex AI para Python

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)
  • experiment_name: dê um nome para o experimento.
  • pipeline_job_display_name: o nome de exibição do job do pipeline.
  • template_path: o caminho para o modelo de pipeline compilado.
  • pipeline_root: especifique um URI do Cloud Storage que sua conta de serviço de pipelines possa acessar. Os artefatos das execuções de pipeline são armazenados na raiz do pipeline.
  • parameter_values: (opcional) os parâmetros de pipeline a serem transmitidos para essa execução. Por exemplo, crie um dict() com os nomes do parâmetro como as chaves do dicionário e os valores de parâmetro como os valores do dicionário.
  • project: o ID do projeto. O projeto do Google Cloud em que o pipeline será executado. Esses IDs estão na página de boas-vindas do console do Google Cloud.
  • location: o local para executar os componentes do pipeline. Deve ser igual local que a instância do TensorBoard usada.

A seguir