Adding stats (run metadata) to the Graph UI.

Todos:
- Add actual memory and cpu data to the info card.
Change: 119050200

tensorflow/tensorboard/components/tf-graph-app/tf-graph-app.html

@@ -45,23 +45,21 @@
   <div class="side">
     <tf-graph-controls
         color-by-params="[[colorByParams]]"
-        has-stats="[[hasStats]]"
+        stats="[[stats]]"
         color-by="{{colorBy}}"
     ></tf-graph-controls>
     <tf-graph-loader id="loader"
         out-graph-hierarchy="{{graphHierarchy}}"
         out-graph="{{graph}}"
-        out-graph-name="{{graphName}}"
-        has-stats="{{hasStats}}"
+        out-stats="{{stats}}"
         progress="{{_progress}}"
     ></tf-graph-loader>
   </div>
   <div class="main">
     <tf-graph-board id="graphboard"
         graph-hierarchy="[[graphHierarchy]]"
         graph="[[graph]]"
-        has-stats="[[hasStats]]"
-        graph-name="[[graphName]]"
+        stats="[[stats]]"
         progress="[[_progress]]"
         color-by="[[colorBy]]"
         color-by-params="{{colorByParams}}"
@@ -77,7 +75,7 @@
 Polymer({
   is: 'tf-graph-app',
   properties: {
-    hasStats: Boolean,
+    stats: Object,
     pbtxt: {
       type: String,
       observer: '_updateGraph',

tensorflow/tensorboard/components/tf-graph-board/tf-graph-board.html

@@ -117,11 +117,11 @@
               basic-graph="[[graph]]"
               hierarchy-params="[[hierarchyParams]]"
               render-hierarchy="{{_renderHierarchy}}"
+              stats="[[stats]]"
               selected-node="{{_selectedNode}}"
               highlighted-node="{{_highlightedNode}}"
               color-by="[[colorBy]]"
               color-by-params="{{colorByParams}}"
-              graph-name="[[graphName]]"
               progress="{{progress}}"
     ></tf-graph>
   </div>
@@ -150,9 +150,7 @@
     // Public API.
     graphHierarchy: Object,
     graph: Object,
-    graphName: String,
-    // True if the graph data has also run-time stats.
-    hasStats: Boolean,
+    stats: Object,
     /**
      * @type {value: number, msg: string}
      *

tensorflow/tensorboard/components/tf-graph-common/lib/common.ts

@@ -239,41 +239,39 @@ export interface TFNode {
 /**
  * TensorFlow stats file definition as defined in the stats proto file.
  */
-export interface TFStats {
-  devStats: {device: string, nodeStats: TFNodeStats[]}[];
+export interface StepStats {
+  dev_stats: {device: string, node_stats: NodeStats[]}[];
 }
 
 /**
  * TensorFlow stats for a node as defined in the stats proto file.
  */
-export interface TFNodeStats {
-  nodeName: string;
+export interface NodeStats {
+  node_name: string;
   // The next 4 properties are currently stored as string in json
   // and must be parsed.
-  allStartMicros: number;
-  opStartRelMicros: number;
-  opEndRelMicros: number;
-  allEndRelMicros: number;
+  all_start_micros: number;
+  op_start_rel_micros: number;
+  op_end_rel_micros: number;
+  all_end_rel_micros: number;
   memory: {
-    allocatorName: string;
-    totalBytes: number; // Stored as string in json and should be parsed.
-    peakBytes: number; // Stored as string in json and should be parsed.
+    allocator_name: string;
+    total_bytes: number; // Stored as string in json and should be parsed.
+    peak_bytes: number; // Stored as string in json and should be parsed.
   }[];
   /** Output sizes recorded for a single execution of a graph node */
   output: TFNodeOutput[];
-  timelineLabel: string;
-  scheduledMicros: string;
-  threadId: string;
+  timeline_label: string;
+  scheduled_micros: string;
+  thread_id: string;
 }
 
 /**
  * Description for the output tensor(s) of an operation in the graph.
  */
 export interface TFNodeOutput {
   slot: number; // Stored as string in json and should be parsed.
-  /** Was the tensor allocated by this Op or a previous computation */
-  allocationType: string;
-  tensorDescription: {
+  tensor_description: {
     /** Data type of tensor elements */
     dtype: string;
     /** Shape of the tensor */
@@ -292,15 +290,15 @@ export interface TFNodeOutput {
       }[];
     };
     /** Information about the size and allocator used for the data */
-    allocationDescription: {
+    allocation_description: {
       // The next 2 properties are stored as string in json and
       // should be parsed.
       /** Total number of bytes requested */
-      requestedBytes: number;
+      requested_bytes: number;
       /** Total number of bytes allocated, if known */
-      allocatedBytes?: number;
+      allocated_bytes?: number;
       /** Name of the allocator used */
-      allocatorName: string;
+      allocator_name: string;
     };
   };
 }

tensorflow/tensorboard/components/tf-graph-common/lib/graph.ts

@@ -376,35 +376,40 @@ export function createMetanode(name: string, opt = {}): Metanode {
  * graph information.
  */
 export function joinStatsInfoWithGraph(graph: SlimGraph,
-    statsJson: TFStats): void {
-  _.each(statsJson.devStats, stats => {
-    _.each(stats.nodeStats, nodeStats => {
+    stats: StepStats): void {
+  _.each(stats.dev_stats, devStats => {
+    _.each(devStats.node_stats, nodeStats => {
       // Lookup the node in the graph by its original name, e.g. A. If not
       // found, lookup by the rewritten name A/(A) in case the name is both
       // a namespace and a node name.
-      let nodeName = nodeStats.nodeName in graph.nodes ?
-          nodeStats.nodeName :
-          nodeStats.nodeName + NAMESPACE_DELIM + "(" + nodeStats.nodeName + ")";
-      if (nodeName in graph.nodes) {
-        // Compute the total bytes used.
-        let totalBytes = 0;
-        if (nodeStats.memory) {
-          _.each(nodeStats.memory, alloc => {
-            if (alloc.totalBytes) {
-              totalBytes += Number(alloc.totalBytes);
-            }
-          });
-        }
-        let outputSize: number[][] = null;
-        if (nodeStats.output) {
-          outputSize = _.map(nodeStats.output, output => {
-            return _.map(output.tensorDescription.shape.dim,
-                dim => Number(dim.size));
-          });
-        }
-        graph.nodes[nodeName].stats = new NodeStats(totalBytes,
-            Number(nodeStats.allEndRelMicros), outputSize);
+      let nodeName = nodeStats.node_name in graph.nodes ?
+          nodeStats.node_name :
+          nodeStats.node_name + NAMESPACE_DELIM + "(" + nodeStats.node_name + ")";
+
+      // Couldn't find a matching node.
+      if (!(nodeName in graph.nodes)) {
+        return;
+      }
+
+      // Compute the total bytes used.
+      let totalBytes = 0;
+      if (nodeStats.memory) {
+        _.each(nodeStats.memory, alloc => {
+          if (alloc.total_bytes) {
+            totalBytes += Number(alloc.total_bytes);
+          }
+        });
+      }
+      let outputSize: number[][] = null;
+      if (nodeStats.output) {
+        outputSize = _.map(nodeStats.output, output => {
+          return _.map(output.tensor_description.shape.dim,
+              dim => Number(dim.size));
+        });
       }
+      graph.nodes[nodeName].device = devStats.device;
+      graph.nodes[nodeName].stats = new NodeStats(totalBytes,
+          Number(nodeStats.all_end_rel_micros), outputSize);
     });
   });
 }
@@ -492,7 +497,6 @@ class MetanodeImpl implements Metanode {
     this.templateId = null;
     /** Metanode which contains this node, if any */
     this.parentNode = null;
-    this.stats = new NodeStats(0, 0, null);
     this.hasNonControlEdges = false;
     this.include = InclusionType.UNSPECIFIED;
   }
@@ -705,7 +709,6 @@ class SeriesNodeImpl implements SeriesNode {
     this.parentNode = null;
     this.deviceHistogram = {};
     this.hasNonControlEdges = false;
-    this.stats = new NodeStats(0, 0, null);
     this.include = InclusionType.UNSPECIFIED;
   }
 }

tensorflow/tensorboard/components/tf-graph-common/lib/hierarchy.ts

@@ -428,6 +428,42 @@ export function build(graph: tf.graph.SlimGraph, params: HierarchyParams,
   });
 };
 
+export function joinAndAggregateStats(h: Hierarchy, stats: StepStats) {
+  // Get all the possible device names.
+  let deviceNames = {};
+  _.each(h.root.leaves(), nodeName => {
+    let leaf = <OpNode> h.node(nodeName);
+    if (leaf.device != null) {
+      deviceNames[leaf.device] = true;
+    }
+  });
+  h.devices = _.keys(deviceNames);
+
+  // Reset stats for each group node.
+  _.each(h.getNodeMap(), (node, nodeName) => {
+    if (node.isGroupNode) {
+      node.stats = new NodeStats(0, 0, null);
+      (<GroupNode>node).deviceHistogram = {};
+    }
+  });
+
+  // Bubble-up the stats and device distribution from leaves to parents.
+  _.each(h.root.leaves(), nodeName => {
+    let leaf = <OpNode> h.node(nodeName);
+    let node = <GroupNode|OpNode> leaf;
+    while (node.parentNode != null) {
+      if (leaf.device != null) {
+        let deviceHistogram = (<GroupNode>node.parentNode).deviceHistogram;
+        deviceHistogram[leaf.device] = (deviceHistogram[leaf.device] || 0) + 1;
+      }
+      if (leaf.stats != null) {
+        node.parentNode.stats.combine(leaf.stats);
+      }
+      node = <GroupNode> node.parentNode;
+    }
+  });
+}
+
 /**
  * Creates the metanodes in the hierarchical graph and assigns parent-child
  * relationship between them.
@@ -446,9 +482,6 @@ function addNodes(h: Hierarchy, graph: SlimGraph) {
       parent.depth = Math.max(parent.depth, path.length - i);
       parent.cardinality += node.cardinality;
       parent.opHistogram[node.op] = (parent.opHistogram[node.op] || 0) + 1;
-      if (node.stats) {
-        parent.stats.combine(node.stats);
-      }
       if (node.device != null) {
         parent.deviceHistogram[node.device] =
             (parent.deviceHistogram[node.device] || 0) + 1;
@@ -623,11 +656,6 @@ function groupSeries(metanode: Metanode, hierarchy: Hierarchy,
       }
       child.parentNode = seriesNode;
       seriesNames[n] = seriesName;
-
-      if (child.stats) {
-        seriesNode.stats.combine(child.stats);
-      }
-
       // Remove now-grouped node from its original parent's metagraph.
       metagraph.removeNode(n);
     });