Configurar el script de entrenamiento

Tu secuencia de comandos de entrenamiento debe estar configurada para escribir registros de TensorBoard. Si ya usas TensorBoard, no tendrás que cambiar el código de entrenamiento de tu modelo.

Para configurar tu script de entrenamiento en TensorFlow 2.x, crea una retrollamada de TensorBoard y asigna a la variable log_dir cualquier ubicación que pueda conectarse a Google Cloud.

La retrollamada de TensorBoard se incluye en la lista de retrollamadas de TensorFlow model.fit.

import tensorflow as tf

def train_tensorflow_model_with_tensorboard(log_dir):
    (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=log_dir,
        histogram_freq=1
    )

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

Los registros de TensorBoard se crean en el directorio especificado y se pueden subir a un experimento de Vertex AI TensorBoard siguiendo las instrucciones de Subir registros de TensorBoard.

Para ver más ejemplos, consulta la documentación de código abierto de TensorBoard.

Siguientes pasos