Look inside functions to see which node types are used.

PiperOrigin-RevId: 163360375
This commit is contained in:
Pete Warden 2017-07-27 10:45:22 -07:00 committed by TensorFlower Gardener
parent 2139e7d8b1
commit 376bb063b7

View File

@ -24,6 +24,7 @@ limitations under the License.
// --in_graph=my_graph.pb
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
@ -254,6 +255,11 @@ Status SummarizeGraph(const GraphDef& graph, const string& graph_path,
for (const NodeDef& node : graph.node()) {
++op_counts[node.op()];
}
for (const FunctionDef& function : graph.library().function()) {
for (const NodeDef& node : function.node_def()) {
++op_counts[node.op()];
}
}
std::vector<std::pair<string, int>> op_counts_vec(op_counts.begin(),
op_counts.end());
std::sort(op_counts_vec.begin(), op_counts_vec.end(),