ONNX graphs simplifier

modules/dnn/src/graph_simplifier.cpp

+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+
+// Copyright (C) 2020, Intel Corporation, all rights reserved.
+// Third party copyrights are property of their respective owners.
+
+#include "precomp.hpp"
+
+#include "graph_simplifier.hpp"
+
+#include <queue>
+
+namespace cv { namespace dnn {
+
+Subgraph::~Subgraph() {}
+
+int Subgraph::addNodeToMatch(const std::string& op, int input_0, int input_1,
+                             int input_2, int input_3)
+{
+    int nodeInputs[] = {input_0, input_1, input_2, input_3};
+    int numInputs = 0;
+    for (int i = 0; i < 4; ++i)
+    {
+        numInputs += (int)(nodeInputs[i] != -1);
+    }
+    return addNodeToMatch(op, std::vector<int>(&nodeInputs[0], &nodeInputs[0] + numInputs));
+}
+
+int Subgraph::addNodeToMatch(const std::string& op, const std::vector<int>& inputs_)
+{
+    for (int i = 0; i < inputs_.size(); ++i)
+    {
+        CV_Assert(inputs_[i] < (int)nodes.size());
+    }
+    nodes.push_back(op);
+    inputs.push_back(inputs_);
+    return nodes.size() - 1;
+}
+
+void Subgraph::setFusedNode(const std::string& op, int input_0, int input_1,
+                            int input_2, int input_3, int input_4, int input_5)
+{
+    int nodeInputs[] = {input_0, input_1, input_2, input_3, input_4, input_5};
+    int numInputs = 0;
+    for (int i = 0; i < 6; ++i)
+    {
+        CV_Assert(nodeInputs[i] < (int)nodes.size());
+        numInputs += (int)(nodeInputs[i] != -1);
+    }
+    setFusedNode(op, std::vector<int>(&nodeInputs[0], &nodeInputs[0] + numInputs));
+}
+
+void Subgraph::setFusedNode(const std::string& op, const std::vector<int>& inputs_)
+{
+    fusedNodeInputs = inputs_;
+    fusedNodeOp = op;
+}
+
+int Subgraph::getInputNodeId(const Ptr<ImportGraphWrapper>& net,
+                             const Ptr<ImportNodeWrapper>& node,
+                             int inpId)
+{
+    CV_Assert(inpId < node->getNumInputs());
+    std::string name = node->getInputName(inpId);
+    // If operation produces several tensors, they are specified by index
+    // after ':' character. In example, "input:0".
+    name = name.substr(0, name.rfind(':'));
+    const int numNodes = net->getNumNodes();
+    for (int i = 0; i < numNodes; ++i)
+    {
+        if (net->getNodeName(i) == name)
+            return i;
+    }
+    CV_Error(Error::StsParseError, "Input node with name " + name + " not found");
+}
+
+bool Subgraph::match(const Ptr<ImportGraphWrapper>& net, int nodeId,
+                     std::vector<int>& matchedNodesIds,
+                     std::vector<int>& targetNodesIds)
+{
+    matchedNodesIds.clear();
+    targetNodesIds.clear();
+
+    std::queue<int> nodesToMatch;
+    std::queue<int> targetNodes;
+    nodesToMatch.push(nodeId);
+    targetNodes.push(nodes.size() - 1);
+    while (!nodesToMatch.empty())
+    {
+        int nodeToMatch = nodesToMatch.front();
+        int targetNodeId = targetNodes.front();
+        nodesToMatch.pop();
+        targetNodes.pop();
+
+        if (std::find(matchedNodesIds.begin(), matchedNodesIds.end(), nodeToMatch) !=
+            matchedNodesIds.end())
+            continue;
+
+        const Ptr<ImportNodeWrapper> node = net->getNode(nodeToMatch);
+        if (node->getType() != nodes[targetNodeId])
+            return false;
+
+        std::vector<int>& inputNodes = inputs[targetNodeId];
+        if (inputNodes.size() != node->getNumInputs())
+            return false;
+
+        for (int j = 0; j < inputNodes.size(); ++j)
+        {
+            if (nodes[inputNodes[j]].empty())  // Unknown input node type.
+                continue;
+            nodeId = getInputNodeId(net, node, j);
+            const Ptr<ImportNodeWrapper> inpNode = net->getNode(nodeId);
+            if (inpNode->getType() != "Const")
+            {
+                nodesToMatch.push(nodeId);
+                targetNodes.push(inputNodes[j]);
+            }
+            else if (nodes[inputNodes[j]] != "Const")
+                return false;
+        }
+        matchedNodesIds.push_back(nodeToMatch);
+        targetNodesIds.push_back(targetNodeId);
+    }
+
+    const int n = matchedNodesIds.size();
+    std::vector<std::pair<int, int> > elements(n);
+    for (int i = 0; i < n; ++i)
+        elements[i] = std::make_pair(matchedNodesIds[i], targetNodesIds[i]);
+    std::sort(elements.begin(), elements.end());
+    for (int i = 0; i < n; ++i)
+    {
+        matchedNodesIds[i] = elements[i].first;
+        targetNodesIds[i] = elements[i].second;
+    }
+    return true;
+}
+
+void Subgraph::replace(const Ptr<ImportGraphWrapper>& net, const std::vector<int>& matchedNodesIds,
+                       const std::vector<int>& targetNodesIds)
+{
+    // Extract names of input nodes.
+    std::vector<std::string> inputsNames(fusedNodeInputs.size());
+    for (int i = 0; i < fusedNodeInputs.size(); ++i)
+    {
+        std::string inpName;
+        // Find input node name looking at inputs of fused nodes.
+        for (int j = 0; j < matchedNodesIds.size() && inpName.empty(); ++j)
+        {
+            Ptr<ImportNodeWrapper> node = net->getNode(matchedNodesIds[j]);
+            std::vector<int>& inpIndices = inputs[targetNodesIds[j]];
+
+            CV_Assert(node->getNumInputs() == inpIndices.size());
+            for (int k = 0; k < inpIndices.size(); ++k)
+            {
+                if (inpIndices[k] == fusedNodeInputs[i])
+                {
+                    inpName = node->getInputName(k);
+                    break;
+                }
+            }
+        }
+        CV_Assert(!inpName.empty());
+        inputsNames[i] = inpName;
+    }
+
+    // Remove matched nodes except the last one. Indices in ascending order are expected.
+    Ptr<ImportNodeWrapper> node = net->getNode(matchedNodesIds.back());
+    for (int i = matchedNodesIds.size() - 2; i >= 0; --i)
+        net->removeNode(matchedNodesIds[i]);
+
+    // Modify the last node to be a fused one.
+    node->setType(fusedNodeOp);
+    node->setInputNames(inputsNames);
+
+    std::vector<Ptr<ImportNodeWrapper> > inputNodes(inputsNames.size());
+    for (int i = 0; i < inputsNames.size(); ++i)
+    {
+        inputNodes[i] = net->getNode(getInputNodeId(net, node, i));
+    }
+    finalize(net, node, inputNodes);
+}
+
+void Subgraph::finalize(const Ptr<ImportGraphWrapper>& net,
+                        const Ptr<ImportNodeWrapper>& fusedNode,
+                        std::vector<Ptr<ImportNodeWrapper> >& inputs) {}
+
+void simplifySubgraphs(const Ptr<ImportGraphWrapper>& net,
+                       const std::vector<Ptr<Subgraph> >& patterns)
+{
+    int numNodes = net->getNumNodes();
+    std::vector<int> matchedNodesIds, targetNodesIds;
+    for (int i = 0; i < numNodes; ++i)
+    {
+        for (int j = 0; j < patterns.size(); ++j)
+        {
+            if (patterns[j]->match(net, i, matchedNodesIds, targetNodesIds))
+            {
+                patterns[j]->replace(net, matchedNodesIds, targetNodesIds);
+                numNodes -= matchedNodesIds.size() - 1;  // #matchedNodes removed and one added.
+                break;
+            }
+        }
+    }
+}
+
+}}  // namespace cv::dnn

modules/dnn/src/graph_simplifier.hpp

+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+
+// Copyright (C) 2020, Intel Corporation, all rights reserved.
+// Third party copyrights are property of their respective owners.
+
+#ifndef __OPENCV_DNN_GRAPH_SIMPLIFIER_HPP__
+#define __OPENCV_DNN_GRAPH_SIMPLIFIER_HPP__
+
+#include <string>
+
+#include <opencv2/core.hpp>
+
+namespace cv { namespace dnn {
+
+class ImportNodeWrapper
+{
+public:
+    virtual ~ImportNodeWrapper() {};
+
+    virtual int getNumInputs() const = 0;
+
+    virtual std::string getInputName(int idx) const = 0;
+
+    virtual std::string getType() const = 0;
+
+    virtual void setType(const std::string& type) = 0;
+
+    virtual void setInputNames(const std::vector<std::string>& inputs) = 0;
+};
+
+class ImportGraphWrapper
+{
+public:
+    virtual ~ImportGraphWrapper() {};
+
+    virtual Ptr<ImportNodeWrapper> getNode(int idx) const = 0;
+
+    virtual int getNumNodes() const = 0;
+
+    virtual std::string getNodeName(int idx) const = 0;
+
+    virtual void removeNode(int idx) = 0;
+};
+
+class Subgraph  // Interface to match and replace subgraphs.
+{
+public:
+    virtual ~Subgraph();
+
+    // Add a node to be matched in the origin graph. Specify ids of nodes that
+    // are expected to be inputs. Returns id of a newly added node.
+    // TODO: Replace inputs to std::vector<int> in C++11
+    int addNodeToMatch(const std::string& op, int input_0 = -1, int input_1 = -1,
+                       int input_2 = -1, int input_3 = -1);
+
+    int addNodeToMatch(const std::string& op, const std::vector<int>& inputs_);
+
+    // Specify resulting node. All the matched nodes in subgraph excluding
+    // input nodes will be fused into this single node.
+    // TODO: Replace inputs to std::vector<int> in C++11
+    void setFusedNode(const std::string& op, int input_0 = -1, int input_1 = -1,
+                      int input_2 = -1, int input_3 = -1, int input_4 = -1,
+                      int input_5 = -1);
+
+    void setFusedNode(const std::string& op, const std::vector<int>& inputs_);
+
+    static int getInputNodeId(const Ptr<ImportGraphWrapper>& net,
+                              const Ptr<ImportNodeWrapper>& node,
+                              int inpId);
+
+    // Match TensorFlow subgraph starting from <nodeId> with a set of nodes to be fused.
+    // Const nodes are skipped during matching. Returns true if nodes are matched and can be fused.
+    virtual bool match(const Ptr<ImportGraphWrapper>& net, int nodeId,
+                       std::vector<int>& matchedNodesIds,
+                       std::vector<int>& targetNodesIds);
+
+    // Fuse matched subgraph.
+    void replace(const Ptr<ImportGraphWrapper>& net, const std::vector<int>& matchedNodesIds,
+                 const std::vector<int>& targetNodesIds);
+
+    virtual void finalize(const Ptr<ImportGraphWrapper>& net,
+                          const Ptr<ImportNodeWrapper>& fusedNode,
+                          std::vector<Ptr<ImportNodeWrapper> >& inputs);
+
+private:
+    std::vector<std::string> nodes;         // Nodes to be matched in the origin graph.
+    std::vector<std::vector<int> > inputs;  // Connections of an every node to it's inputs.
+
+    std::string fusedNodeOp;           // Operation name of resulting fused node.
+    std::vector<int> fusedNodeInputs;  // Inputs of fused node.
+};
+
+void simplifySubgraphs(const Ptr<ImportGraphWrapper>& net,
+                       const std::vector<Ptr<Subgraph> >& patterns);
+
+}}  // namespace dnn, namespace cv
+
+#endif  // __OPENCV_DNN_GRAPH_SIMPLIFIER_HPP__

modules/dnn/src/onnx/onnx_graph_simplifier.cpp

+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+
+// Copyright (C) 2020, Intel Corporation, all rights reserved.
+// Third party copyrights are property of their respective owners.
+
+#include "../precomp.hpp"
+
+#include "../graph_simplifier.hpp"
+#include "onnx_graph_simplifier.hpp"
+
+#include <queue>
+
+namespace cv { namespace dnn {
+CV__DNN_EXPERIMENTAL_NS_BEGIN
+
+// This wrapper can behave differently for fake input nodes and real graph nodes.
+class ONNXNodeWrapper : public ImportNodeWrapper
+{
+public:
+    ONNXNodeWrapper(opencv_onnx::NodeProto* _node = 0) : node(_node) {}
+
+    virtual int getNumInputs() const CV_OVERRIDE
+    {
+        return node ? node->input_size() : 0;
+    }
+
+    virtual std::string getInputName(int idx) const CV_OVERRIDE
+    {
+        CV_Assert_N(node, idx < node->input_size());
+        return node->input(idx);
+    }
+
+    virtual std::string getType() const CV_OVERRIDE
+    {
+        return node ? node->op_type() : "";
+    }
+
+    virtual void setType(const std::string& type) CV_OVERRIDE
+    {
+        CV_Assert(node);
+        node->set_op_type(type);
+    }
+
+    virtual void setInputNames(const std::vector<std::string>& inputs) CV_OVERRIDE
+    {
+        CV_Assert(node);
+        node->clear_input();
+        for (int i = 0; i < inputs.size(); ++i)
+            node->add_input(inputs[i]);
+    }
+
+    opencv_onnx::NodeProto* node;
+};
+
+// ONNX graph's inputs are separate from nodes so we index them before the rest of nodes.
+class ONNXGraphWrapper : public ImportGraphWrapper
+{
+public:
+    ONNXGraphWrapper(opencv_onnx::GraphProto& _net) : net(_net)
+    {
+        numInputs = net.input_size();
+    }
+
+    virtual Ptr<ImportNodeWrapper> getNode(int idx) const CV_OVERRIDE
+    {
+        opencv_onnx::NodeProto* node = 0;
+        if (idx >= numInputs)
+            node = net.mutable_node(idx - numInputs);
+        return makePtr<ONNXNodeWrapper>(node);
+    }
+
+    virtual int getNumNodes() const CV_OVERRIDE
+    {
+        return numInputs + net.node_size();
+    }
+
+    virtual std::string getNodeName(int idx) const CV_OVERRIDE
+    {
+        if (idx < numInputs)
+            return net.input(idx).name();
+        else
+            return net.node(idx - numInputs).output(0);
+    }
+
+    virtual void removeNode(int idx) CV_OVERRIDE
+    {
+        CV_Assert(idx >= numInputs);
+        net.mutable_node()->DeleteSubrange(idx - numInputs, 1);
+    }
+
+private:
+    int numInputs;
+    opencv_onnx::GraphProto& net;
+};
+
+class SoftMaxSubgraph : public Subgraph
+{
+public:
+    SoftMaxSubgraph()
+    {
+        int input = addNodeToMatch("");
+        int inpExp = addNodeToMatch("Exp", input);
+        int sum = addNodeToMatch("ReduceSum", inpExp);
+        addNodeToMatch("Div", inpExp, sum);
+        setFusedNode("Softmax", input);
+    }
+
+    virtual bool match(const Ptr<ImportGraphWrapper>& net, int nodeId,
+                       std::vector<int>& matchedNodesIds,
+                       std::vector<int>& targetNodesIds) CV_OVERRIDE
+    {
+        if (Subgraph::match(net, nodeId, matchedNodesIds, targetNodesIds))
+        {
+            Ptr<ImportNodeWrapper> sum = net->getNode(matchedNodesIds[1]);
+            opencv_onnx::NodeProto* node = sum.dynamicCast<ONNXNodeWrapper>()->node;
+
+            for (int i = 0; i < node->attribute_size(); i++)
+            {
+                opencv_onnx::AttributeProto attr = node->attribute(i);
+                if (attr.name() != "axes")
+                    continue;
+                if (attr.ints_size() != 1)
+                    CV_Error(Error::StsNotImplemented, format("Unexpected number of axes: %d", attr.ints_size()));
+                axis = attr.ints(0);
+                return true;
+            }
+            CV_Error(Error::StsNotImplemented, "Missed axes attribute");
+        }
+        return false;
+    }
+
+    virtual void finalize(const Ptr<ImportGraphWrapper>&,
+                          const Ptr<ImportNodeWrapper>& fusedNode,
+                          std::vector<Ptr<ImportNodeWrapper> >&) CV_OVERRIDE
+    {
+        opencv_onnx::NodeProto* node = fusedNode.dynamicCast<ONNXNodeWrapper>()->node;
+        opencv_onnx::AttributeProto* attr = node->add_attribute();
+        attr->set_name("axis");
+        attr->set_i(axis);
+    }
+
+private:
+    int axis;
+};
+
+void simplifySubgraphs(opencv_onnx::GraphProto& net)
+{
+    std::vector<Ptr<Subgraph> > subgraphs;
+    subgraphs.push_back(makePtr<SoftMaxSubgraph>());
+
+    simplifySubgraphs(Ptr<ImportGraphWrapper>(new ONNXGraphWrapper(net)), subgraphs);
+}
+
+CV__DNN_EXPERIMENTAL_NS_END
+}}  // namespace cv::dnn

modules/dnn/src/onnx/onnx_graph_simplifier.hpp

+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+
+// Copyright (C) 2020, Intel Corporation, all rights reserved.
+// Third party copyrights are property of their respective owners.
+
+#ifndef __OPENCV_DNN_ONNX_SIMPLIFIER_HPP__
+#define __OPENCV_DNN_ONNX_SIMPLIFIER_HPP__
+
+#include "../precomp.hpp"
+
+#if defined(__GNUC__) && __GNUC__ >= 5
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wsuggest-override"
+#endif
+#include "opencv-onnx.pb.h"
+#if defined(__GNUC__) && __GNUC__ >= 5
+#pragma GCC diagnostic pop
+#endif
+
+namespace cv { namespace dnn {
+CV__DNN_EXPERIMENTAL_NS_BEGIN
+
+void simplifySubgraphs(opencv_onnx::GraphProto& net);
+
+CV__DNN_EXPERIMENTAL_NS_END
+}}  // namespace dnn, namespace cv
+
+#endif  // __OPENCV_DNN_ONNX_SIMPLIFIER_HPP__

modules/dnn/src/onnx/onnx_importer.cpp

@@ -26,6 +26,8 @@
 #pragma GCC diagnostic pop
 #endif
 
+#include "onnx_graph_simplifier.hpp"
+
 namespace cv {
 namespace dnn {
 CV__DNN_EXPERIMENTAL_NS_BEGIN
@@ -326,6 +328,9 @@ void ONNXImporter::populateNet(Net dstNet)
 {
     CV_Assert(model_proto.has_graph());
     opencv_onnx::GraphProto graph_proto = model_proto.graph();
+
+    simplifySubgraphs(graph_proto);
+
     std::map<std::string, Mat> constBlobs = getGraphTensors(graph_proto);
     // List of internal blobs shapes.
     std::map<std::string, MatShape> outShapes;

modules/dnn/test/test_onnx_importer.cpp

@@ -396,6 +396,7 @@ TEST_P(Test_ONNX_layers, Softmax)
 {
     testONNXModels("softmax");
     testONNXModels("log_softmax", npy, 0, 0, false, false);
+    testONNXModels("softmax_unfused");
 }
 
 TEST_P(Test_ONNX_layers, Split_EltwiseMax)