Configurer le script d'entraînement

Votre script d'entraînement doit être configuré pour écrire des journaux TensorBoard. Pour les utilisateurs TensorBoard existants, il n'est pas nécessaire de modifier le code d'entraînement du modèle.

Pour configurer votre script d'entraînement dans TensorFlow 2.x, créez un rappel TensorBoard et définissez la variable log_dir sur n'importe quel emplacement pouvant se connecter à Google Cloud.

Le rappel TensorBoard est ensuite inclus dans la liste des rappels model.fit TensorFlow.

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

Les journaux TensorBoard sont créés dans le répertoire spécifié et peuvent être importés dans un test Vertex AI TensorBoard en suivant les instructions d'importation de la page Importer des journaux TensorBoard.

Pour plus d'exemples, consultez la documentation Open Source de TensorBoard.

Étapes suivantes