Criar perfil de desempenho de treinamento de modelo usando o Vertex AI TensorBoard Profiler

Nesta página, mostramos como ativar o Vertex AI TensorBoard Profiler para depurar o desempenho do treinamento de modelo para seus jobs de treinamento personalizados.

Os modelos de treinamento podem ser computacionalmente caros. O TensorBoard Profiler do Vertex AI permite monitorar e otimizar o desempenho do treinamento de modelo, ajudando a entender o consumo de recursos das operações de treinamento. Com essas informações, é possível identificar e corrigir gargalos de desempenho para treinar modelos de maneira mais rápida e mais barata.

Antes de começar

Ativar o criador de perfil do TensorBoard da Vertex AI

Para ativar o Vertex AI TensorBoard Profiler para o job de treinamento, adicione o seguinte ao seu script de treinamento:

  1. Adicione a importação cloud_profiler nas importações de nível superior:

    from google.cloud.aiplatform.training_utils import cloud_profiler

  2. Inicialize o plug-in cloud_profiler adicionando:

    cloud_profiler.init()

Exemplo

Veja um exemplo de script de treinamento:

#!/usr/bin/env python

import tensorflow as tf
import argparse
import os
from google.cloud.aiplatform.training_utils import cloud_profiler
import time

"""Train an mnist model and use cloud_profiler for profiling."""

def _create_model():
    model = tf.keras.models.Sequential(
        [
            tf.keras.layers.Flatten(input_shape=(28, 28)),
            tf.keras.layers.Dense(128, activation="relu"),
            tf.keras.layers.Dropout(0.2),
            tf.keras.layers.Dense(10),
        ]
    )
    return model

def main(args):
    strategy = None
    if args.distributed:
        strategy = tf.distribute.MultiWorkerMirroredStrategy()

    mnist = tf.keras.datasets.mnist

    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0

    if args.distributed:
        strategy = tf.distribute.MultiWorkerMirroredStrategy()
        with strategy.scope():
            model = _create_model()
            model.compile(
                optimizer="adam",
                loss=tf.keras.losses.sparse_categorical_crossentropy,
                metrics=["accuracy"],
            )
    else:
        model = _create_model()
        model.compile(
            optimizer="adam",
            loss=tf.keras.losses.sparse_categorical_crossentropy,
            metrics=["accuracy"],
        )

    # Initialize the profiler.
    cloud_profiler.init()

    # Use AIP_TENSORBOARD_LOG_DIR to update where logs are written to.
    tensorboard_callback = tf.keras.callbacks.TensorBoard(
        log_dir=os.environ["AIP_TENSORBOARD_LOG_DIR"], histogram_freq=1
    )

    model.fit(
        x_train,
        y_train,
        epochs=args.epochs,
        verbose=0,
        callbacks=[tensorboard_callback],
    )
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--epochs", type=int, default=100, help="Number of epochs to run model."
    )
    parser.add_argument(
        "--distributed", action="store_true", help="Use MultiWorkerMirroredStrategy"
    )
    args = parser.parse_args()
    main(args)

Acesse o painel do criador de perfil do TensorBoard da Vertex AI

Depois de configurar o script de treinamento para ativar o Vertex AI TensorBoard Profiler, execute o script de treinamento com uma instância do Vertex AI TensorBoard.

No script de treinamento, verifique as seguintes configurações:

  • Defina BASE_OUTPUT_DIR: como o bucket do Cloud Storage em que você quer armazenar os registros do TensorBoard da Vertex AI, gerados pelo script de treinamento.
  • Defina 'serviceAccount': como a conta de serviço que você criou com os papéis roles/storage.admin e roles/aiplatform.user.

  • Defina 'tensorboard': como o nome totalmente qualificado da instância do TensorBoard da Vertex AI que você quer usar com este job de treinamento. O nome totalmente qualificado tem o seguinte formato:

    projects/PROJECT_NUMBER_OR_ID/locations/REGION/tensorboards/TENSORBOARD_INSTANCE_ID

Há duas maneiras de acessar o painel do TensorBoard Profiler da Vertex AI no Console do Google Cloud:

  • Na página "Jobs personalizados".
  • Na página "Experimentos".

Acessar o painel do criador de perfil na página "Jobs personalizados"

Use esse método para acessar o painel do TensorBoard Profiler da Vertex AI, mesmo que o job de treinamento esteja no estado Concluído.

  1. No Console do Google Cloud, acesse a guia Jobs personalizados na página Treinamento.

    Acessar Jobs personalizados

  2. Clique no nome do job de treinamento que você acabou de criar para acessar a página de detalhes.

  3. Clique em Abrir o TensorBoard.

  4. Clique na guia Perfil.

Acessar o painel do criador de perfil na página "Experimentos"

Use esse método para acessar o painel do TensorBoard Profiler da Vertex AI somente quando o job de treinamento estiver no estado Em execução.

  1. No Console do Google Cloud, acesse a página Experimentos da Vertex AI.

    Acessar os experimentos da Vertex AI

  2. Selecione a região do job de treinamento que você acabou de criar.

  3. Clique em Abrir TensorBoard ao lado do nome do job de treinamento.

  4. Clique na guia Perfil.

Capturar uma sessão de criação de perfil

Para capturar uma sessão de criação de perfil, seu job de treinamento precisa estar no estado Em execução. Na guia Profile da instância do TensorBoard da Vertex AI, execute as seguintes etapas:

  1. Clique em Capturar perfil.

  2. No campo URLs de serviço do perfil ou nome da TPU, digite:

    workerpool0-0

  3. Em Tipo de endereço, selecione Endereço IP.

  4. Clique em Capturar.

Notebook

A seguir