使用 TensorFlow 2.0 Alpha 时无法在 Tensorboard 中看到 keras 模型图

Posted 2023-02-16

【中文标题】使用 TensorFlow 2.0 Alpha 时无法在 Tensorboard 中看到 keras 模型图

【英文标题】：Unable to see keras model graph in Tensorboard when using TensorFlow 2.0 Alpha

【发布时间】：2019-08-04 02:33:06

【问题描述】：

我正在尝试在 TensorFlow 2.0 alpha 上进行自定义训练，同时我正在尝试将一些指标和我的训练图添加到 TensorBoard。考虑以下人为设计的示例

import tensorflow as tf
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.models import Model


def create_model():
    inp = Input((32, ))
    net = Dense(16, activation="relu")(inp)
    net = Dense(8, activation="relu")(net)
    net = Dense(2, activation=None)(net)
    return Model(inp, net)


@tf.function
def grad(model, loss, x, y):
    with tf.GradientTape() as tape:
        y_ = model(x)
        loss_value = loss(y_true=y, y_pred=y_)
    return loss_value, tape.gradient(loss_value, model.trainable_variables)


@tf.function
def train_step(model, loss, optimizer, features, labels):
    loss_value, grads = grad(model, loss, features, labels)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss_value


def train():
    tf.summary.trace_on(graph=True, profiler=True)

    with tf.summary.create_file_writer("model").as_default():
        model = create_model()

        loss = tf.keras.losses.MeanSquaredError()
        optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)

        for i in range(10):
            tf.summary.experimental.set_step(i)

            features = tf.random.normal((16, 32))
            labels = tf.random.normal((16, 2))
            loss_value = train_step(model, loss, optimizer, features, labels)
            print(loss_value)

        tf.summary.trace_export("model", profiler_outdir="model")


if __name__ == "__main__":
    train()

这个，没有正确显示模型图，在做

tensorboard --logdir model

在图表标签中我看到

当我通过 model.fit 或 estimator 进行训练时，我得到了图表。例如，这是我使用model_to_estimator 转换模型时的图表部分

The guide article 不通过 tensorboard 跟踪指标，而且我没有找到任何有关在 alpha 版 (https://www.tensorflow.org/alpha) 上自定义添加和跟踪指标的新工作流程的部分。我设计的实现是基于 tf.summary (https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/summary) 的 API 文档

【参考方案1】：

您可以使用tf.keras.utils.plot_model(model, show_shapes=True,show_dtype=True,rankdir="LR") 绘制 keras 模型图。