Subgraph extraction in ONNX models

ngraph/frontend/onnx_import/include/onnx_import/editor/detail/subgraph_extraction.hpp

@@ -0,0 +1,164 @@
+//*****************************************************************************
+// Copyright 2017-2021 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#pragma once
+
+#include <map>
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace ONNX_NAMESPACE
+{
+    class GraphProto;
+    class NodeProto;
+    class ValueInfoProto;
+} // namespace ONNX_NAMESPACE
+
+namespace ngraph
+{
+    enum class EdgeType
+    {
+        INPUT,
+        OUTPUT
+    };
+
+    template <EdgeType>
+    struct Edge
+    {
+        Edge() = delete;
+        Edge(const int node_idx, std::string tensor_name)
+            : m_node_idx{node_idx}
+            , m_tensor_name{std::move(tensor_name)}
+        {
+        }
+
+        const int m_node_idx;
+        const std::string m_tensor_name;
+    };
+    namespace onnx_import
+    {
+        /// \brief Defines an edge connected to an input of any node in the graph.
+        ///        It consists of a node index in the processed ONNX model and the input name.
+        ///        The index should point to a node in the topological sort of the underlying graph
+        ///        which means it has to be in range:  0 <= node_idx < graph.node_size()
+        ///
+        ///        For a node number 5, with 3 inputs:
+        ///
+        ///            ----(in_A)---->  +--------+
+        ///            ----(in_B)---->  | node 5 |  ----(out)---->
+        ///            ----(in_C)---->  +--------+
+        ///
+        ///        there are 3 possible valid instances of this struct:
+        ///            InputEdge(5, "in_A")
+        ///            InputEdge(5, "in_B")
+        ///            InputEdge(5, "in_C")
+        using InputEdge = Edge<EdgeType::INPUT>;
+
+        /// \brief Defines an edge connected to an output of any node in the graph.
+        ///        It consists of a node index in the processed ONNX model and the output name.
+        ///
+        ///        For a node number 5, with 2 outputs:
+        ///
+        ///                             +--------+  ----(out1)---->
+        ///            ----(in_A)---->  | node 5 |
+        ///                             +--------+  ----(out2)---->
+        ///
+        ///        there are 2 possible valid instances of this struct:
+        ///            OutputEdge(5, "out1")
+        ///            OutputEdge(5, "out2")
+        using OutputEdge = Edge<EdgeType::OUTPUT>;
+
+        /// \brief Subgraph extraction helper structure
+        struct SubgraphExtractor
+        {
+            SubgraphExtractor(ONNX_NAMESPACE::GraphProto& graph);
+
+            /// \brief Adds new inputs to the graph and connects them to the nodes indicated by
+            ///        the provided input edges.
+            void add_new_inputs(const std::vector<InputEdge>& new_inputs);
+
+            /// \brief Adds new outputs to the graph with the same name as the nodes pointed to
+            ///        by the input edges "new_outputs".
+            void add_new_outputs(const std::vector<OutputEdge>& new_outputs);
+
+            /// \brief Extracts the final subgraph by traversing the original model bottom-up
+            ///        starting at each of the provided output edges. The extracted subgraph
+            ///        contains all previously added inputs and potentially a subset of original
+            ///        model's inputs that contribute to the value calculated in the output tensors.
+            ///        In the end the underlying GraphProto is modified and obsolete elements
+            ///        are discarded after this method call has finished.
+            ///
+            /// \param subgraph_outputs A list of expected outputs of the extracted subgraph.
+            void extract_subgraph(std::vector<OutputEdge> subgraph_outputs);
+
+            /// \brief Represents a subgraph of an ONNX model by holding a subset of nodes, inputs,
+            ///        outputs and initializers of the original graph. Objects of this struct can be
+            ///        merged into other instances using the += operator to build a subgraph from
+            ///        smaller clusters.
+            struct SubgraphComponents
+            {
+                SubgraphComponents() = default;
+                SubgraphComponents(const SubgraphComponents&) = delete;
+                SubgraphComponents(SubgraphComponents&&) = default;
+                SubgraphComponents& operator=(const SubgraphComponents&) = delete;
+                SubgraphComponents& operator=(SubgraphComponents&&) = default;
+
+                std::set<int> nodes;
+                std::set<std::string> inputs;
+                std::set<std::string> initializers;
+                std::set<std::string> outputs;
+
+                SubgraphComponents& operator+=(SubgraphComponents&& other)
+                {
+                    nodes.insert(other.nodes.begin(), other.nodes.end());
+                    inputs.insert(other.inputs.begin(), other.inputs.end());
+                    initializers.insert(other.initializers.begin(), other.initializers.end());
+                    outputs.insert(other.outputs.begin(), other.outputs.end());
+                    return *this;
+                }
+            };
+
+        private:
+            ONNX_NAMESPACE::GraphProto& m_onnx_graph;
+
+            // Graph traversal helper: node index -> node inputs (one-to-many)
+            std::unordered_multimap<int, std::string> m_node_inputs;
+            // Number of consumers of all tensors in the graph
+            std::map<std::string, int> m_tensor_consumers;
+
+            /// \brief Replaces the old input edge with a new one in the helper struct.
+            ///        This is used by the output contributors discovery.
+            void replace_input_edge(const InputEdge& old_edge, const InputEdge& new_edge);
+
+            /// \brief Returns a list of edges of each outputs of the graph "m_onnx_graph"
+            std::vector<OutputEdge> all_output_edges() const;
+
+            /// \brief Traverses the graph bottom-up and collects all nodes, inputs and initializers
+            ///        that contribute to an output designated by the provided output edge.
+            ///        A sum of such SubgraphComponents objects forms a target extracted subgraph.
+            SubgraphComponents
+                discover_output_contributors(const OutputEdge& output_edge,
+                                             const SubgraphComponents& already_collected) const;
+
+            /// \brief Modifies the underlying GraphProto object and discards all obsolete elements.
+            ///
+            /// \param subgraph An object describing the subgraph to be extracted (elems to be kept)
+            void extract_subgraph_from_onnx_model(const SubgraphComponents& subgraph);
+        };
+    } // namespace onnx_import
+} // namespace ngraph

ngraph/frontend/onnx_import/include/onnx_import/editor/editor.hpp

@@ -23,6 +23,7 @@
 #include "ngraph/op/constant.hpp"
 #include "ngraph/partial_shape.hpp"
 #include "ngraph/type/element_type.hpp"
+#include "onnx_import/editor/detail/subgraph_extraction.hpp"
 #include "onnx_import/utils/onnx_importer_visibility.hpp"
 
 namespace ONNX_NAMESPACE
@@ -53,7 +54,7 @@ namespace ngraph
             /// \param model_path Path to the file containing the model.
             ONNXModelEditor(const std::string& model_path);
 
-            /// \brief Modifies the in-memory representation of the model (m_model_proto) by setting
+            /// \brief Modifies the in-memory representation of the model by setting
             ///        custom input types for all inputs specified in the provided map.
             ///
             /// \param input_types A collection of pairs {input_name: new_input_type} that should be
@@ -62,7 +63,7 @@ namespace ngraph
             ///                    the inputs specified in its parameter.
             void set_input_types(const std::map<std::string, element::Type_t>& input_types);
 
-            /// \brief Modifies the in-memory representation of the model (m_model_proto) by setting
+            /// \brief Modifies the in-memory representation of the model by setting
             ///        custom input shapes for all inputs specified in the provided map.
             ///
             /// \param input_shapes A collection of pairs {input_name: new_input_shape} that should
@@ -71,6 +72,18 @@ namespace ngraph
             ///                     the inputs specified in its parameter.
             void set_input_shapes(const std::map<std::string, ngraph::PartialShape>& input_shapes);
 
+            /// \brief Extracts a subgraph constrained by input edges and output edges. In the end
+            ///        the underlying ModelProto is modified - obsolete inputs, initializers, nodes
+            ///        and outputs are removed from the in-memory model.
+            ///
+            /// \node Please look at the declaration of InputEdge and OutputEdge for explanation
+            ///       how those objects can be created. If the outputs parameter is empty
+            ///       this method keeps all of the original outputs of the model.
+            ///
+            /// \param inputs A collection of input edges which become new inputs to the graph
+            /// \param outputs A collection of output edges which become new outputs of the graph
+            void cut_graph_fragment(const std::vector<InputEdge>& inputs,
+                                    const std::vector<OutputEdge>& outputs);
             /// \brief Modifies the in-memory representation of the model by setting custom input
             ///        values for inputs specified in the provided map.
             ///
@@ -91,19 +104,28 @@ namespace ngraph
             /// \return A reference to ONNX ModelProto object containing the in-memory model
             ONNX_NAMESPACE::ModelProto& model() const;
 
+            /// \brief Returns a serialized ONNX model, possibly modified by the editor.
+            std::string model_string() const;
+
+            /// \brief Returns a list of all inputs of the in-memory model, including initializers.
+            ///        The returned value might depend on the previous operations executed on an
+            ///        instance of the model editor, in particular the subgraph extraction which
+            ///        can discard some inputs and initializers from the original graph.
+            std::vector<std::string> model_inputs() const;
+
             /// \brief Returns the path to the original model file
             const std::string& model_path() const;
 
-            /// \brief Saves the possibly model held by this class to a file. Serializes in binary
-            /// mode.
+            /// \brief Saves the possibly modified model held by this class to a file.
+            /// Serializes in binary mode.
             ///
             /// \param out_file_path A path to the file where the modified model should be dumped.
             void serialize(const std::string& out_file_path) const;
 
         private:
             const std::string m_model_path;
 
-            class Impl;
+            struct Impl;
             std::unique_ptr<Impl, void (*)(Impl*)> m_pimpl;
         };
     } // namespace onnx_import