I'm trying to show the graph of a deep autoencoder in TensorBoard.
I've called the decorator @tf.function on the function used for training, and I've successfully used model.summary() to print the summary, so the graph should exist.
In order to show the graph I've called summary.trace_on() during the training loop.
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) I'd expect the graph to show in TensorBoard, but that doesn't seem to be the case. What am I doing wrong?
