[Dy2static] Refactor ProgramTranslator save_inference_model API (#24989)

* experimental refactoring, test=develop

* add TranslatedLayer & remove StaticModelRunner, test=develop

* revert tracedlayer change, test=develop

* fix test_mnist unittest error, test=develop

* add doc & examples, test=develop

* polish doc details, test=develop

* add imperative.jit module, test=develop

* change TranslatedLayer pos, test=develop

* adjust jit module import path, test=develop

* polish doc based review result

* add SaveLoadConfig.separate_params to save paraams separately

* add Layer.buffer support, test=develop

* polish doc details based review result, test=develop

* polish details baesd review comments, test=develop

* add empty str check for param, test=develop

* add unittests, test=develop

* polish details based review comment, test=develop

* remove blanks in comment, test=develop

* polish doc details, test=develop

* update imperative doc link, test=develop

* add api attr for load, test=develop

paddle/fluid/operators/run_program_op.h

@@ -49,13 +49,13 @@ static void CheckInputVarStatus(const Variable &var,
       var.IsType<LoDTensor>(), true,
       platform::errors::InvalidArgument(
           "The input variable %s of "
-          "RunProgram(Grad)Op(StaticModelRunner) holds "
+          "RunProgram(Grad)Op holds "
           "wrong type. Expect type is LoDTensor, but receive type is %s.",
           var_name, platform::demangle(framework::ToTypeName(var.Type()))));
   PADDLE_ENFORCE_EQ(
       var.Get<LoDTensor>().IsInitialized(), true,
       platform::errors::InvalidArgument("The tensor in input variable %s of "
-                                        "RunProgram(Grad)Op(StaticModelRunner) "
+                                        "RunProgram(Grad)Op "
                                         "is not initialized.",
                                         var_name));
 }
@@ -68,14 +68,14 @@ static void CheckOutputVarStatus(const Variable &src_var,
         src_var.IsType<LoDTensor>(), true,
         platform::errors::InvalidArgument(
             "The output variable %s get from "
-            "RunProgram(Grad)Op(StaticModelRunner)'s internal scope holds "
+            "RunProgram(Grad)Op's internal scope holds "
             "wrong type. Expect type is LoDTensor, but receive type is %s.",
             var_name,
             platform::demangle(framework::ToTypeName(src_var.Type()))));
     PADDLE_ENFORCE_EQ(src_var.Get<LoDTensor>().IsInitialized(), true,
                       platform::errors::InvalidArgument(
                           "The tensor in output variable %s get from "
-                          "RunProgram(Grad)Op(StaticModelRunner)'s internal "
+                          "RunProgram(Grad)Op's internal "
                           "scope is not initialized.",
                           var_name));
   } else if (dst_var.IsType<SelectedRows>()) {
@@ -83,20 +83,20 @@ static void CheckOutputVarStatus(const Variable &src_var,
         src_var.IsType<SelectedRows>(), true,
         platform::errors::InvalidArgument(
             "The output variable %s get from "
-            "RunProgram(Grad)Op(StaticModelRunner)'s internal scope holds "
+            "RunProgram(Grad)Op's internal scope holds "
             "wrong type. Expect type is SelectedRows, but receive type is %s.",
             var_name,
             platform::demangle(framework::ToTypeName(src_var.Type()))));
     PADDLE_ENFORCE_EQ(src_var.Get<SelectedRows>().value().IsInitialized(), true,
                       platform::errors::InvalidArgument(
                           "The tensor in output variable %s get from "
-                          "RunProgram(Grad)Op(StaticModelRunner)'s "
+                          "RunProgram(Grad)Op's "
                           "internal scope is not initialized.",
                           var_name));
 
   } else {
     PADDLE_THROW(platform::errors::InvalidArgument(
-        "The RunProgram(Grad)Op(StaticModelRunner) only support output "
+        "The RunProgram(Grad)Op only support output "
         "variable of type LoDTensor or SelectedRows, "
         "but received variable %s's type is %s",
         var_name, platform::demangle(framework::ToTypeName(dst_var.Type()))));
@@ -143,7 +143,7 @@ static void ShareVarsFromScope(const std::vector<Variable *> &vars,
     auto *var = scope->FindVar(var_names[i]);
     PADDLE_ENFORCE_NOT_NULL(
         var, platform::errors::NotFound("The output variable %s is not in "
-                                        "RunProgram(Grad)Op(StaticModelRunner)'"
+                                        "RunProgram(Grad)Op'"
                                         "s internal scope.",
                                         var_names[i]));
     CheckOutputVarStatus(*var, *vars[i], var_names[i]);

python/paddle/fluid/dygraph/__init__.py

@@ -44,6 +44,9 @@ from .backward_strategy import *
 from . import jit
 from .jit import *
 
+from . import io
+from .io import *
+
 from . import static_runner
 from .static_runner import StaticModelRunner
 
@@ -63,5 +66,6 @@ __all__ += checkpoint.__all__
 __all__ += learning_rate_scheduler.__all__
 __all__ += backward_strategy.__all__
 __all__ += jit.__all__
+__all__ += io.__all__
 __all__ += rnn.__all__
 __all__ += ['ProgramTranslator']

python/paddle/fluid/dygraph/dygraph_to_static/program_translator.py

@@ -36,6 +36,7 @@ from paddle.fluid.wrapped_decorator import signature_safe_contextmanager
 from paddle.fluid.dygraph.base import param_guard
 from paddle.fluid.data_feeder import check_type
 from paddle.fluid.dygraph.dygraph_to_static.partial_program import partial_program_from
+from paddle.fluid.annotations import deprecated
 
 __all__ = ['ProgramTranslator', 'convert_to_static']
 
@@ -125,6 +126,9 @@ class FunctionSpec(object):
         self._args = args
         self._kwargs = kwargs
 
+        dyfunc = getattr(func, '__wrapped__', func)
+        self._dyfunc_code = inspect.getsource(dyfunc)
+
     def is_method(self):
         return self._args and isinstance(self._args[0], layers.Layer)
 
@@ -198,7 +202,9 @@ class FunctionSpec(object):
         # Note: if dygraph function is a method of class,
         # consider instance info as hash key.
         if self.is_method():
-            return self._dyfunc, self._args[0]
+            # NOTE: we can use Layer's (instance + function code) as hash key.
+            # An instance will not hold two identical methods 
+            return self._dyfunc_code, self._args[0]
         else:
             return self._dyfunc
 
@@ -312,6 +318,17 @@ class ProgramCache(object):
             self._caches[item] = self._build_once(item)
         return self._caches[item]
 
+    def get_program(self, item):
+        if not isinstance(item, FunctionSpec):
+            raise ValueError(
+                "Input item's type should be FunctionSpec, but received %s" %
+                type(item))
+        if item not in self._caches:
+            raise RuntimeError(
+                "Failed to find program for input item, please decorate input function by `@declarative`."
+            )
+        return self._caches[item]
+
     def last(self):
         assert len(
             self._caches) >= 1, "No valid cached program in ProgramCache."
@@ -633,6 +650,7 @@ class ProgramTranslator(object):
         source_code = ast_to_source_code(root_wrapper.node)
         return source_code
 
+    @deprecated(since='2.0', instead="paddle.imperative.jit.save")
     @switch_to_static_graph
     def save_inference_model(self, dirname, feed=None, fetch=None):
         """

python/paddle/fluid/tests/unittests/dygraph_to_static/test_mnist.py

@@ -26,6 +26,7 @@ from paddle.fluid.dygraph import to_variable
 from paddle.fluid.dygraph.nn import Conv2D, Linear, Pool2D
 from paddle.fluid.optimizer import AdamOptimizer
 from paddle.fluid.dygraph.jit import declarative
+from paddle.fluid.dygraph.io import VARIABLE_FILENAME
 from paddle.fluid.dygraph.dygraph_to_static import ProgramTranslator
 
 SEED = 2020
@@ -201,6 +202,9 @@ class TestMNISTWithDeclarative(TestMNIST):
                         self.check_save_inference_model([dy_x_data, y_data],
                                                         prog_trans, to_static,
                                                         prediction)
+                        # new save load check
+                        self.check_jit_save_load(mnist, [dy_x_data], [img],
+                                                 to_static, prediction)
                         break
         return loss_data
 
@@ -224,6 +228,45 @@ class TestMNISTWithDeclarative(TestMNIST):
 
         return np.array(results[0])
 
+    def check_jit_save_load(self, model, inputs, input_spec, to_static, gt_out):
+        if to_static:
+            infer_model_path = "./test_mnist_inference_model_by_jit_save"
+            configs = fluid.dygraph.jit.SaveLoadConfig()
+            configs.output_spec = [gt_out]
+            fluid.dygraph.jit.save(
+                layer=model,
+                model_path=infer_model_path,
+                input_spec=input_spec,
+                configs=configs)
+            # load in static mode
+            static_infer_out = self.jit_load_and_run_inference_static(
+                infer_model_path, inputs)
+            self.assertTrue(np.allclose(gt_out.numpy(), static_infer_out))
+            # load in dygraph mode
+            dygraph_infer_out = self.jit_load_and_run_inference_dygraph(
+                infer_model_path, inputs)
+            self.assertTrue(np.allclose(gt_out.numpy(), dygraph_infer_out))
+
+    @switch_to_static_graph
+    def jit_load_and_run_inference_static(self, model_path, inputs):
+        exe = fluid.Executor(self.place)
+        [inference_program, feed_target_names,
+         fetch_targets] = fluid.io.load_inference_model(
+             dirname=model_path,
+             executor=exe,
+             params_filename=VARIABLE_FILENAME)
+        assert len(inputs) == len(feed_target_names)
+        results = exe.run(inference_program,
+                          feed=dict(zip(feed_target_names, inputs)),
+                          fetch_list=fetch_targets)
+
+        return np.array(results[0])
+
+    def jit_load_and_run_inference_dygraph(self, model_path, inputs):
+        infer_net = fluid.dygraph.jit.load(model_path)
+        pred = infer_net(inputs[0])
+        return pred.numpy()
+
 
 if __name__ == "__main__":
     unittest.main()

python/paddle/fluid/tests/unittests/test_jit_save_load.py

+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+from __future__ import print_function
+
+import unittest
+import numpy as np
+
+import paddle
+import paddle.fluid as fluid
+from paddle.fluid.dygraph import Linear
+from paddle.fluid.dygraph import declarative
+
+BATCH_SIZE = 32
+BATCH_NUM = 20
+SEED = 10
+
+
+def random_batch_reader():
+    def _get_random_images_and_labels(image_shape, label_shape):
+        np.random.seed(SEED)
+        image = np.random.random(size=image_shape).astype('float32')
+        label = np.random.random(size=label_shape).astype('int64')
+        return image, label
+
+    def __reader__():
+        for _ in range(BATCH_NUM):
+            batch_image, batch_label = _get_random_images_and_labels(
+                [BATCH_SIZE, 784], [BATCH_SIZE, 1])
+            yield batch_image, batch_label
+
+    return __reader__
+
+
+class LinearNet(fluid.dygraph.Layer):
+    def __init__(self, in_size, out_size):
+        super(LinearNet, self).__init__()
+        self._linear = Linear(in_size, out_size)
+
+    @declarative
+    def forward(self, x):
+        return self._linear(x)
+
+
+class LinearNetNotDeclarative(fluid.dygraph.Layer):
+    def __init__(self, in_size, out_size):
+        super(LinearNetNotDeclarative, self).__init__()
+        self._linear = Linear(in_size, out_size)
+
+    def forward(self, x):
+        return self._linear(x)
+
+
+class LinearNetReturnLoss(fluid.dygraph.Layer):
+    def __init__(self, in_size, out_size):
+        super(LinearNetReturnLoss, self).__init__()
+        self._linear = Linear(in_size, out_size)
+
+    @declarative
+    def forward(self, x):
+        y = self._linear(x)
+        z = self._linear(y)
+        loss = fluid.layers.mean(z)
+        return z, loss
+
+
+def train(layer):
+    # create optimizer
+    adam = fluid.optimizer.AdamOptimizer(
+        learning_rate=0.1, parameter_list=layer.parameters())
+    # create data loader
+    train_loader = fluid.io.DataLoader.from_generator(capacity=5)
+    train_loader.set_batch_generator(random_batch_reader())
+    # train
+    for data in train_loader():
+        img, label = data
+        label.stop_gradient = True
+
+        cost = layer(img)
+
+        loss = fluid.layers.cross_entropy(cost, label)
+        avg_loss = fluid.layers.mean(loss)
+
+        avg_loss.backward()
+        adam.minimize(avg_loss)
+        layer.clear_gradients()
+    return [img], layer, avg_loss
+
+
+def infer(layer):
+    x = fluid.dygraph.to_variable(np.random.random((1, 784)).astype('float32'))
+    return layer(x)
+
+
+class TestJitSaveLoad(unittest.TestCase):
+    def setUp(self):
+        self.model_path = "model.test_jit_save_load"
+        # enable dygraph mode
+        fluid.enable_dygraph()
+        # config seed
+        fluid.default_main_program().random_seed = SEED
+
+    def train_and_save_model(self):
+        layer = LinearNet(784, 1)
+        example_inputs, layer, _ = train(layer)
+        fluid.dygraph.jit.save(
+            layer=layer, model_path=self.model_path, input_spec=example_inputs)
+        return layer
+
+    def test_save(self):
+        # train and save model
+        self.train_and_save_model()
+
+    def test_load_infernece(self):
+        # train and save model
+        train_layer = self.train_and_save_model()
+        # load model
+        infer_layer = fluid.dygraph.jit.load(self.model_path)
+        train_layer.eval()
+        # inference & compare
+        x = fluid.dygraph.to_variable(
+            np.random.random((1, 784)).astype('float32'))
+        self.assertTrue(
+            np.array_equal(train_layer(x).numpy(), infer_layer(x).numpy()))
+
+    def test_load_finetune(self):
+        # train and save model
+        train_layer = self.train_and_save_model()
+        # load model
+        load_train_layer = fluid.dygraph.jit.load(self.model_path)
+        load_train_layer.train()
+        # train & compare
+        _, _, train_loss = train(train_layer)
+        _, _, load_train_loss = train(load_train_layer)
+        self.assertTrue(
+            np.array_equal(train_loss.numpy(), load_train_loss.numpy()))
+
+    def test_save_get_program_failed(self):
+        layer = LinearNetNotDeclarative(784, 1)
+        example_inputs, layer, _ = train(layer)
+        with self.assertRaises(RuntimeError):
+            fluid.dygraph.jit.save(
+                layer=layer,
+                model_path=self.model_path,
+                input_spec=example_inputs)
+
+
+class TestJitSaveLoadConfig(unittest.TestCase):
+    def setUp(self):
+        # enable dygraph mode
+        fluid.enable_dygraph()
+        # config seed
+        fluid.default_main_program().random_seed = SEED
+
+    def basic_save_load(self, layer, model_path, configs):
+        # 1. train & save
+        example_inputs, train_layer, _ = train(layer)
+        fluid.dygraph.jit.save(
+            layer=train_layer,
+            model_path=model_path,
+            input_spec=example_inputs,
+            configs=configs)
+        # 2. load 
+        infer_layer = fluid.dygraph.jit.load(model_path, configs=configs)
+        train_layer.eval()
+        # 3. inference & compare
+        x = fluid.dygraph.to_variable(
+            np.random.random((1, 784)).astype('float32'))
+        self.assertTrue(
+            np.array_equal(train_layer(x).numpy(), infer_layer(x).numpy()))
+
+    def test_model_filename(self):
+        layer = LinearNet(784, 1)
+        model_path = "model.save_load_config.output_spec"
+        configs = fluid.dygraph.jit.SaveLoadConfig()
+        configs.model_filename = "__simplenet__"
+        self.basic_save_load(layer, model_path, configs)
+
+    def test_params_filename(self):
+        layer = LinearNet(784, 1)
+        model_path = "model.save_load_config.params_filename"
+        configs = fluid.dygraph.jit.SaveLoadConfig()
+        configs.params_filename = "__params__"
+        self.basic_save_load(layer, model_path, configs)
+
+    def test_separate_params(self):
+        layer = LinearNet(784, 1)
+        model_path = "model.save_load_config.separate_params"
+        configs = fluid.dygraph.jit.SaveLoadConfig()
+        configs.separate_params = True
+        self.basic_save_load(layer, model_path, configs)
+
+    def test_output_spec(self):
+        train_layer = LinearNetReturnLoss(8, 8)
+        adam = fluid.optimizer.AdamOptimizer(
+            learning_rate=0.1, parameter_list=train_layer.parameters())
+        x = fluid.dygraph.to_variable(
+            np.random.random((4, 8)).astype('float32'))
+        for i in range(10):
+            out, loss = train_layer(x)
+            loss.backward()
+            adam.minimize(loss)
+            train_layer.clear_gradients()
+
+        model_path = "model.save_load_config.output_spec"
+        configs = fluid.dygraph.jit.SaveLoadConfig()
+        configs.output_spec = [out]
+        fluid.dygraph.jit.save(
+            layer=train_layer,
+            model_path=model_path,
+            input_spec=[x],
+            configs=configs)
+
+        train_layer.eval()
+        infer_layer = fluid.dygraph.jit.load(model_path, configs=configs)
+        x = fluid.dygraph.to_variable(
+            np.random.random((4, 8)).astype('float32'))
+        self.assertTrue(
+            np.array_equal(train_layer(x)[0].numpy(), infer_layer(x).numpy()))
+
+
+if __name__ == '__main__':
+    unittest.main()

python/paddle/imperative/__init__.py

@@ -16,7 +16,7 @@
 __all__ = [
     'BackwardStrategy', 'enabled', 'grad', 'guard', 'LayerList', 'load', 'save',
     'prepare_context', 'to_variable', 'TracedLayer', 'no_grad', 'ParallelEnv',
-    'ProgramTranslator', 'declarative', 'DataParallel'
+    'ProgramTranslator', 'declarative', 'DataParallel', 'TranslatedLayer', 'jit'
 ]
 
 __all__ += [
@@ -31,6 +31,7 @@ from ..fluid.dygraph.checkpoint import save_dygraph as save
 from ..fluid.dygraph.parallel import prepare_context, ParallelEnv, DataParallel
 from ..fluid.dygraph.jit import TracedLayer, declarative
 from ..fluid.dygraph import ProgramTranslator
+from . import jit
 
 from ..fluid.dygraph.learning_rate_scheduler import NoamDecay, PiecewiseDecay, NaturalExpDecay, ExponentialDecay, \
         InverseTimeDecay, PolynomialDecay, CosineDecay

python/paddle/imperative/jit/__init__.py

+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+from ...fluid.dygraph.jit import save, load, SaveLoadConfig
+from ...fluid.dygraph.io import TranslatedLayer
+
+__all__ = ['save', 'load', 'SaveLoadConfig']

python/setup.py.in

@@ -202,6 +202,7 @@ packages=['paddle',
           'paddle.nn.initializer',
           'paddle.metric',
           'paddle.imperative',
+          'paddle.imperative.jit',
           'paddle.tensor',
           ]