TensorFlow 2.0 - 图形未显示在TensorBoard中

Question

我试图在TensorBoard中显示深度自动编码器的图形。

我在用于训练的函数上调用了装饰器@tf.function，并且我已成功使用model.summary()来打印摘要，因此图形应该存在。

为了显示图表，我在训练循环期间调用了summary.trace_on()。

import os
import shutil

import numpy as np
import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Layer

from utils import log_results, mse, preprocess_mnist, train_autoencoder


class DeepEncoder(Layer):
    def __init__(self, intermediate_dims):
        super(DeepEncoder, self).__init__()
        self.layers = [
            Dense(units=i, activation=tf.nn.relu) for i in intermediate_dims
        ]

    def call(self, X):
        Z = X
        for layer in self.layers:
            Z = layer(Z)
        return Z


class DeepDecoder(Layer):
    def __init__(self, original_dim, intermediate_dims):
        super(DeepDecoder, self).__init__()
        self.layers = [
            Dense(units=i, activation=tf.nn.relu)
            for i in reversed(intermediate_dims[:-1])
        ]
        self.layers.append(Dense(units=original_dim))

    def call(self, X):
        Z = X
        for layer in self.layers:
            Z = layer(Z)
        return Z


class DeepAutoEncoder(Model):
    def __init__(self, original_dim, intermediate_dims):
        super(DeepAutoEncoder, self).__init__()
        self.encoder = DeepEncoder(intermediate_dims=intermediate_dims)
        self.decoder = DeepDecoder(
            original_dim=original_dim, intermediate_dims=intermediate_dims)

    def call(self, X):
        return self.decoder(self.encoder(X))

    def encode(self, X):
        return self.encoder(X)

    def decode(self, Z):
        return self.decode(Z)


def test_deep_autoencoder(batch_size,
                          learning_rate,
                          epochs,
                          max_outputs=4,
                          seed=None):

    tf.random.set_seed(seed)

    X_train, X_test, train_dataset, _, _, _ = preprocess_mnist(
        batch_size=batch_size)

    autoencoder = DeepAutoEncoder(
        original_dim=784, intermediate_dims=[1024, 256, 64])
    opt = tf.optimizers.Adam(learning_rate=learning_rate)

    log_path = 'logs/deepautoencoder'
    if os.path.exists(log_path):
        shutil.rmtree(log_path)

    writer = tf.summary.create_file_writer(log_path)

    with writer.as_default():
        with tf.summary.record_if(True):
            for epoch in range(epochs):
                for step, batch in enumerate(train_dataset):
                    train_autoencoder(mse, autoencoder, opt, batch)

                # logs (train)
                train_loss = log_results(
                    model=autoencoder,
                    X=X_train,
                    max_outputs=max_outputs,
                    epoch=epoch,
                    prefix='train')

                # logs (test)
                test_loss = log_results(
                    model=autoencoder,
                    X=X_test,
                    max_outputs=max_outputs,
                    epoch=epoch,
                    prefix='test')

                tf.summary.trace_on()

                writer.flush()

                template = 'Epoch {}, Train loss: {:.5f}, Test loss: {:.5f}'
                print(
                    template.format(epoch + 1, train_loss.numpy(),
                                    test_loss.numpy()))

    if not os.path.exists('saved_models'):
        os.makedirs('saved_models')
    np.savez_compressed('saved_models/deepencoder.npz',
                        *autoencoder.encoder.get_weights())


if __name__ == '__main__':
    test_deep_autoencoder(batch_size=128, learning_rate=1e-3, epochs=20)

我希望图表显示在TensorBoard中，但似乎并非如此。我究竟做错了什么？

Answer 1

您需要在跟踪执行后调用tf.summary.trace_export，以便将跟踪实际写入磁盘。

请参阅trace_on和trace_export的文档