How to visualize the network graph in tensorflow 1.15 with Eager mode using tensorboard?

[Frank-Dz]

Hi~ I want to visualize the NN in Eager mode in tf1.15 (can not switch to 2.0.0). And the implementation is based on low-level API of Tensorflow 1.15. I want to use the tensorboard to visualize it. I write a log tracing code but get the error:

WARNING:tensorflow:
The TensorFlow contrib module will not be included in TensorFlow 2.0.
For more information, please see:
  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md
  * https://github.com/tensorflow/addons
  * https://github.com/tensorflow/io (for I/O related ops)
If you depend on functionality not listed there, please file an issue.

Traceback (most recent call last):
  File "/home/frank/PycharmProjects/reconstruction_NN/my_test.py", line 78, in <module>
    tf.contrib.summary.trace_on(graph=True, profiler=True)
AttributeError: module 'tensorflow.contrib.summary.summary' has no attribute 'trace_on'

Environment information (required)

tensorboard 1.15.0
tensorflow-estimator 1.15.1
tensorflow-gpu 1.15.0
Ubuntu16.04

Issue description

Code:

import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten, Conv2D,Dropout
from tensorflow.keras import Model
tf.compat.v1.enable_eager_execution()
print(tf.__version__)
print(tf.executing_eagerly())

mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train, x_test = x_train / 255.0, x_test / 255.0

batch_size = 32

train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(batch_size)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size)

class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.flatten = Flatten()
        self.d1 = Dense(128, activation='relu')
        self.dropout = Dropout(0.5)
        self.d2 = Dense(10, activation='softmax')

    def call(self, x):
        x = self.flatten(x)
        x = self.d1(x)
        x = self.dropout(x)
        return self.d2(x)

model = MyModel()

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, predictions)

@tf.function
def test_step(images, labels):
    predictions = model(images)
    t_loss = loss_object(labels, predictions)
    test_loss(t_loss)
    test_accuracy(labels, predictions)

EPOCHS = 5
from datetime import *
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = 'logs/func/%s' % stamp
writer = tf.contrib.summary.create_file_writer(logdir)
tf.summary.trace_on(graph=True, profiler=True)

for epoch in range(EPOCHS):

    for images, labels in train_ds:

        train_step(images, labels)

    for test_images, test_labels in test_ds:
        test_step(test_images, test_labels)

    template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}'

    print(template.format(epoch + 1, train_loss.result(),
                              train_accuracy.result() * 100,
                              test_loss.result(),
                              test_accuracy.result() * 100))
with writer.as_default():
  tf.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)

Also see: https://stackoverflow.com/questions/61724827/how-to-visualize-the-network-graph-in-tensorflow-1-15-with-eager-mode-using-tens

Thanks in advance!

[pindinagesh]

@Frank-Dz

tf.contrib is deprecated instead of that you can use writer object as writer = tf.summary.create_file_writer(logdir) and also please refer this documentation, Please feel free to reopen this issue if this still exist. Thanks!