使用TensorFlow 2.0Alpha时，无法在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

在graphs选项卡中，我看到

​

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

​

​

The guide article不通过tensorboard跟踪指标，我也没有在TensorBoard on alpha (https://www.tensorflow.org/alpha)中找到用于自定义添加和跟踪指标的新工作流的任何部分。我设计的实现基于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模型图。