Visualizing Data Flow Graphs in TensorFlow

How an interactive graph visualization made all the difference for deep learning developers.

Photo by Carlos Muza on Unsplash

The last few years have seen tremendous growth in the field of machine learning mainly due to the groundbreaking technique known as deep learning: so-called “neural networks” capable of unsupervised learning. However, as a result, deep learning models have a very complex network structure. Creating, understanding, and debugging such networks is a very complicated task.

To make this process easier, developers often draw diagrams first to help them understand and share the high-level structured modules. Since diagrams are an integral part of the process, developers often wish for a tool that could help them automatically generate these diagrams.

Simplified Diagrams to help build a mental map of the model.

TensorFlow: A Blessing for Developers?

To make the lives of developers easier, Google released TensorFlow, an Open Source Library that makes it easier to deploy deep learning models. The main objective of TensorFlow is to use data flow graphs to represent computation.

However, while TensorFlow makes it easier to create deep learning models, it does not really help them understand its structure. Thus, Kanit Wongsuphasawat et al. introduced TensorFlow Graph Visualizer as part of the TensorFlow library. The main function of it is to convert the low-level data flow diagram into a high-level interactive diagram.

Converting a graph from uninterpretable to highly interactive.

Critical Tasks

A set of five main tasks were identified that the visualization should be able to perform after having elaborate conversations with experts:

1. Visualization should be able to understand the high-level components of the model and their relationships.
2. Recognizing similarities and differences between components in the graph. This would be very helpful in building a mental model as well as detecting bugs.
3. Ability to examine the nested structure of complex high-level components especially when the structure has been auto-generated via an API call.
4. Ability to inspect individual operations such as inputs, outputs, etc so that the developer doesn’t have to refer to code every now and then.
5. Represent quantitative data of the graph such as tensor size, computation time, etc. efficiently.

Transformation Strategy

When converting a low-level data flow diagram into a high-level interactive diagram, TensorFlow Graph Visualizer undertakes the following steps:

1. Extract Non-critical operations: This helps in decluttering the graph. Some operations like initializing a variable, summary operations are common but they don’t help developers in distinguishing modules. Thus we reduce the size of their nodes. This is possible only because these operations are connected to only one other operation thus extracting them doesn’t alter the path between other nodes.

Extraction of Nodes ‘Zeros’ & ‘Histogram Summary’.

2) Building Hierarchical Clustered graph: In this step, nodes are grouped. Users can specify the names of the nodes. Based on these names nodes can be grouped to declutter the graph even further.

Grouping of nodes under ‘Train’ & ’Weights’ Nodes.

3) Extract Auxiliary Nodes: Extracting the nodes at the start and end of the network to the right of the main graph helps greatly in decluttering. These groups are shown using placeholder icons in the main graph. In this way, the user can keep track of these groups in an efficient manner.

Main Graph shows the flow of the Network. While the Auxiliary map consists of extracted nodes.

Usage Patterns & Feedback

Feedback was mainly collected in three ways: (1) a questionnaire for internal users at Google, (2) observing conversations in the TensorFlow mailing list, and (3) collecting public feedback on online forums.

On observing the usage of TensorFlow Graph Visualizer on various platforms certain usage patterns were identified. The key usage patterns that were observed are

1. Users were exploring new models as they were excited to verify that they were indeed getting what they intended to get.
2. Users were actively sharing the screenshots from TensorFlow to explain models. This was also visible on many online forums like Stack Overflow.
3. Users had no problems with renaming the nodes if it improved their visualizations. There were also many blog posts and video tutorials guiding users on how to rename nodes.

Overall users were very happy with the following quote from Quora even saying that it is a step in the right direction.

We believe visualization is really fundamental to the creative process and our ability to develop better models. So, visualization tools like TensorBoard are a great step in the right direction.

Video walkthrough of the working of the TensorFlow Graph Visualizer

Conclusion

Targeting the pain points of the developers and providing an appropriate solution is what made TensorFlow Graph Visualizer a success. Identifying the main functions first and then building on it helped the authors in creating a groundbreaking solution for the problem of creating a data flow diagram for deep learning models.

This blog post is based on the following paper:

• Kanit Wongsuphasawat, Daniel Smilkov, James Wexler, Jimbo Wilson, Dandelion Mané, Doug Fritz, Dilip Krishnan, Fernanda B. Viégas, and Martin Wattenberg. Visualizing Dataflow Graphs of Deep Learning Models in TensorFlow. IEEE Transactions on Visualization and Computer Graphics, 2018. [PDF]