Python API for inference model saving/load (#5020)

JiayiFeng · web-flow · commit 6783dcee9e3e · 2017-10-27T17:48:48.000-07:00
* Add `dump_to_file()` for ProgrameDescBind in pybind

* Update

* Add utility.py

* typo

* Fix bugs

* Move add_feed/fetch_components to untility.py

* Compelete dump

* Follow comments

* Change output of Prune() from inference to pointer

* Expose Prune() to Python

* Compelete save/load API of inference model

* Fix errors

* Debuging

* Compelete unit tests

* follow comments
diff --git a/.gitignore b/.gitignore
@@ -28,3 +28,4 @@ cmake_install.cmake
 paddle/.timestamp
 python/paddlepaddle.egg-info/
 paddle/pybind/pybind.h
+python/paddle/v2/framework/tests/tmp/*
diff --git a/paddle/framework/op_desc.h b/paddle/framework/op_desc.h
@@ -107,6 +107,8 @@ class OpDescBind {
 
   void InferVarType(BlockDescBind *block) const;
 
+  void MarkAsTarget() { desc_.set_is_target(true); }
+
   void Flush();
 
  private:
diff --git a/paddle/framework/program_desc.cc b/paddle/framework/program_desc.cc
@@ -49,6 +49,13 @@ ProgramDescBind::ProgramDescBind(const ProgramDescBind &o) {
   }
 }
 
+ProgramDescBind::ProgramDescBind(const ProgramDesc &desc) {
+  desc_ = desc;
+  for (auto &block_desc : *desc_.mutable_blocks()) {
+    blocks_.emplace_back(new BlockDescBind(this, &block_desc));
+  }
+}
+
 ProgramDescBind::ProgramDescBind(const std::string &binary_str) {
   PADDLE_ENFORCE(desc_.ParseFromString(binary_str),
                  "Fail to parse program_desc from binary string.");
diff --git a/paddle/framework/program_desc.h b/paddle/framework/program_desc.h
@@ -29,6 +29,8 @@ class ProgramDescBind {
  public:
   ProgramDescBind();
 
+  explicit ProgramDescBind(const ProgramDesc &desc);
+
   ProgramDescBind(const ProgramDescBind &o);
 
   explicit ProgramDescBind(const std::string &binary_str);
diff --git a/paddle/framework/prune.cc b/paddle/framework/prune.cc
@@ -46,7 +46,7 @@ bool IsTarget(const OpDesc& op_desc) {
   return false;
 }
 
-void prune_impl(const ProgramDesc& input, ProgramDesc& output, int block_id) {
+void prune_impl(const ProgramDesc& input, ProgramDesc* output, int block_id) {
   // TODO(tonyyang-svail):
   //    - will change to use multiple blocks for RNN op and Cond Op
 
@@ -91,8 +91,8 @@ void prune_impl(const ProgramDesc& input, ProgramDesc& output, int block_id) {
   // we reverse the should_run vector
   std::reverse(should_run.begin(), should_run.end());
 
-  output = input;
-  auto* op_field = output.mutable_blocks(block_id)->mutable_ops();
+  *output = input;
+  auto* op_field = output->mutable_blocks(block_id)->mutable_ops();
   op_field->Clear();
   for (size_t i = 0; i < should_run.size(); ++i) {
     if (should_run[i]) {
@@ -101,7 +101,8 @@ void prune_impl(const ProgramDesc& input, ProgramDesc& output, int block_id) {
   }
 }
 
-void Prune(const ProgramDesc& input, ProgramDesc& output) {
+// TODO(fengjiayi): Prune() could be inplaced to avoid unnecessary copies
+void Prune(const ProgramDesc& input, ProgramDesc* output) {
   prune_impl(input, output, 0);
 }
 
diff --git a/paddle/framework/prune.h b/paddle/framework/prune.h
@@ -20,7 +20,7 @@ limitations under the License. */
 namespace paddle {
 namespace framework {
 
-void Prune(const ProgramDesc& input, ProgramDesc& output);
+void Prune(const ProgramDesc& input, ProgramDesc* output);
 
 }  // namespace framework
 }  // namespace paddle
diff --git a/paddle/framework/prune_test.cc b/paddle/framework/prune_test.cc
@@ -59,11 +59,11 @@ TEST(Prune, one_operator) {
   f::ProgramDesc *pdesc = program.Proto();
   f::ProgramDesc pruned;
 
-  Prune(*pdesc, pruned);
+  Prune(*pdesc, &pruned);
   PADDLE_ENFORCE_EQ(pruned.blocks(0).ops_size(), 0);
 
   pdesc->mutable_blocks(0)->mutable_ops(0)->set_is_target(true);
-  Prune(*pdesc, pruned);
+  Prune(*pdesc, &pruned);
   PADDLE_ENFORCE_EQ(pruned.blocks(0).ops_size(), 1);
 }
 
@@ -81,7 +81,7 @@ TEST(Prune, forward) {
   for (int i = 0; i < pdesc->blocks(0).ops_size(); ++i) {
     f::ProgramDesc pruned;
     pdesc->mutable_blocks(0)->mutable_ops(i)->set_is_target(true);
-    Prune(*pdesc, pruned);
+    Prune(*pdesc, &pruned);
     PADDLE_ENFORCE_EQ(pruned.blocks(0).ops_size(), i + 1);
   }
 }
@@ -100,7 +100,7 @@ TEST(Prune, multi_input_op) {
   pdesc->mutable_blocks(0)->mutable_ops(3)->set_is_target(true);
 
   f::ProgramDesc pruned;
-  Prune(*pdesc, pruned);
+  Prune(*pdesc, &pruned);
   PADDLE_ENFORCE_EQ(pruned.blocks(0).ops_size(), 4);
 }
 
@@ -116,7 +116,7 @@ TEST(Prune, multi_output_op) {
   pdesc->mutable_blocks(0)->mutable_ops(2)->set_is_target(true);
 
   f::ProgramDesc pruned;
-  Prune(*pdesc, pruned);
+  Prune(*pdesc, &pruned);
   PADDLE_ENFORCE_EQ(pruned.blocks(0).ops_size(), 2);
 }
 
@@ -133,6 +133,6 @@ TEST(Prune, multi_target) {
   pdesc->mutable_blocks(0)->mutable_ops(2)->set_is_target(true);
 
   f::ProgramDesc pruned;
-  Prune(*pdesc, pruned);
+  Prune(*pdesc, &pruned);
   PADDLE_ENFORCE_EQ(pruned.blocks(0).ops_size(), 3);
 }
diff --git a/paddle/pybind/CMakeLists.txt b/paddle/pybind/CMakeLists.txt
@@ -1,7 +1,7 @@
 if(WITH_PYTHON)
   cc_library(paddle_pybind SHARED
     SRCS pybind.cc exception.cc protobuf.cc
-    DEPS pybind python backward proto_desc tensor_array paddle_memory executor
+    DEPS pybind python backward proto_desc tensor_array paddle_memory executor prune
     ${GLOB_OP_LIB})
 endif(WITH_PYTHON)
 
diff --git a/paddle/pybind/protobuf.cc b/paddle/pybind/protobuf.cc
@@ -141,6 +141,13 @@ void BindProgramDesc(py::module &m) {
                  desc->SerializeToString(&res),
                  "Serialize ProgramDesc Error. This could be a bug of Paddle.");
              return res;
+           })
+      .def("parse_from_string",
+           [](ProgramDescBind &program_desc, const std::string &data) {
+             ProgramDesc *desc = program_desc.Proto();
+             PADDLE_ENFORCE(desc->ParseFromString(data),
+                            "Fail to parse ProgramDesc from string. This could "
+                            "be a bug of Paddle.");
            });
 }
 
diff --git a/paddle/pybind/pybind.cc b/paddle/pybind/pybind.cc
@@ -19,6 +19,7 @@ limitations under the License. */
 #include "paddle/framework/feed_fetch_method.h"
 #include "paddle/framework/framework.pb.h"
 #include "paddle/framework/lod_tensor.h"
+#include "paddle/framework/prune.h"
 #include "paddle/framework/selected_rows.h"
 #include "paddle/framework/tensor_array.h"
 #include "paddle/operators/cond_op.h"
@@ -237,6 +238,16 @@ All parameter, weight, gradient are variables in Paddle.
     }
     return ret_values;
   });
+  m.def("prune", [](const ProgramDescBind &origin,
+                    const std::vector<std::array<size_t, 2>> &targets) {
+    ProgramDescBind prog_with_targets(origin);
+    for (const auto &t : targets) {
+      prog_with_targets.Block(t[0])->Op(t[1])->MarkAsTarget();
+    }
+    ProgramDesc pruned_desc;
+    Prune(*prog_with_targets.Proto(), &pruned_desc);
+    return new ProgramDescBind(pruned_desc);
+  });
   m.def_submodule(
        "var_names",
        "The module will return special predefined variable name in Paddle")
diff --git a/python/paddle/v2/framework/framework.py b/python/paddle/v2/framework/framework.py
@@ -251,6 +251,8 @@ def find_name(var_list, name):
                 self.desc.set_output(out_proto.name, out_argu_names)
 
         if attrs is not None:
+            if not isinstance(attrs, dict):
+                raise TypeError("'attrs' should be a dict.")
             for attr in proto.attrs:
                 attr_name = attr.name
                 if (not attr_name in attrs) or (attrs[attr_name] is None):
@@ -291,6 +293,14 @@ def output(self, name):
     def output_names(self):
         return self.desc.output_names()
 
+    @property
+    def idx(self):
+        for i, op in enumerate(self.block.ops):
+            if op == self:
+                return i
+        raise ValueError(
+            "Can't find op itself in it's block. It could be a bug of Paddle.")
+
     def has_attr(self, name):
         return self.desc.has_attr(name)
 
@@ -440,10 +450,31 @@ def clone(self):
         p.sync_with_cpp()
         return p
 
+    def prune(self, targets):
+        if not isinstance(targets, list):
+            targets = [targets]
+        targets_idx = []
+        for t in targets:
+            if not isinstance(t, Operator):
+                if isinstance(t, Variable):
+                    t = t.op
+                else:
+                    raise ValueError(
+                        "All targets of prune() can only be Variable or Operator."
+                    )
+
+            targets_idx.append([t.block.idx, t.idx])
+        res = Program()
+        res.desc = core.prune(self.desc, targets_idx)
+        res.blocks = [Block(res, i) for i in xrange(res.desc.num_blocks())]
+        res.sync_with_cpp()
+        return res
+
     @staticmethod
     def parse_from_string(binary_str):
         p = Program()
         p.desc = core.ProgramDesc(binary_str)
+        p.blocks = [Block(p, i) for i in xrange(p.desc.num_blocks())]
         p.sync_with_cpp()
         return p
 
diff --git a/python/paddle/v2/framework/io.py b/python/paddle/v2/framework/io.py
@@ -1,11 +1,12 @@
 import os
+import cPickle as pickle
 
 from paddle.v2.framework.framework import Program, Parameter, g_program, \
     Variable
 
 __all__ = [
     'save_vars', 'save_params', 'save_persistables', 'load_vars', 'load_params',
-    'load_persistables'
+    'load_persistables', "save_inference_model", "load_inference_model"
 ]
 
 
@@ -31,7 +32,7 @@ def _clone_var_in_block_(block, var):
 def save_vars(executor, dirname, program=None, vars=None, predicate=None):
     """
     Save variables to directory by executor.
-    
+
     :param executor: executor that save variable
     :param dirname: directory path
     :param program: program. If vars is None, then filter all variables in this 
@@ -92,7 +93,7 @@ def save_persistables(executor, dirname, program=None):
 def load_vars(executor, dirname, program=None, vars=None, predicate=None):
     """
     Load variables from directory by executor.
-    
+
     :param executor: executor that save variable
     :param dirname: directory path
     :param program: program. If vars is None, then filter all variables in this 
@@ -124,6 +125,7 @@ def load_vars(executor, dirname, program=None, vars=None, predicate=None):
                 inputs={},
                 outputs={"Out": [new_var]},
                 attrs={'file_path': os.path.join(dirname, new_var.name)})
+
         executor.run(load_prog)
 
 
@@ -141,3 +143,88 @@ def load_persistables(executor, dirname, program=None):
     """
     load_vars(
         executor, dirname=dirname, program=program, predicate=is_persistable)
+
+
+def save_inference_model(dirname,
+                         feeded_var_names,
+                         target_vars,
+                         executor,
+                         program=None):
+    """
+    Build a model especially for inference, 
+    and save it to directory by the executor.
+
+    :param dirname: directory path
+    :param feeded_var_names: Names of variables that need to be feeded data during inference
+    :param target_vars: Variables from which we can get inference results.
+    :param executor: executor that save inference model
+    :param program: original program, which will be pruned to build the inference model. 
+    Default g_program.
+
+    :return: None
+    """
+    if program is None:
+        program = g_program
+    if not isinstance(target_vars, list):
+        target_vars = [target_vars]
+
+    if not os.path.isdir(dirname):
+        os.makedirs(dirname)
+
+    pruned_program = program.prune(target_vars)
+    fetch_var_names = [v.name for v in target_vars]
+
+    model_file_name = dirname + "/__model__"
+    with open(model_file_name, "w") as f:
+        pickle.dump({
+            "program_desc_str": pruned_program.desc.serialize_to_string(),
+            "feed_var_names": feeded_var_names,
+            "fetch_var_names": fetch_var_names
+        }, f, -1)
+
+    save_params(executor, dirname, program)
+
+
+def load_persistables_if_exist(executor, dirname, program=None):
+    filenames = next(os.walk(dirname))[2]
+    filenames = set(filenames)
+
+    def _is_presistable_and_exist_(var):
+        if not is_persistable(var):
+            return False
+        else:
+            return var.name in filenames
+
+    load_vars(
+        executor,
+        dirname,
+        program=program,
+        vars=None,
+        predicate=_is_presistable_and_exist_)
+
+
+def load_inference_model(dirname, executor):
+    """
+    Load inference model from a directory
+
+    :param dirname: directory path
+    :param executor: executor that load inference model
+
+    :return: [program, feed_var_names, fetch_var_names]
+             program: program especially for inference.
+             feeded_var_names: Names of variables that need to feed data
+             fetch_vars: Variables from which we can get inference results.
+    """
+    if not os.path.isdir(dirname):
+        raise ValueError("There is no directory named '%s'", dirname)
+
+    model_file_name = dirname + "/__model__"
+    model = pickle.load(open(model_file_name, "r"))
+    program_desc_str = model["program_desc_str"]
+    feed_var_names = model["feed_var_names"]
+    fetch_var_names = model["fetch_var_names"]
+    program = Program.parse_from_string(program_desc_str)
+    load_persistables_if_exist(executor, dirname, program)
+    fetch_vars = [program.global_block().var(name) for name in fetch_var_names]
+
+    return [program, feed_var_names, fetch_vars]
diff --git a/python/paddle/v2/framework/tests/test_inference_model_io.py b/python/paddle/v2/framework/tests/test_inference_model_io.py
diff --git a/python/paddle/v2/framework/tests/test_operator_desc.py b/python/paddle/v2/framework/tests/test_operator_desc.py
diff --git a/python/paddle/v2/framework/tests/test_program.py b/python/paddle/v2/framework/tests/test_program.py

Original file line number	Diff line number	Diff line change
`@@ -49,6 +49,13 @@ ProgramDescBind::ProgramDescBind(const ProgramDescBind &o) {`
`49`	`49`	`}`
`50`	`50`	`}`
`51`	`51`
	`52`	`+ProgramDescBind::ProgramDescBind(const ProgramDesc &desc) {`
	`53`	`+ desc_ = desc;`
	`54`	`+ for (auto &block_desc : *desc_.mutable_blocks()) {`
	`55`	`+ blocks_.emplace_back(new BlockDescBind(this, &block_desc));`
	`56`	`+ }`
	`57`	`+}`
	`58`	`+`
`52`	`59`	`ProgramDescBind::ProgramDescBind(const std::string &binary_str) {`
`53`	`60`	`PADDLE_ENFORCE(desc_.ParseFromString(binary_str),`
`54`	`61`	`"Fail to parse program_desc from binary string.");`
Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ bool IsTarget(const OpDesc& op_desc) {`
`46`	`46`	`return false;`
`47`	`47`	`}`
`48`	`48`
`49`		`-void prune_impl(const ProgramDesc& input, ProgramDesc& output, int block_id) {`
	`49`	`+void prune_impl(const ProgramDesc& input, ProgramDesc* output, int block_id) {`
`50`	`50`	`// TODO(tonyyang-svail):`
`51`	`51`	`// - will change to use multiple blocks for RNN op and Cond Op`
`52`	`52`
`@@ -91,8 +91,8 @@ void prune_impl(const ProgramDesc& input, ProgramDesc& output, int block_id) {`
`91`	`91`	`// we reverse the should_run vector`
`92`	`92`	`std::reverse(should_run.begin(), should_run.end());`
`93`	`93`
`94`		`- output = input;`
`95`		`- auto* op_field = output.mutable_blocks(block_id)->mutable_ops();`
	`94`	`+ *output = input;`
	`95`	`+ auto* op_field = output->mutable_blocks(block_id)->mutable_ops();`
`96`	`96`	`op_field->Clear();`
`97`	`97`	`for (size_t i = 0; i < should_run.size(); ++i) {`
`98`	`98`	`if (should_run[i]) {`
`@@ -101,7 +101,8 @@ void prune_impl(const ProgramDesc& input, ProgramDesc& output, int block_id) {`
`101`	`101`	`}`
`102`	`102`	`}`
`103`	`103`
`104`		`-void Prune(const ProgramDesc& input, ProgramDesc& output) {`
	`104`	`+// TODO(fengjiayi): Prune() could be inplaced to avoid unnecessary copies`
	`105`	`+void Prune(const ProgramDesc& input, ProgramDesc* output) {`
`105`	`106`	`prune_impl(input, output, 0);`
`106`	`107`	`}`
`107`	`108`