TensorBoard hangs when visualizing an Object Detection graph

Viewed 886

I need to visualize the structure of a TensorFlow Object Detection model. I am trying to use TensorBoard in Colab with the code below. When TensorBoard loads the logs, it gets stuck on the step "Namespace hierarchy: Finding similar subgraphs".

!pip install -U tensorflow

import tensorflow as tf

from tensorflow.python.client import session  
from tensorflow.python.framework import ops  
from tensorflow.python.tools import saved_model_utils
from tensorflow.python.framework import importer
from tensorflow.python.summary import summary

!wget http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_resnet50_v1_fpn_640x640_coco17_tpu-8.tar.gz
!tar -xf ssd_resnet50_v1_fpn_640x640_coco17_tpu-8.tar.gz

%load_ext tensorboard

log_dir = '/content/logs'
tag_set = "serve"
model_dir = '/content/ssd_resnet50_v1_fpn_640x640_coco17_tpu-8/saved_model'

with session.Session(graph=ops.Graph()) as sess:
  input_graph_def = saved_model_utils.get_meta_graph_def(model_dir,
                                                        tag_set).graph_def
  importer.import_graph_def(input_graph_def)

  pb_visual_writer = summary.FileWriter(log_dir)
  pb_visual_writer.add_graph(sess.graph)
  print("Model Imported. Visualize by running: "
        "tensorboard --logdir={}".format(log_dir))

%tensorboard --logdir=$log_dir

Here is a link to the notebook: https://colab.research.google.com/drive/1MrbNJYR2ds8RRgIBvgUgAILw0Jwdygui?usp=sharing.

Environment: Browser: Chrome OS: Windows RAM: 8 GB

Eventually, I start getting the errors below.

FYI, I tried to run this same process on a Windows computer with 4 GB of RAM, with a TensorBoard server running in a shell. I accessed TensorBoard using the default URL (outside of a notebook). It failed at the same point in the startup process.

I see that Tensorboard got stuck in 'namespace hierarchy finding similar subgraphs' and I'm struggling to implement tensorboard monitoring into the Mask_RCNN training process asked similar questions, but no answer has been provided.

Thanks in advance -- this would really help in the important project I am doing for my company.

enter image description here

Pages Unresponsive

Could not load JavaScript

1 Answers

There are 3 issues that I believe I now know how to fix.

  1. Insufficient RAM/CPU -> run on a computer with more resources
  2. Complex graph, needs to be enlarged -> Use Zoom to see nodes
  3. Variables not initialized -> Evaluate model to show edges of graph

1. Insufficient RAM/CPU

Running the process on a Windows computer with 24 GB RAM and 8 i7 cores eliminated the crashes for me. Sometimes I get a message about the page becoming unresponsive, but you can just click the "Wait" button and it will load. The load itself was also much faster.

2. Complex graph, needs to be enlarged

Running the above code, TensorBoard displays what appears to be a blank screen. However, there is a faint outline of graph nodes.

Initial view of graph

Zoom in to see the details. The navigation box in the lower right portion of the screen is helpful.

Zooming in

3. Initialize variables

Instead of the code in my question above, use the following:

import tensorflow as tf
import numpy as np

!pip install -U tensorflow

%load_ext tensorboard

# Download model
!wget http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_resnet50_v1_fpn_640x640_coco17_tpu-8.tar.gz
!tar -xf ssd_resnet50_v1_fpn_640x640_coco17_tpu-8.tar.gz

model_dir = '/content/ssd_resnet50_v1_fpn_640x640_coco17_tpu-8/saved_model'
log_dir = '/content/logs'

if os.path.exists(log_dir):
  ! rm -r $log_dir

@tf.function
def f(x):
    imported = tf.saved_model.load(model_dir)
    results = imported(x)
    return results
    
#Initialize variables
imgs = np.zeros((1,640,640,3),dtype=int)
imgs_t = tf.constant(imgs, dtype=tf.dtypes.uint8)
imported_g = f.get_concrete_function(imgs_t).graph

# Export the graph
with session.Session(graph=imported_g) as sess:
  pb_visual_writer = summary.FileWriter(log_dir)
  pb_visual_writer.add_graph(sess.graph)
  print("Model Imported. Visualize by running: "
        "tensorboard --logdir={}".format(log_dir))

The new graph shows the edges of the graph, not just the nodes.

Here's a notebook that you can run: https://colab.research.google.com/gist/mherzog01/d631998cb4d0b0dbcb70492b933a67c8/tensorboard-hangs-during-graph-visualization-solution.ipynb.

Related