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

Question

我正在尝试在 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 ) 上的 TensorBoard 中找到有关自定义添加和跟踪指标的新工作流程的任何部分。我的人为实现基于 tf.summary ( https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/summary ) 的 API 文档

Answer 1

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