Add 'get_num_inputs' to GraphRuntime (#6118)

apache · Jul 24, 2020 · bfa4eae · bfa4eae
1 parent 0a1c4c2
commit bfa4eae
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 0 deletions.
diff --git a/python/tvm/contrib/graph_runtime.py b/python/tvm/contrib/graph_runtime.py
@@ -133,6 +133,7 @@ def __init__(self, module):
         self._get_output = module["get_output"]
         self._get_input = module["get_input"]
         self._get_num_outputs = module["get_num_outputs"]
+        self._get_num_inputs = module["get_num_inputs"]
         self._load_params = module["load_params"]
         self._share_params = module["share_params"]
 
@@ -187,6 +188,16 @@ def get_num_outputs(self):
         """
         return self._get_num_outputs()
 
+    def get_num_inputs(self):
+        """Get the number of inputs to the graph
+
+        Returns
+        -------
+        count : int
+            The number of inputs.
+        """
+        return self._get_num_inputs()
+
     def get_input(self, index, out=None):
         """Get index-th input to out
 

diff --git a/src/runtime/graph/graph_runtime.cc b/src/runtime/graph/graph_runtime.cc
@@ -134,6 +134,12 @@ void GraphRuntime::SetInputZeroCopy(int index, DLTensor* data_ref) {
  * \return The number of outputs from graph.
  */
 int GraphRuntime::NumOutputs() const { return outputs_.size(); }
+/*!
+ * \brief Get the number of inputs
+ *
+ * \return The number of inputs to the graph.
+ */
+int GraphRuntime::NumInputs() const { return input_nodes_.size(); }
 /*!
  * \brief Return NDArray for given input index.
  * \param index The input index.
@@ -433,6 +439,9 @@ PackedFunc GraphRuntime::GetFunction(const std::string& name,
   } else if (name == "get_num_outputs") {
     return PackedFunc(
         [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->NumOutputs(); });
+  } else if (name == "get_num_inputs") {
+    return PackedFunc(
+        [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->NumInputs(); });
   } else if (name == "run") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { this->Run(); });
   } else if (name == "load_params") {

diff --git a/src/runtime/graph/graph_runtime.h b/src/runtime/graph/graph_runtime.h
@@ -124,6 +124,12 @@ class TVM_DLL GraphRuntime : public ModuleNode {
    * \return The number of outputs from graph.
    */
   int NumOutputs() const;
+  /*!
+   * \brief Get the number of inputs
+   *
+   * \return The number of inputs to the graph.
+   */
+  int NumInputs() const;
   /*!
    * \brief Return NDArray for given input index.
    * \param index The input index.