diff options
Diffstat (limited to 'tensorflow/tools/tfprof/internal/tfprof_op.cc')
| -rw-r--r-- | tensorflow/tools/tfprof/internal/tfprof_op.cc | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/tensorflow/tools/tfprof/internal/tfprof_op.cc b/tensorflow/tools/tfprof/internal/tfprof_op.cc index 6a7077c085..77a2593623 100644 --- a/tensorflow/tools/tfprof/internal/tfprof_op.cc +++ b/tensorflow/tools/tfprof/internal/tfprof_op.cc @@ -126,6 +126,7 @@ const ShowMultiNode* TFOp::ShowInternal(const Options& opts, } nodes = SortNodes(nodes, opts); + // pre keeps track of previous visited node. OpNode* pre = nullptr; std::vector<OpNode*> account_nodes; for (auto it = nodes.rbegin(); it != nodes.rend(); ++it) { @@ -170,16 +171,20 @@ const ShowMultiNode* TFOp::ShowInternal(const Options& opts, root_->ResetTotalStats(); if (pre) { root_->AggregateTotalStats(pre); - root_->mutable_proto()->add_children()->MergeFrom(pre->proto()); - pre->mutable_proto()->clear_children(); } } + if (pre) { + root_->mutable_proto()->add_children()->MergeFrom(pre->proto()); + pre->mutable_proto()->clear_children(); + } if (opts.output_type == kOutput[1] || opts.output_type == kOutput[2]) { string display_str = FormatLegend(opts); for (OpNode* node : show_nodes) { display_str += FormatNode(node, root_.get(), opts); } + // In op view, we don't show root (total). But it will still in proto. + // TODO(xpan): Is it the right choice? root_->formatted_str = display_str; } return root_.get(); @@ -201,7 +206,7 @@ int64 TFOp::SearchRoot(const std::vector<OpNode*> nodes, return i; } -string TFOp::FormatNode(OpNode* node, OpNode* root, const Options& opts) { +string TFOp::FormatNode(OpNode* node, OpNode* root, const Options& opts) const { std::vector<string> attrs; if (opts.select.find(kShown[0]) != opts.select.end()) { |