PaddlePaddle
diff --git a/‎python/paddle/jit/api.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎python/paddle/jit/api.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎python/paddle/jit/pir_translated_layer.py
Copy file name to clipboard
+125-25Lines changed: 125 additions & 25 deletions b/‎python/paddle/jit/pir_translated_layer.py
Copy file name to clipboard
+125-25Lines changed: 125 additions & 25 deletions
diff --git a/‎python/paddle/pir/core.py
Copy file name to clipboard
+6-1Lines changed: 6 additions & 1 deletion b/‎python/paddle/pir/core.py
Copy file name to clipboard
+6-1Lines changed: 6 additions & 1 deletion
diff --git a/‎python/paddle/static/io.py
Copy file name to clipboard
+2-4Lines changed: 2 additions & 4 deletions b/‎python/paddle/static/io.py
Copy file name to clipboard
+2-4Lines changed: 2 additions & 4 deletions
@@ -1157,11 +1157,14 @@ def save(layer, path, input_spec=None, **configs):
         with dygraph.guard():
             if use_pir_api():
                 for tensor, value in zip(*concrete_program.parameters):
+                    if not value.persistable:
+                        continue
                     param_or_buffer_tensor = scope.var(value.name).get_tensor()
                     src_tensor = (
                         state_var_dict[tensor.name].value().get_tensor()
                     )
                     param_or_buffer_tensor._share_data_with(src_tensor)
+
             else:
                 for param_or_buffer in concrete_program.parameters:
                     # share to scope
 
@@ -53,18 +53,22 @@ def _generate_unique_var_name(prefix):
     return unique_name.generate(prefix)
 
 
-def _get_pir_persistable_var_names(program):
+from paddle.static.pir_io import get_pir_parameters
+
+
+def _get_pir_parameters_var_names(program):
     persistable_vars = []
     persistable_names = []
     rename_new_old_dict = {}
-    for block in program.blocks:
-        for var in block.all_parameters():
-            if var.persistable:
-                persistable_vars.append(var)
-                origin_name = var.name
-                var.name = _generate_unique_var_name(var.name)
-                rename_new_old_dict[var.name] = origin_name
-                persistable_names.append(var.name)
+    param, opt = get_pir_parameters(program)
+    vars = param + opt
+    for var in vars:
+        persistable_vars.append(var)
+        origin_name = var.name
+        var.name = _generate_unique_var_name(var.name)
+        rename_new_old_dict[var.name] = origin_name
+        persistable_names.append(var.name)
+
     return (
         persistable_vars,
         persistable_names,
@@ -79,8 +83,8 @@ def __init__(self, program, trainable):
         # input, output, persistable,
         self._input_vars = []
         self._output_vars = []
-        self._persistable_vars = []
-        self._persistable_names = []
+        self._parameter_vars = []
+        self._parameter_names = []
 
         self.support_train = trainable
         # append suffix var name dict
@@ -91,11 +95,10 @@ def __init__(self, program, trainable):
 
     def _preprocess(self):
         (
-            self._persistable_vars,
-            self._persistable_names,
+            self._parameter_vars,
+            self._parameter_names,
             self._suffix_varname_dict,
-        ) = _get_pir_persistable_var_names(self._infer_program)
-
+        ) = _get_pir_parameters_var_names(self._infer_program)
         block = self._infer_program.global_block()
         for op in block.ops:
             if op.name() == 'pd_op.data':
@@ -135,11 +138,11 @@ def output_vars(self):
 
     @property
     def persistable_names(self):
-        return self._persistable_names
+        return self._parameter_names
 
     @property
     def persistable_vars(self):
-        return self._persistable_vars
+        return self._parameter_vars
 
 
 # [ PirTranslatedLayer : Run program in dygraph mode ]
@@ -173,10 +176,11 @@ def persistable_vars(self):
 # which need get var info from program desc
 
 
-def _load_pir_persistable_vars(model_path, program_holder, params_filename):
+def _load_pir_parameter_vars(model_path, program_holder, params_filename):
     # construct var dict
     load_var_dict = {}
     load_var_list = []
+    other_var_dict = {}
     load_densetensor_list = []
     persistable_var = program_holder.persistable_vars
     persistable_var_name = program_holder.persistable_names
@@ -193,6 +197,12 @@ def _load_pir_persistable_vars(model_path, program_holder, params_filename):
                 name=var.name,
                 persistable=True,
             )
+
+            new_var.stop_gradient = var.stop_gradient
+            load_var_dict[name] = new_var
+            load_var_list.append(new_var)
+            load_densetensor_list.append(new_var.get_tensor())
+
         else:
             new_var = core.eager.Tensor(
                 dtype=datatype_to_vartype[var.dtype],
@@ -202,11 +212,7 @@ def _load_pir_persistable_vars(model_path, program_holder, params_filename):
                 place=framework._current_expected_place(),
                 persistable=False,
             )
-
-        new_var.stop_gradient = var.stop_gradient
-        load_var_dict[name] = new_var
-        load_var_list.append(new_var)
-        load_densetensor_list.append(new_var.get_tensor())
+            other_var_dict[name] = new_var
 
     # load all vars
     assert params_filename is not None, "params_filename should not be None."
@@ -226,6 +232,7 @@ def _load_pir_persistable_vars(model_path, program_holder, params_filename):
             % var_file_path
         )
 
+    load_var_dict.update(other_var_dict)
     return load_var_dict
 
 
@@ -285,7 +292,7 @@ def _construct_params_and_buffers(model_path, programs, params_filename=None):
         # When saving XX, there is only '*.pdmodel'
         return {}
     else:
-        var_dict = _load_pir_persistable_vars(
+        var_dict = _load_pir_parameter_vars(
             model_path, programs['forward'], params_filename
         )
         model_name = params_filename[: -len(INFER_PARAMS_SUFFIX)]
@@ -305,7 +312,7 @@ def _construct_params_and_buffers(model_path, programs, params_filename=None):
                 continue
 
             var_dict.update(
-                _load_pir_persistable_vars(
+                _load_pir_parameter_vars(
                     model_path, programs[func_name], file_name
                 )
             )
@@ -615,6 +622,99 @@ def eval(self):
         self._is_test = True
         self.training = False
 
+    def program(self, method_name='forward'):
+        """
+        Gets translated program of specified method.
+
+        Args:
+            - method_name (string): method name corresponding to the program
+                to be obtained. Default: 'forward'.
+
+        Returns:
+            Program
+
+        Examples:
+            .. code-block:: python
+
+                >>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
+                >>> import numpy as np
+                >>> import paddle
+                >>> from paddle import nn
+                >>> import paddle.optimizer as opt
+
+                >>> BATCH_SIZE = 16
+                >>> BATCH_NUM = 4
+                >>> EPOCH_NUM = 4
+
+                >>> IMAGE_SIZE = 784
+                >>> CLASS_NUM = 10
+
+                >>> # define a random dataset
+                >>> class RandomDataset(paddle.io.Dataset):
+                ...     def __init__(self, num_samples):
+                ...         self.num_samples = num_samples
+                ...
+                ...     def __getitem__(self, idx):
+                ...         image = np.random.random([IMAGE_SIZE]).astype('float32')
+                ...         label = np.random.randint(0, CLASS_NUM - 1, (1, )).astype('int64')
+                ...         return image, label
+                ...
+                ...     def __len__(self):
+                ...         return self.num_samples
+                ...
+                >>> class LinearNet(nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
+                ...
+                ...     @paddle.jit.to_static
+                ...     def forward(self, x):
+                ...         return self._linear(x)
+                ...
+                >>> def train(layer, loader, loss_fn, opt):
+                ...     for epoch_id in range(EPOCH_NUM):
+                ...         for batch_id, (image, label) in enumerate(loader()):
+                ...             out = layer(image)
+                ...             loss = loss_fn(out, label)
+                ...             loss.backward()
+                ...             opt.step()
+                ...             opt.clear_grad()
+                ...             print("Epoch {} batch {}: loss = {}".format(
+                ...                 epoch_id, batch_id, np.mean(loss.numpy())))
+                ...
+                >>> # create network
+                >>> layer = LinearNet()
+                >>> loss_fn = nn.CrossEntropyLoss()
+                >>> adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())
+                >>> # create data loader
+                >>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
+                >>> loader = paddle.io.DataLoader(dataset,
+                ...     batch_size=BATCH_SIZE,
+                ...     shuffle=True,
+                ...     drop_last=True,
+                ...     num_workers=2
+                ... )
+                >>> # train
+                >>> train(layer, loader, loss_fn, adam)
+
+                >>> # save
+                >>> model_path = "linear.example.model"
+                >>> paddle.jit.save(layer, model_path)
+
+                >>> # load
+                >>> translated_layer = paddle.jit.load(model_path)
+
+                >>> # get program
+                >>> program = translated_layer.program()
+        """
+        # 1. get program holder
+        program_holder = self._get_program_holder(method_name)
+
+        # 2. get inference program desc
+        program = program_holder.infer_program
+
+        return program
+
     def _get_program_holder(self, method_name='forward'):
         program_holder = self._program_holder_dict.get(method_name, None)
         if program_holder is None:
 
@@ -419,9 +419,14 @@ def _convert_into_value(tensor):
     )
 
     if isinstance(tensor, paddle.Tensor):
-        return _global_parameter_recorder.get(
+        value = _global_parameter_recorder.get(
             paddle.pir.core.default_main_program(), tensor
         )
+        NON_PERSISTABLE_VAR_NAME_SUFFIX = "__non_persistable"
+        if tensor.name.endswith(NON_PERSISTABLE_VAR_NAME_SUFFIX):
+            value.persistable = False
+        return value
+
     return tensor
 
 
 
@@ -1038,7 +1038,7 @@ def save_vars(
 
     """
     if in_pir_mode():
-        return save_vars_pir(dirname, main_program, vars, predicate, filename)
+        return save_vars_pir(dirname, main_program, vars, filename)
 
     save_to_memory = False
     if dirname is None and filename is None:
@@ -1206,9 +1206,7 @@ def load_vars(
 
     """
     if in_pir_mode():
-        return load_vars_pir(
-            executor, dirname, main_program, vars, predicate, filename
-        )
+        return load_vars_pir(executor, dirname, main_program, vars, filename)
 
     vars_from_memory = False
     if dirname is not None:
Original file line number	Diff line number	Diff line change
`@@ -419,9 +419,14 @@ def _convert_into_value(tensor):`
`419`	`419`	`)`
`420`	`420`
`421`	`421`	`if isinstance(tensor, paddle.Tensor):`
`422`		`- return _global_parameter_recorder.get(`
	`422`	`+ value = _global_parameter_recorder.get(`
`423`	`423`	`paddle.pir.core.default_main_program(), tensor`
`424`	`424`	`)`
	`425`	`+ NON_PERSISTABLE_VAR_NAME_SUFFIX = "__non_persistable"`
	`426`	`+ if tensor.name.endswith(NON_PERSISTABLE_VAR_NAME_SUFFIX):`
	`427`	`+ value.persistable = False`
	`428`	`+ return value`
	`429`	`+`
`425`	`430`	`return tensor`
`426`	`431`
`427`	`432`