Visualize the graph of a custom model in Tensorflow 2.4

Question

I want to visualize my model graph in Tensorboard to check if I implemented my model correct.
I´m implementing my own model by subclassing the tf.keras.Layer and tf.keras.Model class which looks like this (Kept everything unnecessary away):

class My_Model(tf.keras.Model):
    def __init__(self):
        super(Model_C_1, self).__init__()
        # actually here is much more, but this is not important.

    def build(self, inputs_shape):
        self.conv1 = tf.keras.layers.Conv1D(filters=16)
        # actually here is much more, but this is not important. 

    @tf.function
    def call(self, input, training):
        x = self.conv1(input)
        # actually here is much more, but this is not important. 
        return x

I want to visualize the computation graph, since the model is way more complicated and I´m not sure, if I missed something (Model trains and work, but I want to double check). My training loop (very simplyfied) looks like this:

def train_step(batch, model, params, writer, optimizer):

    data = batch['data']
    with tf.GradientTape() as tape:
        predictions = model(data, training=True)
        loss = loss_object(labels, predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    tf.summary.scalar(name='loss', data=loss, step=optimizer.iterations) 


# I process my data with pandas/numpy first and then
kf = KFold(n_splits=params.n_k_fold_splits)

for split_nr, (train_index, val_index) in enumerate(kf.split(ds)):
    #...
    writer = tf.summary.create_file_writer(params.path_train_log)
    for batch_train in train_ds:
        train_step(batch_train, model, params, writer, optimizer)

Training and metrics visualization in Tensorboard works fine. However, I do not get a graph of my Model in Tensorboard. I tried using the following at the very end of my complete training.

 with writer.as_default():
            tf.summary.trace_export(
              name="My_Trace",
              step=0)

This results in the error: ValueError: Must enable trace before export. I could not figure out where to add the tf.summary.trace_on(graph=True) command mentioned by TF.
Any suggestions?

user_9401215 · Answer

The main problem I've found to getting it to work is that you have to make sure that you are doing the trace before you have ever built the graph. So it has to be run on the first batch in a training example, the first time you are calling the model.

In my example it was, where train_step is my @tf.function

for batch, (input_image, target) in tqdm(train_ds.enumerate(), total=num_steps):
  if batch == 0:
    tf.summary.trace_on(graph=True)
    train_step(input_image, target, epoch)
    with summary_writer.as_default():
      tf.summary.trace_export(name='train_func', step=0)
  else:
    train_step(input_image, target, epoch)

If you don't do it before it's ever been called it seems to work, makes a big tf.events file but doesn't actually do it.

Another options is to call tf.summary.graph before you first call the model.

Visualize the graph of a custom model in Tensorflow 2.4

1 Answers