Look inside functions to see which node types are used.
PiperOrigin-RevId: 163360375
This commit is contained in:
parent
2139e7d8b1
commit
376bb063b7
@ -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(),
|
||||
|
Loading…
Reference in New Issue
Block a user