Merge pull request #16270 from dkurt:dnn_sort_tf_text_graphs

This commit is contained in:
Alexander Alekhin 2020-01-11 17:39:13 +00:00
commit 993b9af756
2 changed files with 13 additions and 5 deletions

View File

@ -950,6 +950,7 @@ void sortByExecutionOrder(tensorflow::GraphDef& net)
for (int i = 0; i < net.node_size(); ++i)
{
const tensorflow::NodeDef& node = net.node(i);
int numInputsInGraph = 0;
for (int j = 0; j < node.input_size(); ++j)
{
std::string inpName = node.input(j);
@ -957,22 +958,25 @@ void sortByExecutionOrder(tensorflow::GraphDef& net)
inpName = inpName.substr(inpName.find('^') + 1);
nodesMapIt = nodesMap.find(inpName);
CV_Assert(nodesMapIt != nodesMap.end());
edges[nodesMapIt->second].push_back(i);
if (nodesMapIt != nodesMap.end())
{
edges[nodesMapIt->second].push_back(i);
numInputsInGraph += 1;
}
}
if (node.input_size() == 0)
if (numInputsInGraph == 0)
nodesToAdd.push_back(i);
else
{
if (node.op() == "Merge" || node.op() == "RefMerge")
{
int numControlEdges = 0;
for (int j = 0; j < node.input_size(); ++j)
for (int j = 0; j < numInputsInGraph; ++j)
numControlEdges += node.input(j)[0] == '^';
numRefsToAdd[i] = numControlEdges + 1;
}
else
numRefsToAdd[i] = node.input_size();
numRefsToAdd[i] = numInputsInGraph;
}
}

View File

@ -715,6 +715,10 @@ void TFImporter::populateNet(Net dstNet)
simplifySubgraphs(netBin);
sortByExecutionOrder(netBin);
}
else
{
sortByExecutionOrder(netTxt);
}
std::set<String> layers_to_ignore;