From 2d1fbe9fd8c53ff19c562f743b73edbcd85ac01d Mon Sep 17 00:00:00 2001 From: wangruting Date: Fri, 26 Apr 2024 07:49:46 +0000 Subject: [PATCH 01/15] add jit load.train --- .../serialize_deserialize/include/interface.h | 7 +- .../serialize_deserialize/src/interface.cc | 12 +- python/paddle/jit/pir_translated_layer.py | 131 +++++++++++------- .../legacy_test/test_io_save_load.py | 14 +- 4 files changed, 102 insertions(+), 62 deletions(-) diff --git a/paddle/fluid/pir/serialize_deserialize/include/interface.h b/paddle/fluid/pir/serialize_deserialize/include/interface.h index 3c4107338aa925..05bc08943bc65e 100644 --- a/paddle/fluid/pir/serialize_deserialize/include/interface.h +++ b/paddle/fluid/pir/serialize_deserialize/include/interface.h @@ -53,13 +53,14 @@ void WriteModule(const pir::Program& program, * deserilize program will be stored. * @param[in] pir_version The current version of the PIR program format. * - * @return Void. The function modifies the 'program' object to contain the data - * read from the file. + * @return bool. The function modifies the 'program' object to contain the data + * read from the file. return bool indicates whether the program can use to + * funtune. * * @note If 'pir_version' is larger than the version of file, will trigger * version compatibility modification rule. */ -void ReadModule(const std::string& file_path, +bool ReadModule(const std::string& file_path, pir::Program* program, const uint64_t& pir_version); diff --git a/paddle/fluid/pir/serialize_deserialize/src/interface.cc b/paddle/fluid/pir/serialize_deserialize/src/interface.cc index 7a55c478c8b1b7..26d5810f0c9157 100644 --- a/paddle/fluid/pir/serialize_deserialize/src/interface.cc +++ b/paddle/fluid/pir/serialize_deserialize/src/interface.cc @@ -23,6 +23,7 @@ namespace pir { #define BASE_CODE "base_code" #define MAGIC "magic" #define PIRVERSION "version" +#define TRAINABLE "trainable" #define PIR "pir" void WriteModule(const pir::Program& program, const std::string& file_path, @@ -41,7 +42,8 @@ void WriteModule(const pir::Program& program, // write base code Json total; - total[BASE_CODE] = {{MAGIC, PIR}, {PIRVERSION, pir_version}}; + total[BASE_CODE] = { + {MAGIC, PIR}, {PIRVERSION, pir_version}, {TRAINABLE, trainable}}; ProgramWriter writer(pir_version, trainable); // write program @@ -63,7 +65,7 @@ void WriteModule(const pir::Program& program, fout.close(); } -void ReadModule(const std::string& file_path, +bool ReadModule(const std::string& file_path, pir::Program* program, const uint64_t& pir_version) { std::ifstream f(file_path); @@ -83,6 +85,12 @@ void ReadModule(const std::string& file_path, ProgramReader reader(pir_version); reader.RecoverProgram(&(data[PROGRAM]), program); + + if (data[BASE_CODE].contains(TRAINABLE)) { + return data[BASE_CODE][TRAINABLE].get(); + } else { + return false; + } } } // namespace pir diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py index 1f506778216331..6d2b2dcb8485eb 100644 --- a/python/paddle/jit/pir_translated_layer.py +++ b/python/paddle/jit/pir_translated_layer.py @@ -36,9 +36,11 @@ def _load_pir_program(model_file_path): program = paddle.static.Program() - paddle.base.core.deserialize_pir_program(model_file_path, program, 1) + trainable = paddle.base.core.deserialize_pir_program( + model_file_path, program, 1 + ) - return program + return program, trainable @switch_to_static_graph @@ -58,7 +60,7 @@ def _get_pir_persistable_var_names(program): class _PirProgramHolder: - def __init__(self, program): + def __init__(self, program, trainable): super().__init__() # input, output, persistable, @@ -70,9 +72,15 @@ def __init__(self, program): # execution scope self._inner_scope = core.Scope() + self.support_train = trainable self._infer_program = program + self._forward_program = program + self._backward_program = paddle.pir.Program() + self._preprocess() + self.prepare_train_program = False + def _preprocess(self): ( self._persistable_vars, @@ -128,10 +136,22 @@ def persistable_vars(self): def scope(self): return self._inner_scope + @property + def forward_program(self): + return self._forward_program + + @property + def backward_program(self): + return self._backward_program + + @property + def train_attr(self): + return self._train_attr -# [ PirTranslatedLayer : Run program in imperative mode ] + +# [ PirTranslatedLayer : Run program in dygraph mode ] # -# DESIGN IDEA: using an special operator `RunProgram`, execute program inside operator. +# DESIGN IDEA: using an special operator `PirRunProgram`, execute program inside operator. # # Op's Inputs: # - the input variable of the user feed @@ -228,9 +248,9 @@ def _construct_program_holders(model_path, model_filename=None): continue else: continue - + program, trainable = _load_pir_program(model_file_path) program_holder_dict[func_name] = _PirProgramHolder( - _load_pir_program(model_file_path) + program, trainable ) else: for _, _, file_names in os.walk(model_path): @@ -242,8 +262,9 @@ def _construct_program_holders(model_path, model_filename=None): method_name = 'forward' else: method_name.replace('model', '') + program, trainable = _load_pir_program(model_file_path) program_holder_dict[func_name] = _PirProgramHolder( - _load_pir_program(model_file_path) + program, trainable ) return program_holder_dict @@ -329,56 +350,60 @@ def _run_dygraph(instance, input, program_holder): tmp_scope_vec = [program_holder.scope] # 2. run program by op + if instance._is_test: + attrs = [ + 'forward_global_block', + program_holder._infer_program.global_block(), + 'backward_global_block', + program_holder.backward_program.global_block(), + 'is_test', + instance._is_test, + 'program_id', + paddle.utils._hash_with_id(program_holder._infer_program, instance), + 'fx', + program_holder.input_vars, + 'fo', + program_holder.output_vars, + 'fp', + program_holder.persistable_vars, + 'fm', + [], + 'no_need_buffers', + [], + ] - forward_program = ( - program_holder._infer_program - if instance._is_test - else program_holder.forward_program - ) + _legacy_C_ops.pir_run_program( + _valid_vars(input_tensors), + _valid_vars(persistable_tensors), + _valid_vars(output_tensors), + tmp_scope_vec, + None, + *attrs, + ) - backward_program = ( - paddle.pir.Program() - if instance._is_test - else program_holder.backward_program - ) + outs = output_tensors + if len(output_tensors) == 1: + outs = output_tensors[0] + return outs - attrs = [ - 'forward_global_block', - forward_program.global_block(), - 'backward_global_block', - backward_program.global_block(), - 'is_test', - instance._is_test, - 'program_id', - paddle.utils._hash_with_id(forward_program, instance), - 'fx', - program_holder.input_vars, - 'fo', - program_holder.output_vars, - 'fp', - program_holder.persistable_vars, - 'fm', - [], - 'no_need_buffers', - [], - ] - - _legacy_C_ops.pir_run_program( - _valid_vars(input_tensors), - _valid_vars(persistable_tensors), - _valid_vars(output_tensors), - tmp_scope_vec, - None, - *attrs, - ) + else: + from paddle.jit.dy2static.pir_partial_program import PartialProgramLayer + + inputs = program_holder._input_vars + outputs = program_holder._output_vars + parameters = (persistable_tensors, program_holder._persistable_vars) + + layer = PartialProgramLayer( + program_holder._infer_program, + inputs, + outputs, + parameters, + ) - outs = output_tensors - if len(output_tensors) == 1: - outs = output_tensors[0] - return outs + return layer(input_tensors) -def _run_static_graph(input, program_holder, trace_program): +def _run_static_graph(program_holder, trace_program): paddle.base.framework.switch_main_program(trace_program) return program_holder.output_vars @@ -594,7 +619,7 @@ def __i_m_p_l__(self, *input): return _run_dygraph(self, input, program_holder) else: return _run_static_graph( - input, program_holder, program_holder.infer_program + program_holder, program_holder.infer_program ) __i_m_p_l__.__name__ = method_name diff --git a/test/deprecated/legacy_test/test_io_save_load.py b/test/deprecated/legacy_test/test_io_save_load.py index 01c1dcb6c7c3ee..075e2856781880 100644 --- a/test/deprecated/legacy_test/test_io_save_load.py +++ b/test/deprecated/legacy_test/test_io_save_load.py @@ -107,12 +107,18 @@ def test_when_train_with_no_grad(self): save_path = os.path.join(self.temp_dir.name, 'train_with_no_grad') paddle.jit.save(net, save_path) - net = paddle.jit.load(save_path) - net.eval() + net1 = paddle.jit.load(save_path) + net1.eval() with paddle.no_grad(): - out2 = net(x) - self.assertEqual(out, out2) + out1 = net1(x) + self.assertEqual(out, out1) + + net2 = paddle.jit.load(save_path) + net2.train() + + out2 = net2(x) + self.assertEqual(out, out2) if __name__ == '__main__': From b572a8aa6073eb538196b6315f5652c39ac7e363 Mon Sep 17 00:00:00 2001 From: wangruting Date: Fri, 26 Apr 2024 10:10:19 +0000 Subject: [PATCH 02/15] modify backward program lost --- python/paddle/jit/pir_translated_layer.py | 4 +- .../legacy_test/test_io_save_load.py | 47 ++++++++++--------- 2 files changed, 28 insertions(+), 23 deletions(-) diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py index 6d2b2dcb8485eb..a411bd990214a6 100644 --- a/python/paddle/jit/pir_translated_layer.py +++ b/python/paddle/jit/pir_translated_layer.py @@ -399,8 +399,8 @@ def _run_dygraph(instance, input, program_holder): outputs, parameters, ) - - return layer(input_tensors) + instance.layer = layer + return instance.layer(input_tensors) def _run_static_graph(program_holder, trace_program): diff --git a/test/deprecated/legacy_test/test_io_save_load.py b/test/deprecated/legacy_test/test_io_save_load.py index 075e2856781880..3375435b76f34b 100644 --- a/test/deprecated/legacy_test/test_io_save_load.py +++ b/test/deprecated/legacy_test/test_io_save_load.py @@ -97,28 +97,33 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() - @test_with_pir_api + # @test_with_pir_api def test_when_train_with_no_grad(self): - paddle.disable_static() - net = paddle.nn.Linear(1024, 1) - net = paddle.jit.to_static(net) - x = paddle.rand([1024], 'float32') - out = net(x) - save_path = os.path.join(self.temp_dir.name, 'train_with_no_grad') - - paddle.jit.save(net, save_path) - net1 = paddle.jit.load(save_path) - net1.eval() - - with paddle.no_grad(): - out1 = net1(x) - self.assertEqual(out, out1) - - net2 = paddle.jit.load(save_path) - net2.train() - - out2 = net2(x) - self.assertEqual(out, out2) + with paddle.pir_utils.IrGuard(): + paddle.disable_static() + net = paddle.nn.Linear(1024, 1) + net = paddle.jit.to_static(net, full_graph=True) + x = paddle.rand([1024], 'float32') + x.stop_gradient = False + out = net(x) + out.backward() + x_grad = x.grad.mean() + save_path = os.path.join(self.temp_dir.name, 'train_with_no_grad') + + paddle.jit.save(net, save_path) + net1 = paddle.jit.load(save_path) + net1.eval() + + with paddle.no_grad(): + out1 = net1(x) + self.assertEqual(out, out1) + + net2 = paddle.jit.load(save_path) + net2.train() + out2 = net2(x) + out2.backward() + self.assertEqual(out, out2) + x_grad2 = x.grad.mean() if __name__ == '__main__': From a894be663142acdfdccda1676a4e513d3fcf1cf5 Mon Sep 17 00:00:00 2001 From: wangruting Date: Fri, 26 Apr 2024 10:21:52 +0000 Subject: [PATCH 03/15] modify --- python/paddle/jit/pir_translated_layer.py | 26 ++++++----------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py index a411bd990214a6..c162606e77e37c 100644 --- a/python/paddle/jit/pir_translated_layer.py +++ b/python/paddle/jit/pir_translated_layer.py @@ -74,8 +74,6 @@ def __init__(self, program, trainable): self.support_train = trainable self._infer_program = program - self._forward_program = program - self._backward_program = paddle.pir.Program() self._preprocess() @@ -136,18 +134,6 @@ def persistable_vars(self): def scope(self): return self._inner_scope - @property - def forward_program(self): - return self._forward_program - - @property - def backward_program(self): - return self._backward_program - - @property - def train_attr(self): - return self._train_attr - # [ PirTranslatedLayer : Run program in dygraph mode ] # @@ -353,9 +339,9 @@ def _run_dygraph(instance, input, program_holder): if instance._is_test: attrs = [ 'forward_global_block', - program_holder._infer_program.global_block(), + program_holder.infer_program.global_block(), 'backward_global_block', - program_holder.backward_program.global_block(), + paddle.pir.Program().global_block(), 'is_test', instance._is_test, 'program_id', @@ -389,12 +375,12 @@ def _run_dygraph(instance, input, program_holder): else: from paddle.jit.dy2static.pir_partial_program import PartialProgramLayer - inputs = program_holder._input_vars - outputs = program_holder._output_vars - parameters = (persistable_tensors, program_holder._persistable_vars) + inputs = program_holder.input_vars + outputs = program_holder.output_vars + parameters = (persistable_tensors, program_holder.persistable_vars) layer = PartialProgramLayer( - program_holder._infer_program, + program_holder.infer_program, inputs, outputs, parameters, From 4051bccb2e36020128106120ae789870185619c5 Mon Sep 17 00:00:00 2001 From: wangruting Date: Sun, 28 Apr 2024 02:56:13 +0000 Subject: [PATCH 04/15] combine eval and train --- python/paddle/jit/pir_translated_layer.py | 98 ++++--------------- .../legacy_test/test_io_save_load.py | 58 ++++++----- 2 files changed, 51 insertions(+), 105 deletions(-) diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py index c162606e77e37c..f01b55eeb7fd60 100644 --- a/python/paddle/jit/pir_translated_layer.py +++ b/python/paddle/jit/pir_translated_layer.py @@ -17,7 +17,6 @@ import numpy as np import paddle -from paddle import _legacy_C_ops from paddle.base import core, framework, unique_name from paddle.base.dygraph.base import switch_to_static_graph from paddle.framework import in_dynamic_mode @@ -69,16 +68,11 @@ def __init__(self, program, trainable): self._persistable_vars = [] self._persistable_names = [] - # execution scope - self._inner_scope = core.Scope() - self.support_train = trainable self._infer_program = program self._preprocess() - self.prepare_train_program = False - def _preprocess(self): ( self._persistable_vars, @@ -130,10 +124,6 @@ def persistable_names(self): def persistable_vars(self): return self._persistable_vars - @property - def scope(self): - return self._inner_scope - # [ PirTranslatedLayer : Run program in dygraph mode ] # @@ -271,10 +261,6 @@ def _construct_params_and_buffers( return var_dict -def _valid_vars(vars): - return vars if vars else None - - def _run_dygraph(instance, input, program_holder): # 1. prepare inputs, outputs, attrs input_tensors = [] @@ -302,10 +288,6 @@ def _run_dygraph(instance, input, program_holder): tensor.name = program_holder.input_vars[i].name input_tensor_names.append(tensor.name) input_tensors.append(tensor) - if instance._input_args_names is None: - instance._input_args_names = [ - ins.name for ins in program_holder.input_vars - ] persistable_tensors = [] for var_name in program_holder.persistable_names: @@ -320,73 +302,31 @@ def _run_dygraph(instance, input, program_holder): % var_name ) - output_tensors = [] + from paddle.jit.dy2static.pir_partial_program import PartialProgramLayer - for i, var in enumerate(program_holder.output_vars): - tensor = core.eager.Tensor( - dtype=datatype_to_vartype[var.dtype], - dims=var.shape, - name=var.name, - type=core.VarDesc.VarType.LOD_TENSOR, - persistable=False, - ) - output_tensors.append(tensor) + inputs = program_holder.input_vars + outputs = program_holder.output_vars + parameters = (persistable_tensors, program_holder.persistable_vars) - # hold forward variables - tmp_scope_vec = [program_holder.scope] + layer = PartialProgramLayer( + program_holder.infer_program, + inputs, + outputs, + parameters, + ) + instance.layer = layer - # 2. run program by op if instance._is_test: - attrs = [ - 'forward_global_block', - program_holder.infer_program.global_block(), - 'backward_global_block', - paddle.pir.Program().global_block(), - 'is_test', - instance._is_test, - 'program_id', - paddle.utils._hash_with_id(program_holder._infer_program, instance), - 'fx', - program_holder.input_vars, - 'fo', - program_holder.output_vars, - 'fp', - program_holder.persistable_vars, - 'fm', - [], - 'no_need_buffers', - [], - ] - - _legacy_C_ops.pir_run_program( - _valid_vars(input_tensors), - _valid_vars(persistable_tensors), - _valid_vars(output_tensors), - tmp_scope_vec, - None, - *attrs, - ) - - outs = output_tensors - if len(output_tensors) == 1: - outs = output_tensors[0] - return outs - + layer.training = False else: - from paddle.jit.dy2static.pir_partial_program import PartialProgramLayer - - inputs = program_holder.input_vars - outputs = program_holder.output_vars - parameters = (persistable_tensors, program_holder.persistable_vars) + if not program_holder.support_train: + raise ValueError( + "The model is not trainable, please check model_file of jit.save." + ) + else: + layer.training = True - layer = PartialProgramLayer( - program_holder.infer_program, - inputs, - outputs, - parameters, - ) - instance.layer = layer - return instance.layer(input_tensors) + return instance.layer(input_tensors) def _run_static_graph(program_holder, trace_program): diff --git a/test/deprecated/legacy_test/test_io_save_load.py b/test/deprecated/legacy_test/test_io_save_load.py index 3375435b76f34b..a9eee37b281ded 100644 --- a/test/deprecated/legacy_test/test_io_save_load.py +++ b/test/deprecated/legacy_test/test_io_save_load.py @@ -97,33 +97,39 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() - # @test_with_pir_api + @test_with_pir_api def test_when_train_with_no_grad(self): - with paddle.pir_utils.IrGuard(): - paddle.disable_static() - net = paddle.nn.Linear(1024, 1) - net = paddle.jit.to_static(net, full_graph=True) - x = paddle.rand([1024], 'float32') - x.stop_gradient = False - out = net(x) - out.backward() - x_grad = x.grad.mean() - save_path = os.path.join(self.temp_dir.name, 'train_with_no_grad') - - paddle.jit.save(net, save_path) - net1 = paddle.jit.load(save_path) - net1.eval() - - with paddle.no_grad(): - out1 = net1(x) - self.assertEqual(out, out1) - - net2 = paddle.jit.load(save_path) - net2.train() - out2 = net2(x) - out2.backward() - self.assertEqual(out, out2) - x_grad2 = x.grad.mean() + paddle.disable_static() + net = paddle.nn.Linear(1024, 1) + net = paddle.jit.to_static(net, full_graph=True) + x = paddle.rand([1024], 'float32') + x.stop_gradient = False + out = net(x) + out.backward() + x_grad = x.grad.mean() + x.clear_grad() + + # jit.save + save_path = os.path.join(self.temp_dir.name, 'train_with_no_grad') + paddle.jit.save(net, save_path) + + # test eval mode + net1 = paddle.jit.load(save_path) + net1.eval() + + with paddle.no_grad(): + out1 = net1(x) + self.assertEqual(out, out1) + + # test train mode + net2 = paddle.jit.load(save_path) + net2.train() + out2 = net2(x) + out2.backward() + self.assertEqual(out, out2) + x_grad2 = x.grad.mean() + if paddle.framework.in_pir_mode(): + self.assertEqual(x_grad, x_grad2) if __name__ == '__main__': From fb1a29f84bf1ff97fb0042ea0c11636dfc0f2331 Mon Sep 17 00:00:00 2001 From: wangruting Date: Sun, 28 Apr 2024 11:56:01 +0000 Subject: [PATCH 05/15] modify 8 case of jit.save.load --- python/paddle/jit/api.py | 38 ++++++++- python/paddle/jit/pir_translated_layer.py | 47 ++++++++++- python/paddle/pir_utils.py | 67 ++++++++++++++++ python/paddle/static/input.py | 2 +- python/paddle/static/pir_io.py | 16 ++-- test/legacy_test/test_jit_save_load.py | 95 ++++++++++++++++------- 6 files changed, 225 insertions(+), 40 deletions(-) diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py index 3c343d5829273e..dc2264e66fbdcc 100644 --- a/python/paddle/jit/api.py +++ b/python/paddle/jit/api.py @@ -537,9 +537,26 @@ def _get_output_vars(outputs, output_spec, with_hook=False): ) result_list = [] if use_pir_api(): + from paddle.autograd.backward_utils import ValueSet + for var in paddle.utils.flatten(outputs): if isinstance(var, paddle.pir.Value): result_list.append(var) + + if output_spec is not None: + if len(output_spec) == len(result_list): + for var in output_spec: + if var not in ValueSet(result_list): + warnings.warn(name_no_exists_error % var.name) + else: + result_set = ValueSet(result_list) + result_list = [] + for var in output_spec: + if var not in result_set: + raise ValueError(name_no_exists_error % var.name) + else: + result_list.append(var) + else: output_vars_dict = OrderedDict() for var in paddle.utils.flatten(outputs): @@ -560,6 +577,7 @@ def _get_output_vars(outputs, output_spec, with_hook=False): raise ValueError(name_no_exists_error % var.name) else: result_list.append(output_vars_dict[var.name]) + return result_list @@ -960,7 +978,9 @@ def save(layer, path, input_spec=None, **configs): for var in paddle.utils.flatten(input_spec): if isinstance(var, paddle.static.InputSpec): inner_input_spec.append(var) - elif isinstance(var, (core.eager.Tensor, Variable)): + elif isinstance( + var, (core.eager.Tensor, Variable, paddle.pir.Value) + ): inner_input_spec.append( paddle.static.InputSpec.from_tensor(var) ) @@ -1193,11 +1213,25 @@ def save(layer, path, input_spec=None, **configs): clone_program = concrete_program.main_program.clone() clone_input_vars = input_vars clone_output_vars = output_vars + else: value_map = paddle.pir.IrMapping() clone_program = concrete_program.main_program.clone(value_map) - clone_input_vars = [value_map.look_up(v) for v in input_vars] + clone_input_vars = [] + for v in input_vars: + if type(v) is paddle.static.InputSpec: + name = v.name + for op in clone_program.global_block().ops: + if ( + op.name() == 'pd_op.data' + and op.attrs()["name"] == name + ): + clone_input_vars.append(op.result(0)) + else: + clone_input_vars.append(value_map.look_up(v)) + clone_output_vars = [value_map.look_up(v) for v in output_vars] + save_inference_model( path_prefix=file_prefix, feed_vars=clone_input_vars, diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py index f01b55eeb7fd60..8cffb527601060 100644 --- a/python/paddle/jit/pir_translated_layer.py +++ b/python/paddle/jit/pir_translated_layer.py @@ -326,7 +326,8 @@ def _run_dygraph(instance, input, program_holder): else: layer.training = True - return instance.layer(input_tensors) + out = instance.layer(input_tensors) + return out def _run_static_graph(program_holder, trace_program): @@ -558,3 +559,47 @@ def train(self): def eval(self): self._is_test = True self.training = False + + def _get_program_holder(self, method_name='forward'): + program_holder = self._program_holder_dict.get(method_name, None) + if program_holder is None: + raise ValueError( + "The method `%s` does not exist in loaded PirTranslatedLayer." + % method_name + ) + return program_holder + + def _input_spec(self, method_name='forward'): + # 1. get program holder + program_holder = self._get_program_holder(method_name) + + # 2. build input spec by input desc + input_spec = [] + for var in program_holder.input_vars: + spec = paddle.static.InputSpec( + shape=var.shape, + dtype=var.dtype, + name=var.name, + ) + input_spec.append(spec) + + return input_spec + + def _output_spec(self, method_name='forward'): + # 1. get program holder + program_holder = self._get_program_holder(method_name) + + # 2. build output spec by output desc + output_spec = [] + for var in program_holder.output_vars: + # NOTE(chenweihang): InputSpec describes a tensor, not just input. + # Maybe the name is not good enough. Here we use InputSpec to + # construct the description of Output tensor + spec = paddle.static.InputSpec( + shape=var.shape, + dtype=var.dtype, + name=var.name, + ) + output_spec.append(spec) + + return output_spec diff --git a/python/paddle/pir_utils.py b/python/paddle/pir_utils.py index d2a93f2b8e5576..2c4b8bf28d67d7 100644 --- a/python/paddle/pir_utils.py +++ b/python/paddle/pir_utils.py @@ -127,6 +127,61 @@ def __exit__(self, exc_type, exc_val, exc_tb): _switch_to_pir_() +class DygraphPirGuard: + def __enter__(self): + self.old_flag = paddle.base.framework.get_flags("FLAGS_enable_pir_api")[ + "FLAGS_enable_pir_api" + ] + if not self.old_flag: + paddle.framework.set_flags({"FLAGS_enable_pir_api": True}) + paddle.base.framework.global_var._use_pir_api_ = True + bind_datatype() + self._switch_to_pir() + + def __exit__(self, exc_type, exc_val, exc_tb): + if not self.old_flag: + paddle.framework.set_flags({"FLAGS_enable_pir_api": False}) + paddle.base.framework.global_var._use_pir_api_ = False + bind_vartype() + self._switch_to_old_ir() + + def _switch_to_pir(self): + if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[ + "FLAGS_enable_pir_api" + ]: + _switch_to_pir_() + + def _switch_to_old_ir(self): + if not paddle.base.framework.get_flags("FLAGS_enable_pir_api")[ + "FLAGS_enable_pir_api" + ]: + _switch_to_old_ir_() + else: + raise RuntimeError( + "IrGuard._switch_to_old_ir only work when paddle.framework.in_pir_mode() is false, \ + please set FLAGS_enable_pir_api = false" + ) + + +class DygraphOldIrGuard: + def __enter__(self): + self.old_flag = paddle.base.framework.get_flags("FLAGS_enable_pir_api")[ + "FLAGS_enable_pir_api" + ] + if self.old_flag: + paddle.framework.set_flags({"FLAGS_enable_pir_api": False}) + paddle.base.framework.global_var._use_pir_api_ = False + bind_vartype() + _switch_to_old_ir_() + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.old_flag: + paddle.framework.set_flags({"FLAGS_enable_pir_api": True}) + paddle.base.framework.global_var._use_pir_api_ = True + bind_datatype() + _switch_to_pir_() + + def test_with_pir_api(func): @wraps(func) def impl(*args, **kwargs): @@ -145,3 +200,15 @@ def impl(*args, **kwargs): func(*args, **kwargs) return impl + + +def test_with_dygraph_pir(func): + @wraps(func) + def impl(*args, **kwargs): + with DygraphOldIrGuard(): + func(*args, **kwargs) + + with DygraphPirGuard(): + func(*args, **kwargs) + + return impl diff --git a/python/paddle/static/input.py b/python/paddle/static/input.py index 4cc2d1b9187459..d2c98d859a4e2f 100644 --- a/python/paddle/static/input.py +++ b/python/paddle/static/input.py @@ -251,7 +251,7 @@ def from_tensor(cls, tensor, name=None): InputSpec(shape=(2, 2), dtype=paddle.float32, name=x, stop_gradient=False) """ - if isinstance(tensor, (Variable, core.eager.Tensor)): + if isinstance(tensor, (Variable, core.eager.Tensor, paddle.pir.Value)): return cls(tensor.shape, tensor.dtype, name or tensor.name) else: raise ValueError( diff --git a/python/paddle/static/pir_io.py b/python/paddle/static/pir_io.py index 8211be49bb2820..bf3dad7f9dbd9c 100644 --- a/python/paddle/static/pir_io.py +++ b/python/paddle/static/pir_io.py @@ -36,7 +36,6 @@ from paddle.base.executor import Executor, global_scope from paddle.base.framework import ( dygraph_not_support, - process_type_promotion, static_only, ) from paddle.base.log_helper import get_logger @@ -219,7 +218,7 @@ def normalize_pir_program(program, feed_vars, fetch_vars, **kwargs): uniq_fetch_vars = [] for var in fetch_vars: if var.dtype != paddle.bool: - var_ = paddle.scale(fetch_vars[0], 1.0) + var_ = paddle.scale(var, 1.0) uniq_fetch_vars.append(var_) fetch_vars = uniq_fetch_vars @@ -658,12 +657,6 @@ def save_pir_inference_model( _check_vars('fetch_vars', fetch_vars) program = _get_valid_program(kwargs.get('program', None)) - - # do type promotion - program = process_type_promotion(program) - - clip_extra = kwargs.get('clip_extra', True) - # serialize and save program program = normalize_pir_program( program, @@ -671,7 +664,12 @@ def save_pir_inference_model( fetch_vars, skip_prune_program=kwargs.get('skip_prune_program', False), ) - paddle.core.serialize_pir_program(program, model_path, 1, True, False, True) + + readable = kwargs.get('readable', False) + trainable = kwargs.get('trainable', True) + paddle.core.serialize_pir_program( + program, model_path, 1, True, readable, trainable + ) # serialize and save params save_dirname = os.path.dirname(params_path) diff --git a/test/legacy_test/test_jit_save_load.py b/test/legacy_test/test_jit_save_load.py index e9fbf29759b401..d69d2071c709d2 100644 --- a/test/legacy_test/test_jit_save_load.py +++ b/test/legacy_test/test_jit_save_load.py @@ -27,6 +27,7 @@ from paddle.jit.api import to_static from paddle.jit.translated_layer import INFER_PARAMS_INFO_SUFFIX from paddle.nn import Linear +from paddle.pir_utils import test_with_dygraph_pir from paddle.static import InputSpec BATCH_SIZE = 32 @@ -155,12 +156,6 @@ def __init__(self, in_size, out_size): self._linear1 = Linear(in_size, out_size) self._linear2 = Linear(in_size, out_size) - @to_static( - input_spec=[ - InputSpec([None, 8], dtype='float32'), - InputSpec([None, 8], dtype='float32'), - ] - ) def forward(self, x, y): x_out = self._linear1(x) y_out = self._linear2(y) @@ -174,12 +169,6 @@ def __init__(self, in_size, out_size): self._linear1 = Linear(in_size, out_size) self._linear2 = Linear(in_size, out_size) - @to_static( - input_spec=( - InputSpec([None, 8], dtype='float32'), - InputSpec([None, 8], dtype='float32'), - ) - ) def forward(self, x, y): x_out = self._linear1(x) y_out = self._linear2(y) @@ -237,12 +226,6 @@ def __init__(self, in_size, out_size): super().__init__() self._linear = Linear(in_size, out_size) - @paddle.jit.to_static( - input_spec=[ - {'img': InputSpec(shape=[None, 8], dtype='float32', name='img')}, - {'label': InputSpec(shape=[None, 1], dtype='int64', name='label')}, - ] - ) def forward(self, img, label): out = self._linear(img['img']) # not return loss to avoid prune output @@ -376,14 +359,16 @@ def train_and_save_model(self, model_path=None): self.assertEqual(orig_input_types, new_input_types) return layer + @test_with_dygraph_pir def test_save_load(self): # train and save model train_layer = self.train_and_save_model() # load model loaded_layer = paddle.jit.load(self.model_path) self.load_and_inference(train_layer, loaded_layer) - self.load_dygraph_state_dict(train_layer) self.load_and_finetune(train_layer, loaded_layer) + if not paddle.framework.use_pir_api(): + self.load_dygraph_state_dict(train_layer) def load_and_inference(self, train_layer, infer_layer): train_layer.eval() @@ -427,6 +412,7 @@ def test_load_dygraph_no_path(self): with self.assertRaises(ValueError): model_dict = paddle.load(model_path) + @test_with_dygraph_pir def test_jit_load_no_path(self): path = os.path.join( self.temp_dir.name, "test_jit_save_load.no_path/model_path" @@ -444,6 +430,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_nest_output(self): x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) @@ -465,11 +452,28 @@ def test_nest_output(self): class TestSaveLoadWithDictInput(unittest.TestCase): + @test_with_dygraph_pir def test_dict_input(self): # NOTE: This net cannot be executed, it is just # a special case for exporting models in model validation # We DO NOT recommend this writing way of Layer net = LinearNetWithDictInput(8, 8) + net = paddle.jit.to_static( + net, + input_spec=[ + { + 'img': InputSpec( + shape=[None, 8], dtype=paddle.float32, name='img' + ) + }, + { + 'label': InputSpec( + shape=[None, 1], dtype=paddle.int64, name='label' + ) + }, + ], + full_graph=True, + ) # net.forward.concrete_program.inputs: # (<__main__.LinearNetWithDictInput object at 0x7f2655298a98>, # {'img': var img : base.VarType.LOD_TENSOR.shape(-1, 8).astype(VarType.FP32)}, @@ -484,7 +488,11 @@ def test_dict_input(self): layer=net, path=path, input_spec=[ - {'img': InputSpec(shape=[None, 8], dtype='float32', name='img')} + { + 'img': InputSpec( + shape=[None, 8], dtype=paddle.float32, name='img' + ) + } ], ) @@ -495,10 +503,12 @@ def test_dict_input(self): # loaded_net._input_spec(): # [InputSpec(shape=(-1, 8), dtype=VarType.FP32, name=img)] self.assertEqual(len(loaded_net._input_spec()), 1) + self.assertEqual(len(loaded_net._output_spec()), 1) temp_dir.cleanup() class TestSaveLoadWithDictInputNoPrune(unittest.TestCase): + @test_with_dygraph_pir def test_dict_input(self): net = LinearNetWithDictInputNoPrune(8, 8) temp_dir = tempfile.TemporaryDirectory() @@ -539,11 +549,14 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_with_input_spec(self): net = LinearNetReturnLoss(8, 8) # set x.shape = [None, 8] net.forward = to_static( - net.forward, input_spec=[InputSpec([None, 8], name='x')] + net.forward, + input_spec=[InputSpec([None, 8], name='x')], + full_graph=True, ) model_path = os.path.join( @@ -552,7 +565,10 @@ def test_with_input_spec(self): # check inputs and outputs self.assertTrue(len(net.forward.inputs) == 1) input_x = net.forward.inputs[0] - self.assertTrue(input_x.shape == (-1, 8)) + if paddle.framework.use_pir_api(): + self.assertTrue(input_x.shape == [-1, 8]) + else: + self.assertTrue(input_x.shape == (-1, 8)) self.assertTrue(input_x.name == 'x') # 1. prune loss @@ -564,8 +580,17 @@ def test_with_input_spec(self): x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) pred = infer_layer(x) + @test_with_dygraph_pir def test_multi_in_out(self): net = LinearNetMultiInput(8, 8) + net = paddle.jit.to_static( + net, + input_spec=[ + InputSpec([None, 8], dtype='float32'), + InputSpec([None, 8], dtype='float32'), + ], + full_graph=True, + ) model_path = os.path.join( self.temp_dir.name, "multi_inout.output_spec1/model" @@ -574,8 +599,12 @@ def test_multi_in_out(self): self.assertTrue(len(net.forward.inputs) == 2) input_x = net.forward.inputs[0] input_y = net.forward.inputs[1] - self.assertTrue(input_x.shape == (-1, 8)) - self.assertTrue(input_y.shape == (-1, 8)) + if paddle.framework.use_pir_api(): + self.assertTrue(input_x.shape == [-1, 8]) + self.assertTrue(input_y.shape == [-1, 8]) + else: + self.assertTrue(input_x.shape == (-1, 8)) + self.assertTrue(input_y.shape == (-1, 8)) # 2. prune loss output_spec = net.forward.outputs[:2] @@ -602,9 +631,17 @@ def test_multi_in_out(self): # 4. assert pred_x == pred_xx np.testing.assert_allclose(pred_x.numpy(), pred_xx.numpy(), rtol=1e-05) + @test_with_dygraph_pir def test_multi_in_out1(self): net = LinearNetMultiInput1(8, 8) - + net = paddle.jit.to_static( + net, + input_spec=( + InputSpec([None, 8], dtype='float32'), + InputSpec([None, 8], dtype='float32'), + ), + full_graph=True, + ) model_path = os.path.join( self.temp_dir.name, "multi_inout1.output_spec1/model" ) @@ -612,8 +649,12 @@ def test_multi_in_out1(self): self.assertTrue(len(net.forward.inputs) == 2) input_x = net.forward.inputs[0] input_y = net.forward.inputs[1] - self.assertTrue(input_x.shape == (-1, 8)) - self.assertTrue(input_y.shape == (-1, 8)) + if paddle.framework.use_pir_api(): + self.assertTrue(input_x.shape == [-1, 8]) + self.assertTrue(input_y.shape == [-1, 8]) + else: + self.assertTrue(input_x.shape == (-1, 8)) + self.assertTrue(input_y.shape == (-1, 8)) # 2. prune loss output_spec = net.forward.outputs[:2] From 25efe76374a48e566a7dfa26a51604a44ca5d1ae Mon Sep 17 00:00:00 2001 From: wangruting Date: Mon, 29 Apr 2024 11:02:48 +0000 Subject: [PATCH 06/15] modify jit_save_load case --- paddle/fluid/pybind/pir.cc | 35 +- python/paddle/jit/api.py | 23 +- python/paddle/jit/pir_translated_layer.py | 34 +- test/legacy_test/test_jit_save_load.py | 2136 --------------------- 4 files changed, 76 insertions(+), 2152 deletions(-) delete mode 100644 test/legacy_test/test_jit_save_load.py diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc index 20fc28ec128127..fb21ce98446a81 100644 --- a/paddle/fluid/pybind/pir.cc +++ b/paddle/fluid/pybind/pir.cc @@ -782,6 +782,35 @@ std::string GetValueName(Value value) { } } +void SetValueName(Value value, const std::string name) { + pir::Operation *define_op = value.defining_op(); + if (define_op->isa()) { + define_op->set_attribute( + "parameter_name", + pir::StrAttribute::get(pir::IrContext::Instance(), name)); + } else if (define_op->isa()) { + define_op->set_attribute( + "name", pir::StrAttribute::get(pir::IrContext::Instance(), name)); + } else if (auto block_arg = value.dyn_cast()) { + PADDLE_THROW( + phi::errors::InvalidArgument("Can Not set name for BlockArgument! ")); + } else if (value.first_use()) { + auto nextOp = value.first_use().owner(); + if (nextOp->isa<::pir::ShadowOutputOp>()) { + nextOp->set_attribute( + "output_name", + pir::StrAttribute::get(pir::IrContext::Instance(), name)); + } else { + PADDLE_THROW(phi::errors::InvalidArgument( + "Currently, we can only set name of Value which is " + "shadowoutput ")); + } + } else { + PADDLE_THROW(phi::errors::InvalidArgument( + "Currently, we can only set name of Value that " + "is persistable")); + } +} const phi::DDim &GetValueDims(Value value) { if (!value.type()) { PADDLE_THROW(phi::errors::InvalidArgument("The type of value is nullptr.")); @@ -913,8 +942,10 @@ void BindValue(py::module *m) { return ss.str(); } }) - .def_property_readonly("name", - [](Value self) { return GetValueName(self); }) + .def_property( + "name", + [](Value self) { return GetValueName(self); }, + [](Value self, const std::string &name) { SetValueName(self, name); }) .def_property_readonly( "has_name", [](Value self) { diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py index dc2264e66fbdcc..d7402c3ea5f3c7 100644 --- a/python/paddle/jit/api.py +++ b/python/paddle/jit/api.py @@ -531,6 +531,11 @@ def _get_output_vars(outputs, output_spec, with_hook=False): "in configs.output_spec is the output tensor of " "Layer.forward method." ) + output_spec_is_not_value_error = ( + "tensor `%s` is not support in pir mode, " + "because pir value has no name sometimes, especially as ouptut," + " so we can't check tensor's name with output var name" + ) if output_spec and with_hook: raise RuntimeError( "Currently not support specify output_spec while founding pre/post hooks in your outermost layer." @@ -546,16 +551,24 @@ def _get_output_vars(outputs, output_spec, with_hook=False): if output_spec is not None: if len(output_spec) == len(result_list): for var in output_spec: - if var not in ValueSet(result_list): - warnings.warn(name_no_exists_error % var.name) + if not isinstance(var, paddle.pir.Value): + warnings.warn(output_spec_is_not_value_error % var.name) + else: + if var not in ValueSet(result_list): + warnings.warn(name_no_exists_error % var.name) else: result_set = ValueSet(result_list) result_list = [] for var in output_spec: - if var not in result_set: - raise ValueError(name_no_exists_error % var.name) + if not isinstance(var, paddle.pir.Value): + raise ValueError( + output_spec_is_not_value_error % var.name + ) else: - result_list.append(var) + if var not in result_set: + raise ValueError(name_no_exists_error % var.name) + else: + result_list.append(var) else: output_vars_dict = OrderedDict() diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py index 4ebe1b12c13b68..fa005b2ba1f106 100644 --- a/python/paddle/jit/pir_translated_layer.py +++ b/python/paddle/jit/pir_translated_layer.py @@ -47,6 +47,11 @@ def _generate_unique_var_name(prefix): return unique_name.generate_with_ignorable_key(prefix) +@switch_to_static_graph +def _generate_unique_var_name(prefix): + return unique_name.generate(prefix) + + def _get_pir_persistable_var_names(program): persistable_vars = [] persistable_names = [] @@ -54,6 +59,7 @@ def _get_pir_persistable_var_names(program): for var in block.all_parameters(): if var.persistable: persistable_vars.append(var) + var.name = _generate_unique_var_name(var.name) persistable_names.append(var.name) return persistable_vars, persistable_names @@ -78,7 +84,6 @@ def _preprocess(self): self._persistable_vars, self._persistable_names, ) = _get_pir_persistable_var_names(self._infer_program) - block = self._infer_program.global_block() for op in block.ops: if op.name() == 'pd_op.data': @@ -163,16 +168,27 @@ def _load_pir_persistable_vars(model_path, program_holder, params_filename): load_densetensor_list = [] persistable_var = program_holder.persistable_vars persistable_var_name = program_holder.persistable_names + for name, var in sorted(zip(persistable_var_name, persistable_var)): - new_var = core.eager.Tensor( - dtype=datatype_to_vartype[var.dtype], - dims=var.shape, - name=var.name, - type=core.VarDesc.VarType.LOD_TENSOR, - place=framework._current_expected_place(), - persistable=False, - ) + if var.persistable: + # use default shape and dtype + new_var = framework.EagerParamBase( + shape=var.shape, # only to pass check, this shape is not meaningful + dtype=core.VarDesc.VarType.FP32, + name=var.name, + persistable=True, + ) + else: + new_var = core.eager.Tensor( + dtype=datatype_to_vartype[var.dtype], + dims=var.shape, + name=var.name, + type=core.VarDesc.VarType.LOD_TENSOR, + 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()) diff --git a/test/legacy_test/test_jit_save_load.py b/test/legacy_test/test_jit_save_load.py deleted file mode 100644 index afa5046b101eb5..00000000000000 --- a/test/legacy_test/test_jit_save_load.py +++ /dev/null @@ -1,2136 +0,0 @@ -# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. -# 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. - -import copy -import os -import pickle -import shutil -import tempfile -import unittest - -import numpy as np - -import paddle -from paddle import base -from paddle.base import unique_name -from paddle.jit.api import to_static -from paddle.jit.translated_layer import INFER_PARAMS_INFO_SUFFIX -from paddle.nn import Linear -from paddle.pir_utils import test_with_dygraph_pir -from paddle.static import InputSpec - -BATCH_SIZE = 32 -BATCH_NUM = 10 -SEED = 10 - - -def random_batch_reader(input_size, label_size): - def _get_random_inputs_and_labels(input_size, label_size): - np.random.seed(SEED) - input = np.random.random(size=input_size).astype('float32') - label = np.random.random(size=label_size).astype('int64') - return input, label - - def __reader__(): - for _ in range(BATCH_NUM): - batch_input, batch_label = _get_random_inputs_and_labels( - [BATCH_SIZE, input_size], [BATCH_SIZE, label_size] - ) - yield batch_input, batch_label - - return __reader__ - - -class LinearNet(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - @to_static - def forward(self, x): - return self._linear(x) - - -class LinearNetWithInputSpec(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - @to_static(input_spec=[InputSpec(shape=[None, 784], dtype='float32')]) - def forward(self, x): - return self._linear(x) - - -class LinearNetNotDeclarative(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - def forward(self, x): - return self._linear(x) - - -class LinerNetWithLabel(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - @to_static( - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name="image"), - InputSpec(shape=[None, 1], dtype='int64', name="label"), - ] - ) - def forward(self, x, label): - out = self._linear(x) - loss = paddle.nn.functional.cross_entropy( - out, label, reduction='none', use_softmax=False - ) - avg_loss = paddle.mean(loss) - return out, avg_loss - - -class LinerNetWithPruneInput(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - @to_static( - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name="image"), - InputSpec(shape=[None, 1], dtype='int64', name="label"), - ] - ) - def forward(self, x, label): - out = self._linear(x) - loss = paddle.nn.functional.cross_entropy( - out, label, reduction='none', use_softmax=False - ) - avg_loss = paddle.mean(loss) - return out - - -class LinerNetWithUselessInput(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - @to_static( - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name="image"), - InputSpec(shape=[None, 1], dtype='int64', name="label"), - ] - ) - def forward(self, x, label): - out = self._linear(x) - return out - - -class LinearNetReturnLoss(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - @to_static - def forward(self, x): - y = self._linear(x) - z = self._linear(y) - loss = paddle.mean(z) - return z, loss - - -class LinearNetMultiInput(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear1 = Linear(in_size, out_size) - self._linear2 = Linear(in_size, out_size) - - def forward(self, x, y): - x_out = self._linear1(x) - y_out = self._linear2(y) - loss = paddle.mean(x_out + y_out) - return x_out, y_out, loss - - -class LinearNetMultiInput1(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear1 = Linear(in_size, out_size) - self._linear2 = Linear(in_size, out_size) - - def forward(self, x, y): - x_out = self._linear1(x) - y_out = self._linear2(y) - loss = paddle.mean(x_out + y_out) - return x_out, y_out, loss - - -class MultiLoadingLinearNet(paddle.nn.Layer): - def __init__(self, size, model_path): - super().__init__() - self._linear = Linear(size, size) - self._load_linear1 = paddle.jit.load(model_path) - self._load_linear2 = paddle.jit.load(model_path) - - @to_static - def forward(self, x): - tmp1 = self._linear(x) - tmp2 = self._load_linear1(tmp1) - tmp3 = self._load_linear2(tmp2) - y = self._linear(tmp3) - return y - - -class LinearNetReturnHidden(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear_1 = Linear(in_size, out_size) - self._linear_2 = Linear(in_size, out_size) - - @to_static - def forward(self, x): - y = self._linear_1(x) - z = self._linear_2(y) - loss = paddle.mean(z) - return y, loss - - -class LinearNetWithNestOut(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear_1 = Linear(in_size, out_size) - self._linear_2 = Linear(in_size, out_size) - - @to_static - def forward(self, x): - y = self._linear_1(x) - z = self._linear_2(y) - out = y + z - loss = paddle.mean(out) - return y, [(z, loss), out] - - -class LinearNetWithDictInput(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - def forward(self, img, label): - out = self._linear(img['img']) - # not return loss to avoid prune output - loss = paddle.nn.functional.cross_entropy(out, label['label']) - return out - - -class LinearNetWithDictInputNoPrune(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear = Linear(in_size, out_size) - - def forward(self, img): - out = self._linear(img['img'] + img['img2']) - return out - - -class EmptyLayer(paddle.nn.Layer): - def __init__(self): - super().__init__() - - @paddle.jit.to_static - def forward(self, x): - return x - - -class NoParamLayer(paddle.nn.Layer): - def __init__(self): - super().__init__() - - @paddle.jit.to_static - def forward(self, x, y): - return x + y - - -class LinearNetWithMultiStaticFunc(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear_0 = Linear(in_size, out_size) - self._linear_1 = Linear(in_size, out_size) - self._scale = paddle.to_tensor([9.9]) - - @paddle.jit.to_static - def forward(self, x): - return self._linear_0(x) - - @paddle.jit.to_static - def forward_no_param(self, x): - return x - - @paddle.jit.to_static - def forward_general(self, x): - return self._linear_0(x) + self._linear_1(x) * self._scale - - -class LinearNetWithNonLexicographicalOrderDict(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear_u = Linear(in_size, out_size) - self._linear_v = Linear(in_size, out_size) - self._linear_w = Linear(in_size, out_size) - self._linear_p = Linear(in_size, out_size) - - def forward(self, x): - u = self._linear_u(x) - v = self._linear_v(x) - w = self._linear_w(x) - p = self._linear_p(x) - return { - "u": u, - "v": v, - "w": w, - "p": p, - } - - -class LinearNetWithNestedNonLexicographicalOrderDict(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self._linear_u = Linear(in_size, out_size) - self._linear_v = Linear(in_size, out_size) - self._linear_w = Linear(in_size, out_size) - self._linear_p = Linear(in_size, out_size) - self._linear_y = Linear(in_size, out_size) - self._linear_x = Linear(in_size, out_size) - - def forward(self, x_): - u = self._linear_u(x_) - v = self._linear_v(x_) - w = self._linear_w(x_) - p = self._linear_p(x_) - - x = self._linear_p(x_) - y = self._linear_p(x_) - return { - "u": u, - "v": v, - "w": w, - "p": p, - "a": { - "x": x, - "y": y, - }, - } - - -def train(layer, input_size=784, label_size=1): - # create optimizer - sgd = paddle.optimizer.SGD( - learning_rate=0.01, parameters=layer.parameters() - ) - # create data loader - train_loader = base.io.DataLoader.from_generator(capacity=5) - train_loader.set_batch_generator( - random_batch_reader(input_size, label_size) - ) - # train - for data in train_loader(): - img, label = data - label.stop_gradient = True - - cost = layer(img) - - loss = paddle.nn.functional.cross_entropy( - cost, label, reduction='none', use_softmax=True - ) - avg_loss = paddle.mean(loss) - - avg_loss.backward() - sgd.minimize(avg_loss) - layer.clear_gradients() - return [img], layer, avg_loss - - -def train_with_label(layer, input_size=784, label_size=1): - # create optimizer - sgd = paddle.optimizer.SGD( - learning_rate=0.01, parameters=layer.parameters() - ) - # create data loader - train_loader = base.io.DataLoader.from_generator(capacity=5) - train_loader.set_batch_generator( - random_batch_reader(input_size, label_size) - ) - # train - for data in train_loader(): - img, label = data - label.stop_gradient = True - - out, avg_loss = layer(img, label) - - avg_loss.backward() - sgd.minimize(avg_loss) - layer.clear_gradients() - return out - - -class TestJitSaveLoad(unittest.TestCase): - def setUp(self): - self.temp_dir = tempfile.TemporaryDirectory() - self.model_path = os.path.join( - self.temp_dir.name, "test_jit_save_load/model" - ) - # enable dygraph mode - base.enable_dygraph() - # config seed - paddle.seed(SEED) - paddle.framework.random._manual_program_seed(SEED) - - def tearDown(self): - self.temp_dir.cleanup() - - def train_and_save_model(self, model_path=None): - layer = LinearNet(784, 1) - example_inputs, layer, _ = train(layer) - final_model_path = model_path if model_path else self.model_path - orig_input_types = [type(x) for x in example_inputs] - paddle.jit.save( - layer=layer, path=final_model_path, input_spec=example_inputs - ) - new_input_types = [type(x) for x in example_inputs] - self.assertEqual(orig_input_types, new_input_types) - return layer - - @test_with_dygraph_pir - def test_save_load(self): - # train and save model - train_layer = self.train_and_save_model() - # load model - loaded_layer = paddle.jit.load(self.model_path) - self.load_and_inference(train_layer, loaded_layer) - self.load_and_finetune(train_layer, loaded_layer) - if not paddle.framework.use_pir_api(): - self.load_dygraph_state_dict(train_layer) - - def load_and_inference(self, train_layer, infer_layer): - train_layer.eval() - infer_layer.eval() - # inference & compare - x = paddle.to_tensor(np.random.random((1, 784)).astype('float32')) - np.testing.assert_array_equal( - train_layer(x).numpy(), infer_layer(x).numpy() - ) - - def load_and_finetune(self, train_layer, load_train_layer): - train_layer.train() - load_train_layer.train() - # train & compare - img0, _, train_loss = train(train_layer) - img1, _, load_train_loss = train(load_train_layer) - np.testing.assert_array_equal( - train_loss.numpy(), load_train_loss.numpy() - ) - - def load_dygraph_state_dict(self, train_layer): - train_layer.eval() - # construct new model - new_layer = LinearNet(784, 1) - orig_state_dict = new_layer.state_dict() - load_state_dict = paddle.load(self.model_path) - for structured_name in orig_state_dict: - self.assertTrue(structured_name in load_state_dict) - new_layer.set_state_dict(load_state_dict) - new_layer.eval() - # inference & compare - x = paddle.to_tensor(np.random.random((1, 784)).astype('float32')) - np.testing.assert_array_equal( - train_layer(x).numpy(), new_layer(x).numpy() - ) - - def test_load_dygraph_no_path(self): - model_path = os.path.join( - self.temp_dir.name, "test_jit_save_load.no_path/model_path" - ) - with self.assertRaises(ValueError): - model_dict = paddle.load(model_path) - - @test_with_dygraph_pir - def test_jit_load_no_path(self): - path = os.path.join( - self.temp_dir.name, "test_jit_save_load.no_path/model_path" - ) - with self.assertRaises(ValueError): - loaded_layer = paddle.jit.load(path) - - -class TestSaveLoadWithNestOut(unittest.TestCase): - def setUp(self): - # enable dygraph mode - base.enable_dygraph() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - @test_with_dygraph_pir - def test_nest_output(self): - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - - net = LinearNetWithNestOut(8, 8) - dy_outs = paddle.utils.flatten(net(x)) - net = to_static(net, input_spec=[InputSpec([None, 8], name='x')]) - - model_path = os.path.join(self.temp_dir.name, "net_with_nest_out/model") - paddle.jit.save(net, model_path) - - load_net = paddle.jit.load(model_path) - load_outs = paddle.utils.flatten(load_net(x)) - - self.assertTrue(len(dy_outs) == 4) - for dy_out, load_out in zip(dy_outs, load_outs): - np.testing.assert_allclose( - dy_out.numpy(), load_out.numpy(), rtol=1e-05 - ) - - -class TestSaveLoadWithNonLexicographicalOrderDict(unittest.TestCase): - def setUp(self): - # enable dygraph mode - base.enable_dygraph() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_output_same_order(self): - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - - model = LinearNetWithNonLexicographicalOrderDict(8, 8) - - dy_output_dict = model(x) - - st_model = paddle.jit.to_static(model) - st_output_dict = st_model(x) - - paddle.jit.save(st_model, "./test_jit_save_load") - loaded_model = paddle.jit.load("./test_jit_save_load") - loaded_output_seq = loaded_model(x) - - self.assertTrue(len(dy_output_dict) == 4) - self.assertTrue(len(st_output_dict) == 4) - self.assertTrue(len(loaded_output_seq) == 4) - - # 1. check whether output dict of dygraph and static graph is same - for (dy_key, dy_out), (st_key, st_out) in zip( - dy_output_dict.items(), st_output_dict.items() - ): - self.assertTrue(dy_key == st_key) - np.testing.assert_allclose( - dy_out.numpy(), st_out.numpy(), rtol=1e-05 - ) - - # 2. check whether flattened output has same order of original dict - dy_output_seq = paddle.utils.flatten(dy_output_dict) - - self.assertTrue(len(dy_output_seq) == 4) - for dy_out, flattened_dy_out in zip( - dy_output_dict.values(), dy_output_seq - ): - np.testing.assert_allclose( - dy_out.numpy(), flattened_dy_out.numpy(), rtol=1e-05 - ) - - # 3. check whether flattened output of loaded static graph has same order of dynamic's - for dy_out, loaded_out in zip( - dy_output_dict.values(), loaded_output_seq - ): - np.testing.assert_allclose( - dy_out.numpy(), loaded_out.numpy(), rtol=1e-05 - ) - - -class TestSaveLoadWithNestedNonLexicographicalOrderDict(unittest.TestCase): - def setUp(self): - # enable dygraph mode - base.enable_dygraph() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_nested_output_same_order(self): - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - - model = LinearNetWithNestedNonLexicographicalOrderDict(8, 8) - - dy_output_dict = model(x) - - st_model = paddle.jit.to_static(model) - st_output_dict = st_model(x) - - paddle.jit.save(st_model, "./test_jit_save_load2") - loaded_model = paddle.jit.load("./test_jit_save_load2") - loaded_output_seq = loaded_model(x) - - self.assertTrue(len(dy_output_dict) == 5) - self.assertTrue(len(st_output_dict) == 5) - self.assertTrue(len(loaded_output_seq) == 6) - - # check whether flattened output of loaded static graph has same order of dynamic's - dy_output_tensors = [] - for v in dy_output_dict.values(): - if isinstance(v, dict): - dy_output_tensors.extend(list(v.values())) - else: - dy_output_tensors.append(v) - - for dy_out, loaded_out in zip(dy_output_tensors, loaded_output_seq): - np.testing.assert_allclose( - dy_out.numpy(), loaded_out.numpy(), rtol=1e-05 - ) - - -class TestUtilsMapAndPack(unittest.TestCase): - def setUp(self): - # enable dygraph mode - base.enable_dygraph() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_utils_map_structure(self): - nested_list = [ - { - "d": paddle.to_tensor([1.0]), - "a": paddle.to_tensor([2.0]), - "c": paddle.to_tensor([3.0]), - "tmp": { - "b": paddle.to_tensor([4.0]), - }, - }, - [paddle.to_tensor([5.0]), paddle.to_tensor([6.0])], - [], - [ - paddle.to_tensor([7.0]), - [ - paddle.to_tensor([8.0]), - [paddle.to_tensor([9.0]), [paddle.to_tensor([10.0])]], - ], - ], - ] - FACTOR = 2 - expected_list = [ - { - "d": paddle.to_tensor([1.0]) * FACTOR, - "a": paddle.to_tensor([2.0]) * FACTOR, - "c": paddle.to_tensor([3.0]) * FACTOR, - "tmp": { - "b": paddle.to_tensor([4.0]) * FACTOR, - }, - }, - [ - paddle.to_tensor([5.0]) * FACTOR, - paddle.to_tensor([6.0]) * FACTOR, - ], - [], - [ - paddle.to_tensor([7.0]) * FACTOR, - [ - paddle.to_tensor([8.0]) * FACTOR, - [ - paddle.to_tensor([9.0]) * FACTOR, - [paddle.to_tensor([10.0]) * FACTOR], - ], - ], - ], - ] - mapped_list = paddle.utils.map_structure( - lambda x: x * FACTOR, nested_list - ) - - # test paddle.utils. - def dfs(obj1, obj2): - self.assertTrue(type(obj1) == type(obj2)) - if isinstance(obj1, list): - for i in range(len(obj1)): - dfs(obj1[i], obj2[i]) - elif isinstance(obj1, dict): - self.assertTrue(list(obj1.keys()) == list(obj2.keys())) - for k in obj1: - dfs(obj1[k], obj2[k]) - elif isinstance(obj1, paddle.Tensor): - np.testing.assert_allclose( - obj1.numpy(), obj2.numpy(), rtol=1e-05 - ) - else: - raise ValueError(f"Unsupported type: {type(obj1)} in dfs") - - dfs(expected_list, mapped_list) - - def test_utils_pack_sequence_as(self): - nested_list = [ - { - "d": paddle.to_tensor([1.0]), - "a": paddle.to_tensor([2.0]), - "c": paddle.to_tensor([3.0]), - "tmp": { - "b": paddle.to_tensor([4.0]), - }, - }, - [paddle.to_tensor([5.0]), paddle.to_tensor([6.0])], - [], - [ - paddle.to_tensor([7.0]), - [ - paddle.to_tensor([8.0]), - [paddle.to_tensor([9.0]), [paddle.to_tensor([10.0])]], - ], - ], - ] - - def dfs(obj1, obj2): - self.assertTrue(type(obj1) == type(obj2)) - if isinstance(obj1, list): - for i in range(len(obj1)): - dfs(obj1[i], obj2[i]) - elif isinstance(obj1, dict): - self.assertTrue(list(obj1.keys()) == list(obj2.keys())) - for k in obj1: - dfs(obj1[k], obj2[k]) - elif isinstance(obj1, paddle.Tensor): - np.testing.assert_allclose( - obj1.numpy(), obj2.numpy(), rtol=1e-05 - ) - else: - raise ValueError(f"Unsupported type: {type(obj1)} in dfs") - - nested_list_copy = copy.deepcopy(nested_list) - nested_list_copy_pack_back = paddle.utils.pack_sequence_as( - nested_list_copy, paddle.utils.flatten(nested_list) - ) - - dfs(nested_list_copy, nested_list_copy_pack_back) - - -class TestSaveLoadWithDictInput(unittest.TestCase): - @test_with_dygraph_pir - def test_dict_input(self): - # NOTE: This net cannot be executed, it is just - # a special case for exporting models in model validation - # We DO NOT recommend this writing way of Layer - net = LinearNetWithDictInput(8, 8) - net = paddle.jit.to_static( - net, - input_spec=[ - { - 'img': InputSpec( - shape=[None, 8], dtype=paddle.float32, name='img' - ) - }, - { - 'label': InputSpec( - shape=[None, 1], dtype=paddle.int64, name='label' - ) - }, - ], - full_graph=True, - ) - # net.forward.concrete_program.inputs: - # (<__main__.LinearNetWithDictInput object at 0x7f2655298a98>, - # {'img': var img : base.VarType.LOD_TENSOR.shape(-1, 8).astype(VarType.FP32)}, - # {'label': var label : base.VarType.LOD_TENSOR.shape(-1, 1).astype(VarType.INT64)}) - self.assertEqual(len(net.forward.concrete_program.inputs), 3) - temp_dir = tempfile.TemporaryDirectory() - path = os.path.join( - temp_dir.name, "test_jit_save_load_with_dict_input/model" - ) - # prune inputs - paddle.jit.save( - layer=net, - path=path, - input_spec=[ - { - 'img': InputSpec( - shape=[None, 8], dtype=paddle.float32, name='img' - ) - } - ], - ) - - img = paddle.randn(shape=[4, 8], dtype='float32') - loaded_net = paddle.jit.load(path) - loaded_out = loaded_net(img) - - # loaded_net._input_spec(): - # [InputSpec(shape=(-1, 8), dtype=VarType.FP32, name=img)] - self.assertEqual(len(loaded_net._input_spec()), 1) - self.assertEqual(len(loaded_net._output_spec()), 1) - temp_dir.cleanup() - - -class TestSaveLoadWithDictInputNoPrune(unittest.TestCase): - @test_with_dygraph_pir - def test_dict_input(self): - net = LinearNetWithDictInputNoPrune(8, 8) - temp_dir = tempfile.TemporaryDirectory() - path = os.path.join( - temp_dir.name, "test_jit_save_load_with_dict_input_no_prune/model" - ) - # prune inputs - paddle.jit.save( - layer=net, - path=path, - input_spec=[ - { - 'img': InputSpec( - shape=[None, 8], dtype='float32', name='img' - ), - 'img2': InputSpec( - shape=[None, 8], dtype='float32', name='img2' - ), - } - ], - ) - - img = paddle.randn(shape=[4, 8], dtype='float32') - img2 = paddle.randn(shape=[4, 8], dtype='float32') - loaded_net = paddle.jit.load(path) - loaded_out = loaded_net(img, img2) - - self.assertEqual(len(loaded_net._input_spec()), 2) - temp_dir.cleanup() - - -class TestSaveLoadWithInputSpec(unittest.TestCase): - def setUp(self): - # enable dygraph mode - base.enable_dygraph() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - @test_with_dygraph_pir - def test_with_input_spec(self): - net = LinearNetReturnLoss(8, 8) - # set x.shape = [None, 8] - net.forward = to_static( - net.forward, - input_spec=[InputSpec([None, 8], name='x')], - full_graph=True, - ) - - model_path = os.path.join( - self.temp_dir.name, "input_spec.output_spec/model" - ) - # check inputs and outputs - self.assertTrue(len(net.forward.inputs) == 1) - input_x = net.forward.inputs[0] - if paddle.framework.use_pir_api(): - self.assertTrue(input_x.shape == [-1, 8]) - else: - self.assertTrue(input_x.shape == (-1, 8)) - self.assertTrue(input_x.name == 'x') - - # 1. prune loss - output_spec = net.forward.outputs[:1] - paddle.jit.save(net, model_path, output_spec=output_spec) - - # 2. load to infer - infer_layer = paddle.jit.load(model_path) - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - pred = infer_layer(x) - - @test_with_dygraph_pir - def test_multi_in_out(self): - net = LinearNetMultiInput(8, 8) - net = paddle.jit.to_static( - net, - input_spec=[ - InputSpec([None, 8], dtype='float32'), - InputSpec([None, 8], dtype='float32'), - ], - full_graph=True, - ) - - model_path = os.path.join( - self.temp_dir.name, "multi_inout.output_spec1/model" - ) - # 1. check inputs and outputs - self.assertTrue(len(net.forward.inputs) == 2) - input_x = net.forward.inputs[0] - input_y = net.forward.inputs[1] - if paddle.framework.use_pir_api(): - self.assertTrue(input_x.shape == [-1, 8]) - self.assertTrue(input_y.shape == [-1, 8]) - else: - self.assertTrue(input_x.shape == (-1, 8)) - self.assertTrue(input_y.shape == (-1, 8)) - - # 2. prune loss - output_spec = net.forward.outputs[:2] - paddle.jit.save(net, model_path, output_spec=output_spec) - - # 3. load to infer - infer_layer = paddle.jit.load(model_path) - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - y = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - # 4. predict - pred_x, pred_y = infer_layer(x, y) - - # 1. prune y and loss - model_path = os.path.join( - self.temp_dir.name, "multi_inout.output_spec2/model" - ) - output_spec = net.forward.outputs[:1] - paddle.jit.save(net, model_path, [input_x], output_spec=output_spec) - # 2. load again - infer_layer2 = paddle.jit.load(model_path) - # 3. predict - pred_xx = infer_layer2(x) - - # 4. assert pred_x == pred_xx - np.testing.assert_allclose(pred_x.numpy(), pred_xx.numpy(), rtol=1e-05) - - @test_with_dygraph_pir - def test_multi_in_out1(self): - net = LinearNetMultiInput1(8, 8) - net = paddle.jit.to_static( - net, - input_spec=( - InputSpec([None, 8], dtype='float32'), - InputSpec([None, 8], dtype='float32'), - ), - full_graph=True, - ) - model_path = os.path.join( - self.temp_dir.name, "multi_inout1.output_spec1/model" - ) - # 1. check inputs and outputs - self.assertTrue(len(net.forward.inputs) == 2) - input_x = net.forward.inputs[0] - input_y = net.forward.inputs[1] - if paddle.framework.use_pir_api(): - self.assertTrue(input_x.shape == [-1, 8]) - self.assertTrue(input_y.shape == [-1, 8]) - else: - self.assertTrue(input_x.shape == (-1, 8)) - self.assertTrue(input_y.shape == (-1, 8)) - - # 2. prune loss - output_spec = net.forward.outputs[:2] - paddle.jit.save(net, model_path, output_spec=output_spec) - - # 3. load to infer - infer_layer = paddle.jit.load(model_path) - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - y = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - # 4. predict - pred_x, pred_y = infer_layer(x, y) - - # 1. prune y and loss - model_path = os.path.join( - self.temp_dir.name, "multi_inout1.output_spec2/model" - ) - output_spec = net.forward.outputs[:1] - paddle.jit.save( - net, - model_path, - net.forward.inputs, - output_spec=output_spec, - input_names_after_prune=[input_x.name], - ) - # 2. load again - infer_layer2 = paddle.jit.load(model_path) - # 3. predict - pred_xx = infer_layer2(x) - - # 4. assert pred_x == pred_xx - np.testing.assert_allclose(pred_x.numpy(), pred_xx.numpy(), rtol=1e-05) - - -class TestJitSaveLoadConfig(unittest.TestCase): - def setUp(self): - # enable dygraph mode - base.enable_dygraph() - # config seed - paddle.seed(SEED) - paddle.framework.random._manual_program_seed(SEED) - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_output_spec(self): - train_layer = LinearNetReturnLoss(8, 8) - adam = paddle.optimizer.Adam( - learning_rate=0.1, parameters=train_layer.parameters() - ) - x = paddle.to_tensor(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 = os.path.join( - self.temp_dir.name, "save_load_config.output_spec" - ) - output_spec = [out] - paddle.jit.save( - layer=train_layer, - path=model_path, - input_spec=[x], - output_spec=output_spec, - ) - - train_layer.eval() - infer_layer = paddle.jit.load(model_path) - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - np.testing.assert_array_equal( - train_layer(x)[0].numpy(), infer_layer(x).numpy() - ) - - def test_save_no_support_config_error(self): - layer = LinearNet(784, 1) - path = os.path.join(self.temp_dir.name, "no_support_config_test") - with self.assertRaises(ValueError): - paddle.jit.save(layer=layer, path=path, model_filename="") - - def test_load_empty_model_filename_error(self): - path = os.path.join(self.temp_dir.name, "error_model_filename_test") - with self.assertRaises(ValueError): - paddle.jit.load(path, model_filename="") - - def test_load_empty_params_filename_error(self): - path = os.path.join(self.temp_dir.name, "error_params_filename_test") - with self.assertRaises(ValueError): - paddle.jit.load(path, params_filename="") - - def test_load_with_no_support_config(self): - path = os.path.join(self.temp_dir.name, "no_support_config_test") - with self.assertRaises(ValueError): - paddle.jit.load(path, separate_params=True) - - -class TestJitMultipleLoading(unittest.TestCase): - def setUp(self): - self.linear_size = 4 - self.temp_dir = tempfile.TemporaryDirectory() - self.model_path = os.path.join( - self.temp_dir.name, "jit_multi_load/model" - ) - # enable dygraph mode - base.enable_dygraph() - # config seed - paddle.seed(SEED) - paddle.framework.random._manual_program_seed(SEED) - # train and save base model - self.train_and_save_orig_model() - - def tearDown(self): - self.temp_dir.cleanup() - - def train_and_save_orig_model(self): - layer = LinearNet(self.linear_size, self.linear_size) - example_inputs, layer, _ = train(layer, self.linear_size, 1) - paddle.jit.save( - layer=layer, path=self.model_path, input_spec=example_inputs - ) - - def test_load_model_retransform_inference(self): - multi_loaded_layer = MultiLoadingLinearNet( - self.linear_size, self.model_path - ) - state_dict = multi_loaded_layer.state_dict() - name_set = set() - for _, var in state_dict.items(): - self.assertTrue(var.name not in name_set) - name_set.add(var.name) - - -class TestJitPruneModelAndLoad(unittest.TestCase): - def setUp(self): - self.linear_size = 4 - self.temp_dir = tempfile.TemporaryDirectory() - self.model_path = os.path.join( - self.temp_dir.name, "jit_prune_model_and_load/model" - ) - # enable dygraph mode - base.enable_dygraph() - # config seed - paddle.seed(SEED) - paddle.framework.random._manual_program_seed(SEED) - - def tearDown(self): - self.temp_dir.cleanup() - - def train_and_save(self): - train_layer = LinearNetReturnHidden(8, 8) - adam = paddle.optimizer.Adam( - learning_rate=0.1, parameters=train_layer.parameters() - ) - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - for i in range(10): - hidden, loss = train_layer(x) - loss.backward() - adam.minimize(loss) - train_layer.clear_gradients() - - output_spec = [hidden] - paddle.jit.save( - layer=train_layer, - path=self.model_path, - input_spec=[x], - output_spec=output_spec, - ) - - return train_layer - - def test_load_pruned_model(self): - train_layer = self.train_and_save() - train_layer.eval() - - infer_layer = paddle.jit.load(self.model_path) - - x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) - np.testing.assert_array_equal( - train_layer(x)[0].numpy(), infer_layer(x).numpy() - ) - - def test_load_var_not_in_extra_var_info(self): - self.train_and_save() - - # chage extra var info - var_info_path = self.model_path + INFER_PARAMS_INFO_SUFFIX - with open(var_info_path, 'rb') as f: - extra_var_info = pickle.load(f) - extra_var_info.clear() - with open(var_info_path, 'wb') as f: - pickle.dump(extra_var_info, f, protocol=2) - - with self.assertRaises(RuntimeError): - paddle.jit.load(self.model_path) - - -class TestJitSaveMultiCases(unittest.TestCase): - def setUp(self): - # enable dygraph mode - base.enable_dygraph() - # config seed - paddle.seed(SEED) - paddle.framework.random._manual_program_seed(SEED) - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def verify_inference_correctness( - self, layer, model_path, with_label_and_loss=False, with_label=False - ): - layer.eval() - loaded_layer = paddle.jit.load(model_path) - loaded_layer.eval() - # inference & compare - x = paddle.to_tensor(np.random.random((1, 784)).astype('float32')) - if with_label_and_loss: - y = paddle.to_tensor(np.random.random((1, 1)).astype('int64')) - pred, _ = layer(x, y) - pred = pred.numpy() - elif with_label: - y = paddle.to_tensor(np.random.random((1, 1)).astype('int64')) - pred = layer(x, y) - pred = pred.numpy() - else: - pred = layer(x).numpy() - loaded_pred = loaded_layer(x).numpy() - np.testing.assert_array_equal( - pred, - loaded_pred, - err_msg=f'Result diff when load and inference:\nlayer result:\n{pred}\nloaded layer result:\n{loaded_pred}', - ) - - def test_no_prune_to_static_after_train(self): - layer = LinearNet(784, 1) - - train(layer) - - model_path = os.path.join( - self.temp_dir.name, "test_no_prune_to_static_after_train/model" - ) - paddle.jit.save(layer, model_path) - - self.verify_inference_correctness(layer, model_path) - - def test_no_prune_to_static_no_train(self): - layer = LinearNetWithInputSpec(784, 1) - - model_path = os.path.join( - self.temp_dir.name, "test_no_prune_to_static_no_train/model" - ) - paddle.jit.save(layer, model_path) - - self.verify_inference_correctness(layer, model_path) - - def test_no_prune_no_to_static_after_train(self): - layer = LinearNetNotDeclarative(784, 1) - - train(layer) - - model_path = os.path.join( - self.temp_dir.name, "test_no_prune_no_to_static_after_train/model" - ) - paddle.jit.save( - layer, - model_path, - input_spec=[InputSpec(shape=[None, 784], dtype='float32')], - ) - - self.verify_inference_correctness(layer, model_path) - - def test_no_prune_no_to_static_after_train_with_examples(self): - layer = LinearNetNotDeclarative(784, 1) - - example_inputs, _, _ = train(layer) - - model_path = os.path.join( - self.temp_dir.name, - "test_no_prune_no_to_static_after_train_with_examples/model", - ) - paddle.jit.save(layer=layer, path=model_path, input_spec=example_inputs) - - self.verify_inference_correctness(layer, model_path) - - def test_no_prune_no_to_static_no_train(self): - layer = LinearNetNotDeclarative(784, 1) - - model_path = os.path.join( - self.temp_dir.name, "test_no_prune_no_to_static_no_train/model" - ) - paddle.jit.save( - layer, - model_path, - input_spec=[InputSpec(shape=[None, 784], dtype='float32')], - ) - - self.verify_inference_correctness(layer, model_path) - - def test_prune_to_static_after_train(self): - layer = LinerNetWithLabel(784, 1) - - out = train_with_label(layer) - - model_path = os.path.join( - self.temp_dir.name, "test_prune_to_static_after_train/model" - ) - paddle.jit.save( - layer, - model_path, - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name="image"), - True, - ], - output_spec=[out], - input_names_after_prune=["image"], - ) - - self.verify_inference_correctness( - layer, model_path, with_label_and_loss=True - ) - - def test_prune_to_static_no_train(self): - layer = LinerNetWithLabel(784, 1) - - model_path = os.path.join( - self.temp_dir.name, "test_prune_to_static_no_train/model" - ) - # TODO: no train, cannot get output_spec var here - # now only can use index - output_spec = layer.forward.outputs[:1] - paddle.jit.save( - layer, - model_path, - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name="image"), - True, - ], - output_spec=output_spec, - input_names_after_prune=["image"], - ) - - self.verify_inference_correctness( - layer, model_path, with_label_and_loss=True - ) - - def test_prune_input_to_static_no_train(self): - layer = LinerNetWithPruneInput(784, 1) - - model_path = os.path.join( - self.temp_dir.name, "test_prune_input_to_static_no_train/model" - ) - paddle.jit.save( - layer, - model_path, - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name="image") - ], - ) - - self.verify_inference_correctness(layer, model_path, with_label=True) - - def test_prune_useless_input_to_static_no_train(self): - layer = LinerNetWithUselessInput(784, 1) - - model_path = os.path.join( - self.temp_dir.name, - "test_prune_useless_input_to_static_no_train/model", - ) - paddle.jit.save( - layer, - model_path, - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name="image") - ], - ) - - self.verify_inference_correctness(layer, model_path, with_label=True) - - def test_no_prune_input_spec_name_warning(self): - layer = LinearNetWithInputSpec(784, 1) - - train(layer) - - model_path = os.path.join( - self.temp_dir.name, "test_no_prune_input_spec_name_warning/model" - ) - paddle.jit.save( - layer, - model_path, - input_spec=[InputSpec(shape=[None, 784], dtype='float32')], - ) - paddle.jit.save( - layer, - model_path, - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name='feed_input') - ], - ) - - self.verify_inference_correctness(layer, model_path) - - def test_not_prune_output_spec_name_warning(self): - layer = LinearNet(784, 1) - - train(layer) - - model_path = os.path.join( - self.temp_dir.name, "test_not_prune_output_spec_name_warning/model" - ) - out = paddle.to_tensor(np.random.random((1, 1)).astype('float')) - paddle.jit.save(layer, model_path, output_spec=[out]) - - self.verify_inference_correctness(layer, model_path) - - def test_prune_input_spec_name_error(self): - layer = LinerNetWithLabel(784, 1) - - model_path = os.path.join( - self.temp_dir.name, "test_prune_input_spec_name_error/model" - ) - with self.assertRaises(ValueError): - paddle.jit.save( - layer, - model_path, - input_spec=[InputSpec(shape=[None, 784], dtype='float32')], - ) - with self.assertRaises(ValueError): - paddle.jit.save( - layer, - model_path, - input_spec=[ - InputSpec( - shape=[None, 784], dtype='float32', name='feed_input' - ) - ], - ) - - def test_prune_output_spec_name_error(self): - layer = LinerNetWithLabel(784, 1) - - train_with_label(layer) - - model_path = os.path.join( - self.temp_dir.name, "test_prune_to_static_after_train/model" - ) - out = paddle.to_tensor(np.random.random((1, 1)).astype('float')) - with self.assertRaises(ValueError): - paddle.jit.save( - layer, - model_path, - input_spec=[ - InputSpec(shape=[None, 784], dtype='float32', name="image"), - True, - ], - output_spec=[out], - input_names_after_prune=["image"], - ) - - -class TestJitSaveLoadEmptyLayer(unittest.TestCase): - def setUp(self): - self.temp_dir = tempfile.TemporaryDirectory() - self.model_path = os.path.join( - self.temp_dir.name, "jit_save_load_empty_layer/model" - ) - # enable dygraph mode - paddle.disable_static() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_save_load_empty_layer(self): - layer = EmptyLayer() - x = paddle.to_tensor(np.random.random(10).astype('float32')) - out = layer(x) - paddle.jit.save(layer, self.model_path) - load_layer = paddle.jit.load(self.model_path) - load_out = load_layer(x) - np.testing.assert_array_equal(out, load_out) - - -class TestJitSaveLoadNoParamLayer(unittest.TestCase): - def setUp(self): - self.temp_dir = tempfile.TemporaryDirectory() - self.model_path = os.path.join( - self.temp_dir.name, "jit_save_load_no_param_layer/model" - ) - # enable dygraph mode - paddle.disable_static() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_save_load_no_param_layer(self): - layer = NoParamLayer() - x = paddle.to_tensor(np.random.random(5).astype('float32')) - y = paddle.to_tensor(np.random.random(5).astype('float32')) - out = layer(x, y) - paddle.jit.save(layer, self.model_path) - load_layer = paddle.jit.load(self.model_path) - load_out = load_layer(x, y) - np.testing.assert_array_equal(out, load_out) - - -class TestJitSaveLoadMultiMethods(unittest.TestCase): - def setUp(self): - # enable dygraph mode - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_save_load_inference(self): - model_path_inference = os.path.join( - self.temp_dir.name, "jit_save_load_multi_methods/model" - ) - IMAGE_SIZE = 224 - layer = LinearNetWithMultiStaticFunc(IMAGE_SIZE, 10) - inps = paddle.randn([1, IMAGE_SIZE]) - result_origin = {} - for func in dir(layer): - if func.startswith('forward'): - result_origin[func] = getattr(layer, func, None)(inps) - paddle.jit.save(layer, model_path_inference) - load_net = paddle.jit.load(model_path_inference) - for func, result in result_origin.items(): - self.assertTrue( - float( - (result - getattr(load_net, func, None)(inps)).abs().max() - ) - < 1e-5 - ) - - def test_jit_save_load_multi_methods_inputspec(self): - model_path = os.path.join( - self.temp_dir.name, 'jit_save_load_multi_methods/model' - ) - layer = LinearNetWithMultiStaticFunc(784, 1) - with self.assertRaises(ValueError): - paddle.jit.save( - layer, model_path, input_spec=[InputSpec(shape=[None, 784])] - ) - - def test_parse_name(self): - model_path_inference = os.path.join( - self.temp_dir.name, "jit_save_load_parse_name/model" - ) - IMAGE_SIZE = 224 - layer = LinearNet(IMAGE_SIZE, 1) - inps = paddle.randn([1, IMAGE_SIZE]) - layer(inps) - paddle.jit.save(layer, model_path_inference) - paddle.jit.save(layer, model_path_inference + '_v2') - load_net = paddle.jit.load(model_path_inference) - - self.assertFalse(hasattr(load_net, 'v2')) - - -class LayerSaved(paddle.nn.Layer): - def __init__(self, in_size, out_size): - super().__init__() - self.hidden = 100 - self._linear_0 = Linear(in_size, self.hidden) - self._linear_1_0 = Linear(self.hidden, self.hidden) - self._linear_1_1 = Linear(self.hidden, self.hidden) - self._linear_2 = Linear(self.hidden, out_size) - self._scale = paddle.to_tensor([9.9]) - - @paddle.jit.to_static - def forward(self, x): - y = self._linear_0(x) - # Multiple blocks - if paddle.shape(x)[0] == 1: - y = self._linear_1_0(y) - else: - y += self._linear_1_1(y + self._scale) - return self._linear_2(y) - - -class Net(paddle.nn.Layer): - def __init__(self): - super().__init__() - self.fc1 = paddle.nn.Linear(4, 4) - self.fc2 = paddle.nn.Linear(4, 4) - self.bias = 0.4 - self.flag = paddle.ones([2], dtype="int32") - - @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) - def log_softmax(self, input): - return paddle.nn.functional.log_softmax(input, axis=-1) - - @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) - def forward(self, x): - out = self.fc1(x) - out = paddle.nn.functional.relu(out) - out = paddle.mean(out) - return out - - @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) - def infer(self, input): - out = self.fc2(input) - out = out + self.bias - out = paddle.mean(out) - return out - - # For extra Python float - @paddle.jit.to_static(property=True) - def fbias(self): - return self.bias + 1 - - @paddle.jit.to_static(property=True) - def down_sampling(self): - return 4 - - @paddle.jit.to_static(property=True) - def fstr(self): - return "save str property" - - @paddle.jit.to_static(property=True) - def ints(self): - return [10, 20] - - @paddle.jit.to_static(property=True) - def floats(self): - return [1.1, 2.2] - - @paddle.jit.to_static(property=True) - def strs(self): - return ["hello", "world"] - - -class NetTensor(paddle.nn.Layer): - def __init__(self): - super().__init__() - self.fc1 = paddle.nn.Linear(4, 4) - self.fc2 = paddle.nn.Linear(4, 4) - self.bias = 0.4 - self.flag = paddle.ones([2], dtype="int32") - - @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) - def forward(self, x): - out = self.fc1(x) - out = paddle.nn.functional.relu(out) - out = paddle.mean(out) - return out - - @paddle.jit.to_static(property=True) - def fflag(self): - return True - - -class TestJitSaveCombineProperty(unittest.TestCase): - def setUp(self): - # enable dygraph mode - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_save_combine_property(self): - model_path = os.path.join( - self.temp_dir.name, "test_jit_save_combine/model" - ) - # Use new namespace - with unique_name.guard(): - net = Net() - # save - paddle.jit.save(net, model_path, combine_params=True) - - def test_jit_save_tensor_property(self): - model_path = os.path.join( - self.temp_dir.name, "test_jit_save_combine/model" - ) - # Use new namespace - with unique_name.guard(): - net = NetTensor() - - paddle.jit.save(net, model_path, combine_params=True) - - -class LayerLoadFinetune(paddle.nn.Layer): - def __init__(self, in_size, out_size, load_path): - super().__init__() - # Test duplicate name - self._linear_0 = Linear(in_size, in_size) - self._linear_1_0 = Linear(out_size, in_size) - self._linear_1_1 = Linear(out_size, in_size) - self._linear_2 = Linear(out_size, out_size) - self._scale = paddle.to_tensor([9.9]) - - # Load multiple times - self._load_l1 = paddle.jit.load(load_path) - self._load_l2 = paddle.jit.load(load_path) - - @paddle.jit.to_static - def forward(self, x): - y = self._linear_0(x) - y = self._load_l1(y) - # Multiple blocks - if paddle.shape(x)[0] == 1: - y = self._linear_1_0(y) - y = self._load_l1(y) - else: - y += self._linear_1_1(x + self._scale) - y = self._load_l2(y) - y = self._linear_1_0(y) - y = self._load_l1(y) - y = self._linear_1_0(y) - # Use the same layer multiple times. - y = self._load_l1(y) - return y - - -class TestJitSaveLoadSaveWithoutRunning(unittest.TestCase): - def setUp(self): - # enable dygraph mode - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_save_load_finetune_load(self): - model_path = os.path.join( - self.temp_dir.name, "test_jit_save_load_save_without_running/model" - ) - IMAGE_SIZE = 224 - inps0 = paddle.randn([1, IMAGE_SIZE]) - inps1 = paddle.randn([2, IMAGE_SIZE]) - # Use new namespace - with unique_name.guard(): - layer_save = LayerSaved(IMAGE_SIZE, IMAGE_SIZE) - # save - paddle.jit.save( - layer_save, - model_path, - input_spec=[ - paddle.static.InputSpec( - shape=[None, IMAGE_SIZE], dtype='float32' - ) - ], - ) - result_00 = layer_save(inps0) - result_01 = layer_save(inps1) - # load and save without running - with unique_name.guard(): - layer_load = paddle.jit.load(model_path) - paddle.jit.save( - layer_load, - model_path, - input_spec=[ - paddle.static.InputSpec( - shape=[None, IMAGE_SIZE], dtype='float32' - ) - ], - ) - # reload - layer_reload = paddle.jit.load(model_path) - result_10 = layer_reload(inps0) - result_11 = layer_reload(inps1) - - self.assertTrue(float((result_00 - result_10).abs().max()) < 1e-5) - self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5) - - -class TestJitSaveLoadFinetuneLoad(unittest.TestCase): - def setUp(self): - # enable dygraph mode - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_save_load_finetune_load(self): - model_path = os.path.join( - self.temp_dir.name, "test_jit_save_load_finetune_load/model" - ) - IMAGE_SIZE = 224 - inps0 = paddle.randn([1, IMAGE_SIZE]) - inps1 = paddle.randn([2, IMAGE_SIZE]) - # Use new namespace - with unique_name.guard(): - layer_save = LayerSaved(IMAGE_SIZE, IMAGE_SIZE) - layer_save(inps0) - # save - paddle.jit.save(layer_save, model_path) - # load - with unique_name.guard(): - layer_load = LayerLoadFinetune(IMAGE_SIZE, IMAGE_SIZE, model_path) - # train - train(layer_load, input_size=IMAGE_SIZE) - result_00 = layer_load(inps0) - result_01 = layer_load(inps1) - # save - paddle.jit.save(layer_load, model_path) - # load - layer_finetune = paddle.jit.load(model_path) - result_10 = layer_finetune(inps0) - result_11 = layer_finetune(inps1) - - self.assertTrue(float((result_00 - result_10).abs().max()) < 1e-5) - self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5) - - -# NOTE(weixin): When there are multiple test functions in an -# `unittest.TestCase`, functions will affect each other, -# and there is a risk of random failure. -# So divided into three TestCase: TestJitSaveLoadFunctionCase1, -# TestJitSaveLoadFunctionCase2, TestJitSaveLoadFunctionCase3. -class TestJitSaveLoadFunctionCase1(unittest.TestCase): - def setUp(self): - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_save_load_static_function(self): - @paddle.jit.to_static - def fun(inputs): - return paddle.tanh(inputs) - - path = os.path.join( - self.temp_dir.name, 'test_jit_save_load_function_1/func' - ) - inps = paddle.rand([3, 6]) - origin = fun(inps) - - paddle.jit.save(fun, path) - load_func = paddle.jit.load(path) - - load_result = load_func(inps) - self.assertTrue((load_result - origin).abs().max() < 1e-10) - - -class TestJitSaveLoadFunctionCase2(unittest.TestCase): - def setUp(self): - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_save_load_function_input_spec(self): - @paddle.jit.to_static( - input_spec=[ - InputSpec(shape=[None, 6], dtype='float32', name='x'), - ] - ) - def fun(inputs): - return paddle.nn.functional.relu(inputs) - - path = os.path.join( - self.temp_dir.name, 'test_jit_save_load_function_2/func' - ) - inps = paddle.rand([3, 6]) - origin = fun(inps) - - paddle.jit.save(fun, path) - load_func = paddle.jit.load(path) - load_result = load_func(inps) - self.assertTrue((load_result - origin).abs().max() < 1e-10) - - -class TestJitSaveLoadFunctionCase3(unittest.TestCase): - def setUp(self): - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_save_load_function_function(self): - def fun(inputs): - return paddle.tanh(inputs) - - path = os.path.join( - self.temp_dir.name, 'test_jit_save_load_function_3/func' - ) - inps = paddle.rand([3, 6]) - origin = fun(inps) - - paddle.jit.save( - fun, - path, - input_spec=[ - InputSpec(shape=[None, 6], dtype='float32', name='x'), - ], - ) - load_func = paddle.jit.load(path) - - load_result = load_func(inps) - self.assertTrue((load_result - origin).abs().max() < 1e-10) - - -class TestJitSaveLoadFunctionWithParamCase1(unittest.TestCase): - def setUp(self): - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_save_load_function(self): - class LinearNet(paddle.nn.Layer): - def __init__(self): - super().__init__() - self._linear = paddle.nn.Linear(5, 6) - - def forward(self, x): - return paddle.tanh(x) - - def anothor_forward(self, x): - return self._linear(x) - - layer = LinearNet() - - inps = paddle.rand([3, 5]) - origin = layer.anothor_forward(inps) - - func = paddle.jit.to_static( - layer.anothor_forward, [paddle.static.InputSpec(shape=[-1, 5])] - ) - path = os.path.join( - self.temp_dir.name, - 'test_jit_save_load_function_with_params_case1/func', - ) - paddle.jit.save(func, path) - load_func = paddle.jit.load(path) - - load_result = load_func(inps) - np.testing.assert_array_equal(load_result.numpy(), origin.numpy()) - - -class TestJitSaveLoadFunctionWithParamCase2(unittest.TestCase): - def setUp(self): - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_save_load_function(self): - class LinearNet(paddle.nn.Layer): - def __init__(self): - super().__init__() - self._linear = paddle.nn.Linear(5, 6) - - def forward(self, x): - return paddle.tanh(x) - - @paddle.jit.to_static(input_spec=[InputSpec(shape=[-1, 5])]) - def anothor_forward(self, x): - return self._linear(x) - - layer = LinearNet() - - inps = paddle.rand([3, 5]) - - path = os.path.join( - self.temp_dir.name, - 'test_jit_save_load_function_with_params_case2/func', - ) - paddle.jit.save(layer.anothor_forward, path) - origin_result = layer.anothor_forward(inps) - load_func = paddle.jit.load(path) - - load_result = load_func(inps) - - np.testing.assert_array_equal( - origin_result.numpy(), load_result.numpy() - ) - - -class TestJitSaveLoadFunctionWithParamCase3(unittest.TestCase): - def setUp(self): - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_save_load_function(self): - class LinearNet(paddle.nn.Layer): - def __init__(self): - super().__init__() - self._linear = paddle.nn.Linear(5, 6) - - def forward(self, x): - return paddle.tanh(x) - - @paddle.jit.to_static - def anothor_forward(self, x): - return self._linear(x) - - layer = LinearNet() - - inps = paddle.rand([3, 5]) - origin = layer.anothor_forward(inps) - - path = os.path.join( - self.temp_dir.name, - 'test_jit_save_load_function_with_params_case3/func', - ) - paddle.jit.save(layer.anothor_forward, path) - load_func = paddle.jit.load(path) - - load_result = load_func(inps) - np.testing.assert_array_equal(load_result.numpy(), origin.numpy()) - - -class TestJitSaveLoadDataParallel(unittest.TestCase): - def setUp(self): - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def verify_inference_correctness(self, layer, path): - layer.eval() - loaded_layer = paddle.jit.load(path) - loaded_layer.eval() - # inference & compare - x = paddle.to_tensor(np.random.random((1, 784)).astype('float32')) - pred = layer(x).numpy() - loaded_pred = loaded_layer(x).numpy() - np.testing.assert_array_equal( - pred, - loaded_pred, - err_msg=f'Result diff when load and inference:\nlayer result:\n{pred}\nloaded layer result:\n{loaded_pred}', - ) - - def test_jit_save_data_parallel_with_inputspec(self): - layer = LinearNetNotDeclarative(784, 1) - layer = paddle.DataParallel(layer) - path = os.path.join( - self.temp_dir.name, "jit_save_data_parallel_with_inputspec/model" - ) - paddle.jit.save( - layer=layer, path=path, input_spec=[InputSpec(shape=[None, 784])] - ) - - self.verify_inference_correctness(layer, path) - - def test_jit_save_data_parallel_with_to_static(self): - layer = LinearNetWithInputSpec(784, 1) - layer = paddle.DataParallel(layer) - - path = os.path.join( - self.temp_dir.name, "jit_save_data_parallel_with_to_static/model" - ) - paddle.jit.save(layer, path) - - self.verify_inference_correctness(layer, path) - - -class InputSepcLayer(paddle.nn.Layer): - ''' - A layer with InputSpec to test InputSpec compatibility - ''' - - @paddle.jit.to_static( - input_spec=[ - InputSpec(shape=[None, 8], dtype='float32', name='x'), - InputSpec(shape=[None, 1], dtype='float64', name='y'), - ] - ) - def forward(self, x, y): - return x, y - - -class TestInputSpecCompatibility(unittest.TestCase): - def setUp(self): - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def _assert_input_spec_layer_return(self, expect_layer, test_layer): - input_x = paddle.uniform([8, 8], dtype='float32') - input_y = paddle.uniform([8, 1], dtype='float64') - expected_result = expect_layer(input_x, input_y) - test_result = test_layer(input_x, input_y) - np.testing.assert_allclose( - expected_result[0].numpy(), test_result[0].numpy() - ) - np.testing.assert_allclose( - expected_result[1].numpy(), test_result[1].numpy() - ) - - def test_jit_save_compatible_input_sepc(self): - layer = InputSepcLayer() - save_dir = os.path.join( - self.temp_dir.name, "jit_save_compatible_input_spec" - ) - path = save_dir + "/model" - - paddle.jit.save(layer=layer, path=path) - no_input_spec_layer = paddle.jit.load(path) - self._assert_input_spec_layer_return(layer, no_input_spec_layer) - shutil.rmtree(save_dir) - - paddle.jit.save( - layer=layer, - path=path, - input_spec=[ - InputSpec(shape=[None, 8], dtype='float32', name='x'), - InputSpec(shape=[None, 1], dtype='float64', name='y'), - ], - ) - same_input_spec_layer = paddle.jit.load(path) - self._assert_input_spec_layer_return(layer, same_input_spec_layer) - shutil.rmtree(save_dir) - - paddle.jit.save( - layer=layer, - path=path, - input_spec=[ - InputSpec(shape=[8, 8], dtype='float32'), - InputSpec(shape=[8, -1], dtype='float64'), - ], - ) - compatible_input_spec_layer = paddle.jit.load(path) - self._assert_input_spec_layer_return(layer, compatible_input_spec_layer) - shutil.rmtree(save_dir) - - def test_jit_save_incompatible_input_sepc(self): - layer = InputSepcLayer() - save_dir = os.path.join( - self.temp_dir.name, "jit_save_compatible_input_spec" - ) - path = save_dir + "/model" - - with self.assertRaises(ValueError): - # type mismatch - paddle.jit.save( - layer=layer, - path=path, - input_spec=[ - InputSpec(shape=[None, 8], dtype='float64'), - InputSpec(shape=[None, 1], dtype='float64'), - ], - ) - - with self.assertRaises(ValueError): - # shape len mismatch - paddle.jit.save( - layer=layer, - path=path, - input_spec=[ - InputSpec(shape=[None, 8, 1], dtype='float32'), - InputSpec(shape=[None, 1], dtype='float64'), - ], - ) - - with self.assertRaises(ValueError): - # shape mismatch - paddle.jit.save( - layer=layer, - path=path, - input_spec=[ - InputSpec(shape=[None, 8], dtype='float32'), - InputSpec(shape=[None, 2], dtype='float64'), - ], - ) - if os.path.exists(save_dir): - shutil.rmtree(save_dir) - - -class NotJitForward(paddle.nn.Layer): - def __init__(self): - super().__init__() - - def forward(self, x, y): - return x + y - - -class TestNotJitForward(unittest.TestCase): - def setUp(self): - self.temp_dir = tempfile.TemporaryDirectory() - - def tearDown(self): - self.temp_dir.cleanup() - - def test_jit_not_save_forward(self): - layer = NotJitForward() - - save_dir = os.path.join(self.temp_dir.name, "jit_not_save_forward") - path = save_dir + "/model" - - paddle.jit.save(layer=layer, path=path, skip_forward=True) - - self.assertTrue(not os.path.exists(path + ".pdmodel")) - self.assertTrue(not os.path.exists(path + ".pdparam")) - - with self.assertRaises(ValueError): - paddle.jit.load(path=path) - - shutil.rmtree(save_dir) - - -if __name__ == '__main__': - unittest.main() From 266ad66fe0db978774d9a55cefb78cb3309eb0eb Mon Sep 17 00:00:00 2001 From: wangruting Date: Mon, 29 Apr 2024 11:05:36 +0000 Subject: [PATCH 07/15] rename jit_save_load --- test/legacy_test/test_jit_save_load_rename.py | 2205 +++++++++++++++++ 1 file changed, 2205 insertions(+) create mode 100644 test/legacy_test/test_jit_save_load_rename.py diff --git a/test/legacy_test/test_jit_save_load_rename.py b/test/legacy_test/test_jit_save_load_rename.py new file mode 100644 index 00000000000000..789d4277a7b318 --- /dev/null +++ b/test/legacy_test/test_jit_save_load_rename.py @@ -0,0 +1,2205 @@ +# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. +# 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. + +import copy +import os +import pickle +import tempfile +import unittest + +import numpy as np + +import paddle +from paddle import base +from paddle.jit.api import to_static +from paddle.jit.translated_layer import INFER_PARAMS_INFO_SUFFIX +from paddle.nn import Linear +from paddle.pir_utils import test_with_dygraph_pir +from paddle.static import InputSpec + +BATCH_SIZE = 32 +BATCH_NUM = 10 +SEED = 10 + + +def random_batch_reader(input_size, label_size): + def _get_random_inputs_and_labels(input_size, label_size): + np.random.seed(SEED) + input = np.random.random(size=input_size).astype('float32') + label = np.random.random(size=label_size).astype('int64') + return input, label + + def __reader__(): + for _ in range(BATCH_NUM): + batch_input, batch_label = _get_random_inputs_and_labels( + [BATCH_SIZE, input_size], [BATCH_SIZE, label_size] + ) + yield batch_input, batch_label + + return __reader__ + + +class LinearNet(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + @to_static + def forward(self, x): + return self._linear(x) + + +class LinearNetWithInputSpec(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + @to_static( + input_spec=[InputSpec(shape=[None, 784], dtype='float32')], + full_graph=True, + ) + def forward(self, x): + return self._linear(x) + + +class LinearNetNotDeclarative(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + def forward(self, x): + return self._linear(x) + + +class LinerNetWithLabel(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + def forward(self, x, label): + out = self._linear(x) + loss = paddle.nn.functional.cross_entropy( + out, label, reduction='none', use_softmax=False + ) + avg_loss = paddle.mean(loss) + return out, avg_loss + + +class LinerNetWithPruneInput(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + def forward(self, x, label): + out = self._linear(x) + loss = paddle.nn.functional.cross_entropy( + out, label, reduction='none', use_softmax=False + ) + avg_loss = paddle.mean(loss) + return out + + +class LinerNetWithUselessInput(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + def forward(self, x, label): + out = self._linear(x) + return out + + +class LinearNetReturnLoss(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + @to_static + def forward(self, x): + y = self._linear(x) + z = self._linear(y) + loss = paddle.mean(z) + return z, loss + + +class LinearNetMultiInput(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear1 = Linear(in_size, out_size) + self._linear2 = Linear(in_size, out_size) + + def forward(self, x, y): + x_out = self._linear1(x) + y_out = self._linear2(y) + loss = paddle.mean(x_out + y_out) + return x_out, y_out, loss + + +class LinearNetMultiInput1(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear1 = Linear(in_size, out_size) + self._linear2 = Linear(in_size, out_size) + + def forward(self, x, y): + x_out = self._linear1(x) + y_out = self._linear2(y) + loss = paddle.mean(x_out + y_out) + return x_out, y_out, loss + + +class MultiLoadingLinearNet(paddle.nn.Layer): + def __init__(self, size, model_path): + super().__init__() + self._linear = Linear(size, size) + self._load_linear1 = paddle.jit.load(model_path) + self._load_linear2 = paddle.jit.load(model_path) + + @to_static + def forward(self, x): + tmp1 = self._linear(x) + tmp2 = self._load_linear1(tmp1) + tmp3 = self._load_linear2(tmp2) + y = self._linear(tmp3) + return y + + +class LinearNetReturnHidden(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear_1 = Linear(in_size, out_size) + self._linear_2 = Linear(in_size, out_size) + + @to_static + def forward(self, x): + y = self._linear_1(x) + z = self._linear_2(y) + loss = paddle.mean(z) + return y, loss + + +class LinearNetWithNestOut(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear_1 = Linear(in_size, out_size) + self._linear_2 = Linear(in_size, out_size) + + @to_static + def forward(self, x): + y = self._linear_1(x) + z = self._linear_2(y) + out = y + z + loss = paddle.mean(out) + return y, [(z, loss), out] + + +class LinearNetWithDictInput(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + def forward(self, img, label): + out = self._linear(img['img']) + # not return loss to avoid prune output + loss = paddle.nn.functional.cross_entropy(out, label['label']) + return out + + +class LinearNetWithDictInputNoPrune(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear = Linear(in_size, out_size) + + def forward(self, img): + out = self._linear(img['img'] + img['img2']) + return out + + +class EmptyLayer(paddle.nn.Layer): + def __init__(self): + super().__init__() + + @paddle.jit.to_static + def forward(self, x): + return x + + +class NoParamLayer(paddle.nn.Layer): + def __init__(self): + super().__init__() + + @paddle.jit.to_static + def forward(self, x, y): + return x + y + + +class LinearNetWithMultiStaticFunc(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear_0 = Linear(in_size, out_size) + self._linear_1 = Linear(in_size, out_size) + self._scale = paddle.to_tensor([9.9]) + + def forward(self, x): + return self._linear_0(x) + + def forward_no_param(self, x): + return x + + def forward_general(self, x): + return self._linear_0(x) + self._linear_1(x) * self._scale + + +class LinearNetWithNonLexicographicalOrderDict(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear_u = Linear(in_size, out_size) + self._linear_v = Linear(in_size, out_size) + self._linear_w = Linear(in_size, out_size) + self._linear_p = Linear(in_size, out_size) + + def forward(self, x): + u = self._linear_u(x) + v = self._linear_v(x) + w = self._linear_w(x) + p = self._linear_p(x) + return { + "u": u, + "v": v, + "w": w, + "p": p, + } + + +class LinearNetWithNestedNonLexicographicalOrderDict(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self._linear_u = Linear(in_size, out_size) + self._linear_v = Linear(in_size, out_size) + self._linear_w = Linear(in_size, out_size) + self._linear_p = Linear(in_size, out_size) + self._linear_y = Linear(in_size, out_size) + self._linear_x = Linear(in_size, out_size) + + def forward(self, x_): + u = self._linear_u(x_) + v = self._linear_v(x_) + w = self._linear_w(x_) + p = self._linear_p(x_) + + x = self._linear_p(x_) + y = self._linear_p(x_) + return { + "u": u, + "v": v, + "w": w, + "p": p, + "a": { + "x": x, + "y": y, + }, + } + + +def train(layer, input_size=784, label_size=1): + # create optimizer + sgd = paddle.optimizer.SGD( + learning_rate=0.01, parameters=layer.parameters() + ) + # create data loader + train_loader = base.io.DataLoader.from_generator(capacity=5) + train_loader.set_batch_generator( + random_batch_reader(input_size, label_size) + ) + # train + for data in train_loader(): + img, label = data + label.stop_gradient = True + + cost = layer(img) + + loss = paddle.nn.functional.cross_entropy( + cost, label, reduction='none', use_softmax=True + ) + avg_loss = paddle.mean(loss) + + avg_loss.backward() + sgd.minimize(avg_loss) + layer.clear_gradients() + return [img], layer, avg_loss + + +def train_with_label(layer, input_size=784, label_size=1): + # create optimizer + sgd = paddle.optimizer.SGD( + learning_rate=0.01, parameters=layer.parameters() + ) + # create data loader + train_loader = base.io.DataLoader.from_generator(capacity=5) + train_loader.set_batch_generator( + random_batch_reader(input_size, label_size) + ) + # train + for data in train_loader(): + img, label = data + label.stop_gradient = True + + out, avg_loss = layer(img, label) + + avg_loss.backward() + sgd.minimize(avg_loss) + layer.clear_gradients() + return out + + +class TestJitSaveLoad(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.model_path = os.path.join( + self.temp_dir.name, "test_jit_save_load/model" + ) + # enable dygraph mode + base.enable_dygraph() + # config seed + paddle.seed(SEED) + paddle.framework.random._manual_program_seed(SEED) + + def tearDown(self): + self.temp_dir.cleanup() + + def train_and_save_model(self, model_path=None): + layer = LinearNet(784, 1) + example_inputs, layer, _ = train(layer) + final_model_path = model_path if model_path else self.model_path + orig_input_types = [type(x) for x in example_inputs] + paddle.jit.save( + layer=layer, path=final_model_path, input_spec=example_inputs + ) + new_input_types = [type(x) for x in example_inputs] + self.assertEqual(orig_input_types, new_input_types) + return layer + + @test_with_dygraph_pir + def test_save_load(self): + # train and save model + train_layer = self.train_and_save_model() + # load model + loaded_layer = paddle.jit.load(self.model_path) + self.load_and_inference(train_layer, loaded_layer) + self.load_and_finetune(train_layer, loaded_layer) + if not paddle.framework.use_pir_api(): + self.load_dygraph_state_dict(train_layer) + + def load_and_inference(self, train_layer, infer_layer): + train_layer.eval() + infer_layer.eval() + # inference & compare + x = paddle.to_tensor(np.random.random((1, 784)).astype('float32')) + np.testing.assert_array_equal( + train_layer(x).numpy(), infer_layer(x).numpy() + ) + + def load_and_finetune(self, train_layer, load_train_layer): + train_layer.train() + load_train_layer.train() + # train & compare + img0, _, train_loss = train(train_layer) + img1, _, load_train_loss = train(load_train_layer) + np.testing.assert_array_equal( + train_loss.numpy(), load_train_loss.numpy() + ) + + def load_dygraph_state_dict(self, train_layer): + train_layer.eval() + # construct new model + new_layer = LinearNet(784, 1) + orig_state_dict = new_layer.state_dict() + load_state_dict = paddle.load(self.model_path) + for structured_name in orig_state_dict: + self.assertTrue(structured_name in load_state_dict) + new_layer.set_state_dict(load_state_dict) + new_layer.eval() + # inference & compare + x = paddle.to_tensor(np.random.random((1, 784)).astype('float32')) + np.testing.assert_array_equal( + train_layer(x).numpy(), new_layer(x).numpy() + ) + + def test_load_dygraph_no_path(self): + model_path = os.path.join( + self.temp_dir.name, "test_jit_save_load.no_path/model_path" + ) + with self.assertRaises(ValueError): + model_dict = paddle.load(model_path) + + @test_with_dygraph_pir + def test_jit_load_no_path(self): + path = os.path.join( + self.temp_dir.name, "test_jit_save_load.no_path/model_path" + ) + with self.assertRaises(ValueError): + loaded_layer = paddle.jit.load(path) + + +class TestSaveLoadWithNestOut(unittest.TestCase): + def setUp(self): + # enable dygraph mode + base.enable_dygraph() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + @test_with_dygraph_pir + def test_nest_output(self): + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + + net = LinearNetWithNestOut(8, 8) + dy_outs = paddle.utils.flatten(net(x)) + net = to_static(net, input_spec=[InputSpec([None, 8], name='x')]) + + model_path = os.path.join(self.temp_dir.name, "net_with_nest_out/model") + paddle.jit.save(net, model_path) + + load_net = paddle.jit.load(model_path) + load_outs = paddle.utils.flatten(load_net(x)) + + self.assertTrue(len(dy_outs) == 4) + for dy_out, load_out in zip(dy_outs, load_outs): + np.testing.assert_allclose( + dy_out.numpy(), load_out.numpy(), rtol=1e-05 + ) + + +class TestSaveLoadWithNonLexicographicalOrderDict(unittest.TestCase): + def setUp(self): + # enable dygraph mode + base.enable_dygraph() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_output_same_order(self): + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + + model = LinearNetWithNonLexicographicalOrderDict(8, 8) + + dy_output_dict = model(x) + + st_model = paddle.jit.to_static(model) + st_output_dict = st_model(x) + + paddle.jit.save(st_model, "./test_jit_save_load") + loaded_model = paddle.jit.load("./test_jit_save_load") + loaded_output_seq = loaded_model(x) + + self.assertTrue(len(dy_output_dict) == 4) + self.assertTrue(len(st_output_dict) == 4) + self.assertTrue(len(loaded_output_seq) == 4) + + # 1. check whether output dict of dygraph and static graph is same + for (dy_key, dy_out), (st_key, st_out) in zip( + dy_output_dict.items(), st_output_dict.items() + ): + self.assertTrue(dy_key == st_key) + np.testing.assert_allclose( + dy_out.numpy(), st_out.numpy(), rtol=1e-05 + ) + + # 2. check whether flattened output has same order of original dict + dy_output_seq = paddle.utils.flatten(dy_output_dict) + + self.assertTrue(len(dy_output_seq) == 4) + for dy_out, flattened_dy_out in zip( + dy_output_dict.values(), dy_output_seq + ): + np.testing.assert_allclose( + dy_out.numpy(), flattened_dy_out.numpy(), rtol=1e-05 + ) + + # 3. check whether flattened output of loaded static graph has same order of dynamic's + for dy_out, loaded_out in zip( + dy_output_dict.values(), loaded_output_seq + ): + np.testing.assert_allclose( + dy_out.numpy(), loaded_out.numpy(), rtol=1e-05 + ) + + +class TestSaveLoadWithNestedNonLexicographicalOrderDict(unittest.TestCase): + def setUp(self): + # enable dygraph mode + base.enable_dygraph() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_nested_output_same_order(self): + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + + model = LinearNetWithNestedNonLexicographicalOrderDict(8, 8) + + dy_output_dict = model(x) + + st_model = paddle.jit.to_static(model) + st_output_dict = st_model(x) + + paddle.jit.save(st_model, "./test_jit_save_load2") + loaded_model = paddle.jit.load("./test_jit_save_load2") + loaded_output_seq = loaded_model(x) + + self.assertTrue(len(dy_output_dict) == 5) + self.assertTrue(len(st_output_dict) == 5) + self.assertTrue(len(loaded_output_seq) == 6) + + # check whether flattened output of loaded static graph has same order of dynamic's + dy_output_tensors = [] + for v in dy_output_dict.values(): + if isinstance(v, dict): + dy_output_tensors.extend(list(v.values())) + else: + dy_output_tensors.append(v) + + for dy_out, loaded_out in zip(dy_output_tensors, loaded_output_seq): + np.testing.assert_allclose( + dy_out.numpy(), loaded_out.numpy(), rtol=1e-05 + ) + + +class TestUtilsMapAndPack(unittest.TestCase): + def setUp(self): + # enable dygraph mode + base.enable_dygraph() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_utils_map_structure(self): + nested_list = [ + { + "d": paddle.to_tensor([1.0]), + "a": paddle.to_tensor([2.0]), + "c": paddle.to_tensor([3.0]), + "tmp": { + "b": paddle.to_tensor([4.0]), + }, + }, + [paddle.to_tensor([5.0]), paddle.to_tensor([6.0])], + [], + [ + paddle.to_tensor([7.0]), + [ + paddle.to_tensor([8.0]), + [paddle.to_tensor([9.0]), [paddle.to_tensor([10.0])]], + ], + ], + ] + FACTOR = 2 + expected_list = [ + { + "d": paddle.to_tensor([1.0]) * FACTOR, + "a": paddle.to_tensor([2.0]) * FACTOR, + "c": paddle.to_tensor([3.0]) * FACTOR, + "tmp": { + "b": paddle.to_tensor([4.0]) * FACTOR, + }, + }, + [ + paddle.to_tensor([5.0]) * FACTOR, + paddle.to_tensor([6.0]) * FACTOR, + ], + [], + [ + paddle.to_tensor([7.0]) * FACTOR, + [ + paddle.to_tensor([8.0]) * FACTOR, + [ + paddle.to_tensor([9.0]) * FACTOR, + [paddle.to_tensor([10.0]) * FACTOR], + ], + ], + ], + ] + mapped_list = paddle.utils.map_structure( + lambda x: x * FACTOR, nested_list + ) + + # test paddle.utils. + def dfs(obj1, obj2): + self.assertTrue(type(obj1) == type(obj2)) + if isinstance(obj1, list): + for i in range(len(obj1)): + dfs(obj1[i], obj2[i]) + elif isinstance(obj1, dict): + self.assertTrue(list(obj1.keys()) == list(obj2.keys())) + for k in obj1: + dfs(obj1[k], obj2[k]) + elif isinstance(obj1, paddle.Tensor): + np.testing.assert_allclose( + obj1.numpy(), obj2.numpy(), rtol=1e-05 + ) + else: + raise ValueError(f"Unsupported type: {type(obj1)} in dfs") + + dfs(expected_list, mapped_list) + + def test_utils_pack_sequence_as(self): + nested_list = [ + { + "d": paddle.to_tensor([1.0]), + "a": paddle.to_tensor([2.0]), + "c": paddle.to_tensor([3.0]), + "tmp": { + "b": paddle.to_tensor([4.0]), + }, + }, + [paddle.to_tensor([5.0]), paddle.to_tensor([6.0])], + [], + [ + paddle.to_tensor([7.0]), + [ + paddle.to_tensor([8.0]), + [paddle.to_tensor([9.0]), [paddle.to_tensor([10.0])]], + ], + ], + ] + + def dfs(obj1, obj2): + self.assertTrue(type(obj1) == type(obj2)) + if isinstance(obj1, list): + for i in range(len(obj1)): + dfs(obj1[i], obj2[i]) + elif isinstance(obj1, dict): + self.assertTrue(list(obj1.keys()) == list(obj2.keys())) + for k in obj1: + dfs(obj1[k], obj2[k]) + elif isinstance(obj1, paddle.Tensor): + np.testing.assert_allclose( + obj1.numpy(), obj2.numpy(), rtol=1e-05 + ) + else: + raise ValueError(f"Unsupported type: {type(obj1)} in dfs") + + nested_list_copy = copy.deepcopy(nested_list) + nested_list_copy_pack_back = paddle.utils.pack_sequence_as( + nested_list_copy, paddle.utils.flatten(nested_list) + ) + + dfs(nested_list_copy, nested_list_copy_pack_back) + + +class TestSaveLoadWithDictInput(unittest.TestCase): + @test_with_dygraph_pir + def test_dict_input(self): + # NOTE: This net cannot be executed, it is just + # a special case for exporting models in model validation + # We DO NOT recommend this writing way of Layer + net = LinearNetWithDictInput(8, 8) + net = paddle.jit.to_static( + net, + input_spec=[ + { + 'img': InputSpec( + shape=[None, 8], dtype=paddle.float32, name='img' + ) + }, + { + 'label': InputSpec( + shape=[None, 1], dtype=paddle.int64, name='label' + ) + }, + ], + full_graph=True, + ) + # net.forward.concrete_program.inputs: + # (<__main__.LinearNetWithDictInput object at 0x7f2655298a98>, + # {'img': var img : base.VarType.LOD_TENSOR.shape(-1, 8).astype(VarType.FP32)}, + # {'label': var label : base.VarType.LOD_TENSOR.shape(-1, 1).astype(VarType.INT64)}) + self.assertEqual(len(net.forward.concrete_program.inputs), 3) + temp_dir = tempfile.TemporaryDirectory() + path = os.path.join( + temp_dir.name, "test_jit_save_load_with_dict_input/model" + ) + # prune inputs + paddle.jit.save( + layer=net, + path=path, + input_spec=[ + { + 'img': InputSpec( + shape=[None, 8], dtype=paddle.float32, name='img' + ) + } + ], + ) + + img = paddle.randn(shape=[4, 8], dtype='float32') + loaded_net = paddle.jit.load(path) + loaded_out = loaded_net(img) + + # loaded_net._input_spec(): + # [InputSpec(shape=(-1, 8), dtype=VarType.FP32, name=img)] + self.assertEqual(len(loaded_net._input_spec()), 1) + self.assertEqual(len(loaded_net._output_spec()), 1) + temp_dir.cleanup() + + +class TestSaveLoadWithDictInputNoPrune(unittest.TestCase): + @test_with_dygraph_pir + def test_dict_input(self): + net = LinearNetWithDictInputNoPrune(8, 8) + temp_dir = tempfile.TemporaryDirectory() + path = os.path.join( + temp_dir.name, "test_jit_save_load_with_dict_input_no_prune/model" + ) + # prune inputs + paddle.jit.save( + layer=net, + path=path, + input_spec=[ + { + 'img': InputSpec( + shape=[None, 8], dtype='float32', name='img' + ), + 'img2': InputSpec( + shape=[None, 8], dtype='float32', name='img2' + ), + } + ], + ) + + img = paddle.randn(shape=[4, 8], dtype='float32') + img2 = paddle.randn(shape=[4, 8], dtype='float32') + loaded_net = paddle.jit.load(path) + loaded_out = loaded_net(img, img2) + + self.assertEqual(len(loaded_net._input_spec()), 2) + temp_dir.cleanup() + + +class TestSaveLoadWithInputSpec(unittest.TestCase): + def setUp(self): + # enable dygraph mode + base.enable_dygraph() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + @test_with_dygraph_pir + def test_with_input_spec(self): + net = LinearNetReturnLoss(8, 8) + # set x.shape = [None, 8] + net.forward = to_static( + net.forward, + input_spec=[InputSpec([None, 8], name='x')], + full_graph=True, + ) + + model_path = os.path.join( + self.temp_dir.name, "input_spec.output_spec/model" + ) + # check inputs and outputs + self.assertTrue(len(net.forward.inputs) == 1) + input_x = net.forward.inputs[0] + if paddle.framework.use_pir_api(): + self.assertTrue(input_x.shape == [-1, 8]) + else: + self.assertTrue(input_x.shape == (-1, 8)) + self.assertTrue(input_x.name == 'x') + + # 1. prune loss + output_spec = net.forward.outputs[:1] + paddle.jit.save(net, model_path, output_spec=output_spec) + + # 2. load to infer + infer_layer = paddle.jit.load(model_path) + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + pred = infer_layer(x) + + @test_with_dygraph_pir + def test_multi_in_out(self): + net = LinearNetMultiInput(8, 8) + net = paddle.jit.to_static( + net, + input_spec=[ + InputSpec([None, 8], dtype='float32'), + InputSpec([None, 8], dtype='float32'), + ], + full_graph=True, + ) + + model_path = os.path.join( + self.temp_dir.name, "multi_inout.output_spec1/model" + ) + # 1. check inputs and outputs + self.assertTrue(len(net.forward.inputs) == 2) + input_x = net.forward.inputs[0] + input_y = net.forward.inputs[1] + if paddle.framework.use_pir_api(): + self.assertTrue(input_x.shape == [-1, 8]) + self.assertTrue(input_y.shape == [-1, 8]) + else: + self.assertTrue(input_x.shape == (-1, 8)) + self.assertTrue(input_y.shape == (-1, 8)) + + # 2. prune loss + output_spec = net.forward.outputs[:2] + paddle.jit.save(net, model_path, output_spec=output_spec) + + # 3. load to infer + infer_layer = paddle.jit.load(model_path) + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + y = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + # 4. predict + pred_x, pred_y = infer_layer(x, y) + + # 1. prune y and loss + model_path = os.path.join( + self.temp_dir.name, "multi_inout.output_spec2/model" + ) + output_spec = net.forward.outputs[:1] + paddle.jit.save(net, model_path, [input_x], output_spec=output_spec) + # 2. load again + infer_layer2 = paddle.jit.load(model_path) + # 3. predict + pred_xx = infer_layer2(x) + + # 4. assert pred_x == pred_xx + np.testing.assert_allclose(pred_x.numpy(), pred_xx.numpy(), rtol=1e-05) + + @test_with_dygraph_pir + def test_multi_in_out1(self): + net = LinearNetMultiInput1(8, 8) + net = paddle.jit.to_static( + net, + input_spec=( + InputSpec([None, 8], dtype='float32'), + InputSpec([None, 8], dtype='float32'), + ), + full_graph=True, + ) + model_path = os.path.join( + self.temp_dir.name, "multi_inout1.output_spec1/model" + ) + # 1. check inputs and outputs + self.assertTrue(len(net.forward.inputs) == 2) + input_x = net.forward.inputs[0] + input_y = net.forward.inputs[1] + if paddle.framework.use_pir_api(): + self.assertTrue(input_x.shape == [-1, 8]) + self.assertTrue(input_y.shape == [-1, 8]) + else: + self.assertTrue(input_x.shape == (-1, 8)) + self.assertTrue(input_y.shape == (-1, 8)) + + # 2. prune loss + output_spec = net.forward.outputs[:2] + paddle.jit.save(net, model_path, output_spec=output_spec) + + # 3. load to infer + infer_layer = paddle.jit.load(model_path) + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + y = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + # 4. predict + pred_x, pred_y = infer_layer(x, y) + + # 1. prune y and loss + model_path = os.path.join( + self.temp_dir.name, "multi_inout1.output_spec2/model" + ) + output_spec = net.forward.outputs[:1] + paddle.jit.save( + net, + model_path, + net.forward.inputs, + output_spec=output_spec, + input_names_after_prune=[input_x.name], + ) + # 2. load again + infer_layer2 = paddle.jit.load(model_path) + # 3. predict + pred_xx = infer_layer2(x) + + # 4. assert pred_x == pred_xx + np.testing.assert_allclose(pred_x.numpy(), pred_xx.numpy(), rtol=1e-05) + + +class TestJitSaveLoadConfig(unittest.TestCase): + def setUp(self): + # enable dygraph mode + base.enable_dygraph() + # config seed + paddle.seed(SEED) + paddle.framework.random._manual_program_seed(SEED) + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + @test_with_dygraph_pir + def test_output_spec(self): + train_layer = LinearNetReturnLoss(8, 8) + train_layer.forward = to_static( + train_layer.forward, + input_spec=[InputSpec([None, 8], name='x')], + full_graph=True, + ) + adam = paddle.optimizer.Adam( + learning_rate=0.1, parameters=train_layer.parameters() + ) + x = paddle.to_tensor(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 = os.path.join( + self.temp_dir.name, "save_load_config.output_spec" + ) + output_spec = train_layer.forward.outputs[:1] + paddle.jit.save( + layer=train_layer, + path=model_path, + input_spec=[x], + output_spec=output_spec, + ) + + train_layer.eval() + infer_layer = paddle.jit.load(model_path) + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + np.testing.assert_array_equal( + train_layer(x)[0].numpy(), infer_layer(x).numpy() + ) + + @test_with_dygraph_pir + def test_save_no_support_config_error(self): + layer = LinearNet(784, 1) + path = os.path.join(self.temp_dir.name, "no_support_config_test") + with self.assertRaises(ValueError): + paddle.jit.save(layer=layer, path=path, model_filename="") + + @test_with_dygraph_pir + def test_load_empty_model_filename_error(self): + path = os.path.join(self.temp_dir.name, "error_model_filename_test") + + with self.assertRaises(ValueError): + paddle.jit.load(path, model_filename="") + + @test_with_dygraph_pir + def test_load_empty_params_filename_error(self): + path = os.path.join(self.temp_dir.name, "error_params_filename_test") + with self.assertRaises(ValueError): + paddle.jit.load(path, params_filename="") + + @test_with_dygraph_pir + def test_load_with_no_support_config(self): + path = os.path.join(self.temp_dir.name, "no_support_config_test") + with self.assertRaises(ValueError): + paddle.jit.load(path, separate_params=True) + + +class TestJitMultipleLoading(unittest.TestCase): + def setUp(self): + self.linear_size = 4 + self.temp_dir = tempfile.TemporaryDirectory() + self.model_path = os.path.join( + self.temp_dir.name, "jit_multi_load/model" + ) + # enable dygraph mode + base.enable_dygraph() + # config seed + paddle.seed(SEED) + paddle.framework.random._manual_program_seed(SEED) + # train and save base model + self.train_and_save_orig_model() + + def tearDown(self): + self.temp_dir.cleanup() + + @test_with_dygraph_pir + def train_and_save_orig_model(self): + layer = LinearNet(self.linear_size, self.linear_size) + example_inputs, layer, _ = train(layer, self.linear_size, 1) + paddle.jit.save( + layer=layer, path=self.model_path, input_spec=example_inputs + ) + + @test_with_dygraph_pir + def test_load_model_retransform_inference(self): + multi_loaded_layer = MultiLoadingLinearNet( + self.linear_size, self.model_path + ) + state_dict = multi_loaded_layer.state_dict() + name_set = set() + for name, var in state_dict.items(): + # print("name = ", name, " var.name = ", var.name) + self.assertTrue(var.name not in name_set) + name_set.add(var.name) + + +class TestJitPruneModelAndLoad(unittest.TestCase): + def setUp(self): + self.linear_size = 4 + self.temp_dir = tempfile.TemporaryDirectory() + self.model_path = os.path.join( + self.temp_dir.name, "jit_prune_model_and_load/model" + ) + # enable dygraph mode + base.enable_dygraph() + # config seed + paddle.seed(SEED) + paddle.framework.random._manual_program_seed(SEED) + + def tearDown(self): + self.temp_dir.cleanup() + + def train_and_save(self): + train_layer = LinearNetReturnHidden(8, 8) + train_layer = to_static( + train_layer, + input_spec=[InputSpec([None, 8], name='x')], + full_graph=True, + ) + adam = paddle.optimizer.Adam( + learning_rate=0.1, parameters=train_layer.parameters() + ) + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + for i in range(10): + hidden, loss = train_layer(x) + loss.backward() + adam.minimize(loss) + train_layer.clear_gradients() + + output_spec = train_layer.forward.outputs[:1] + paddle.jit.save( + layer=train_layer, + path=self.model_path, + input_spec=[x], + output_spec=output_spec, + ) + + return train_layer + + @test_with_dygraph_pir + def test_load_pruned_model(self): + train_layer = self.train_and_save() + train_layer.eval() + + infer_layer = paddle.jit.load(self.model_path) + + x = paddle.to_tensor(np.random.random((4, 8)).astype('float32')) + np.testing.assert_array_equal( + train_layer(x)[0].numpy(), infer_layer(x).numpy() + ) + + # pir has no need to save extra var info, param always saved with program, + # and trainable info saved in program's op attr + def test_load_var_not_in_extra_var_info(self): + self.train_and_save() + + # chage extra var info + var_info_path = self.model_path + INFER_PARAMS_INFO_SUFFIX + with open(var_info_path, 'rb') as f: + extra_var_info = pickle.load(f) + extra_var_info.clear() + with open(var_info_path, 'wb') as f: + pickle.dump(extra_var_info, f, protocol=2) + + with self.assertRaises(RuntimeError): + paddle.jit.load(self.model_path) + + +class TestJitSaveMultiCases(unittest.TestCase): + def setUp(self): + # enable dygraph mode + base.enable_dygraph() + # config seed + paddle.seed(SEED) + paddle.framework.random._manual_program_seed(SEED) + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def verify_inference_correctness( + self, layer, model_path, with_label_and_loss=False, with_label=False + ): + layer.eval() + loaded_layer = paddle.jit.load(model_path) + loaded_layer.eval() + # inference & compare + x = paddle.to_tensor(np.random.random((1, 784)).astype('float32')) + if with_label_and_loss: + y = paddle.to_tensor(np.random.random((1, 1)).astype('int64')) + pred, _ = layer(x, y) + pred = pred.numpy() + elif with_label: + y = paddle.to_tensor(np.random.random((1, 1)).astype('int64')) + pred = layer(x, y) + pred = pred.numpy() + else: + pred = layer(x).numpy() + loaded_pred = loaded_layer(x).numpy() + np.testing.assert_array_equal( + pred, + loaded_pred, + err_msg=f'Result diff when load and inference:\nlayer result:\n{pred}\nloaded layer result:\n{loaded_pred}', + ) + + @test_with_dygraph_pir + def test_no_prune_to_static_after_train(self): + layer = LinearNet(784, 1) + + train(layer) + + model_path = os.path.join( + self.temp_dir.name, "test_no_prune_to_static_after_train/model" + ) + paddle.jit.save(layer, model_path) + + self.verify_inference_correctness(layer, model_path) + + @test_with_dygraph_pir + def test_no_prune_to_static_no_train(self): + layer = LinearNetWithInputSpec(784, 1) + + model_path = os.path.join( + self.temp_dir.name, "test_no_prune_to_static_no_train/model" + ) + paddle.jit.save(layer, model_path) + + self.verify_inference_correctness(layer, model_path) + + @test_with_dygraph_pir + def test_no_prune_no_to_static_after_train(self): + layer = LinearNetNotDeclarative(784, 1) + + train(layer) + + model_path = os.path.join( + self.temp_dir.name, "test_no_prune_no_to_static_after_train/model" + ) + paddle.jit.save( + layer, + model_path, + input_spec=[InputSpec(shape=[None, 784], dtype='float32')], + ) + + self.verify_inference_correctness(layer, model_path) + + @test_with_dygraph_pir + def test_no_prune_no_to_static_after_train_with_examples(self): + layer = LinearNetNotDeclarative(784, 1) + + example_inputs, _, _ = train(layer) + + model_path = os.path.join( + self.temp_dir.name, + "test_no_prune_no_to_static_after_train_with_examples/model", + ) + paddle.jit.save(layer=layer, path=model_path, input_spec=example_inputs) + + self.verify_inference_correctness(layer, model_path) + + @test_with_dygraph_pir + def test_no_prune_no_to_static_no_train(self): + layer = LinearNetNotDeclarative(784, 1) + + model_path = os.path.join( + self.temp_dir.name, "test_no_prune_no_to_static_no_train/model" + ) + paddle.jit.save( + layer, + model_path, + input_spec=[InputSpec(shape=[None, 784], dtype='float32')], + ) + + self.verify_inference_correctness(layer, model_path) + + @test_with_dygraph_pir + def test_prune_to_static_after_train(self): + layer = LinerNetWithLabel(784, 1) + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image"), + InputSpec(shape=[None, 1], dtype='int64', name="label"), + ], + full_graph=True, + ) + out = train_with_label(layer) + + model_path = os.path.join( + self.temp_dir.name, "test_prune_to_static_after_train/model" + ) + paddle.jit.save( + layer, + model_path, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image"), + ], + output_spec=layer.forward.outputs[:1], + input_names_after_prune=["image"], + ) + + self.verify_inference_correctness( + layer, model_path, with_label_and_loss=True + ) + + @test_with_dygraph_pir + def test_prune_to_static_no_train(self): + layer = LinerNetWithLabel(784, 1) + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image"), + InputSpec(shape=[None, 1], dtype='int64', name="label"), + ], + full_graph=True, + ) + model_path = os.path.join( + self.temp_dir.name, "test_prune_to_static_no_train/model" + ) + # TODO: no train, cannot get output_spec var here + # now only can use index + output_spec = layer.forward.outputs[:1] + paddle.jit.save( + layer, + model_path, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image"), + ], + output_spec=output_spec, + input_names_after_prune=["image"], + ) + + self.verify_inference_correctness( + layer, model_path, with_label_and_loss=True + ) + + @test_with_dygraph_pir + def test_prune_input_to_static_no_train(self): + layer = LinerNetWithPruneInput(784, 1) + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image"), + InputSpec(shape=[None, 1], dtype='int64', name="label"), + ], + full_graph=True, + ) + model_path = os.path.join( + self.temp_dir.name, "test_prune_input_to_static_no_train/model" + ) + paddle.jit.save( + layer, + model_path, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image") + ], + ) + + self.verify_inference_correctness(layer, model_path, with_label=True) + + @test_with_dygraph_pir + def test_prune_useless_input_to_static_no_train(self): + layer = LinerNetWithUselessInput(784, 1) + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image"), + InputSpec(shape=[None, 1], dtype='int64', name="label"), + ], + full_graph=True, + ) + model_path = os.path.join( + self.temp_dir.name, + "test_prune_useless_input_to_static_no_train/model", + ) + paddle.jit.save( + layer, + model_path, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image") + ], + ) + + self.verify_inference_correctness(layer, model_path, with_label=True) + + @test_with_dygraph_pir + def test_no_prune_input_spec_name_warning(self): + layer = LinearNetWithInputSpec(784, 1) + + train(layer) + + model_path = os.path.join( + self.temp_dir.name, "test_no_prune_input_spec_name_warning/model" + ) + paddle.jit.save( + layer, + model_path, + input_spec=[InputSpec(shape=[None, 784], dtype='float32')], + ) + paddle.jit.save( + layer, + model_path, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name='feed_input') + ], + ) + + self.verify_inference_correctness(layer, model_path) + + @test_with_dygraph_pir + def test_not_prune_output_spec_name_warning(self): + layer = LinearNet(784, 1) + + train(layer) + + model_path = os.path.join( + self.temp_dir.name, "test_not_prune_output_spec_name_warning/model" + ) + out = paddle.to_tensor(np.random.random((1, 1)).astype('float')) + paddle.jit.save(layer, model_path, output_spec=[out]) + + self.verify_inference_correctness(layer, model_path) + + @test_with_dygraph_pir + def test_prune_input_spec_name_error(self): + layer = LinerNetWithLabel(784, 1) + + model_path = os.path.join( + self.temp_dir.name, "test_prune_input_spec_name_error/model" + ) + with self.assertRaises(ValueError): + paddle.jit.save( + layer, + model_path, + input_spec=[InputSpec(shape=[None, 784], dtype='float32')], + ) + with self.assertRaises(ValueError): + paddle.jit.save( + layer, + model_path, + input_spec=[ + InputSpec( + shape=[None, 784], dtype='float32', name='feed_input' + ) + ], + ) + + @test_with_dygraph_pir + def test_prune_output_spec_name_error(self): + layer = LinerNetWithLabel(784, 1) + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image"), + InputSpec(shape=[None, 1], dtype='int64', name="label"), + ], + full_graph=True, + ) + train_with_label(layer) + + model_path = os.path.join( + self.temp_dir.name, "test_prune_to_static_after_train/model" + ) + out = paddle.to_tensor(np.random.random((1, 1)).astype('float')) + with self.assertRaises(ValueError): + paddle.jit.save( + layer, + model_path, + input_spec=[ + InputSpec(shape=[None, 784], dtype='float32', name="image"), + True, + ], + output_spec=[out], + input_names_after_prune=["image"], + ) + + +class TestJitSaveLoadEmptyLayer(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.model_path = os.path.join( + self.temp_dir.name, "jit_save_load_empty_layer/model" + ) + # enable dygraph mode + paddle.disable_static() + + def tearDown(self): + self.temp_dir.cleanup() + + @test_with_dygraph_pir + def test_save_load_empty_layer(self): + layer = EmptyLayer() + x = paddle.to_tensor(np.random.random(10).astype('float32')) + out = layer(x) + paddle.jit.save(layer, self.model_path) + load_layer = paddle.jit.load(self.model_path) + load_out = load_layer(x) + np.testing.assert_array_equal(out, load_out) + + +class TestJitSaveLoadNoParamLayer(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.model_path = os.path.join( + self.temp_dir.name, "jit_save_load_no_param_layer/model" + ) + # enable dygraph mode + paddle.disable_static() + + def tearDown(self): + self.temp_dir.cleanup() + + @test_with_dygraph_pir + def test_save_load_no_param_layer(self): + layer = NoParamLayer() + x = paddle.to_tensor(np.random.random(5).astype('float32')) + y = paddle.to_tensor(np.random.random(5).astype('float32')) + out = layer(x, y) + paddle.jit.save(layer, self.model_path) + load_layer = paddle.jit.load(self.model_path) + load_out = load_layer(x, y) + np.testing.assert_array_equal(out, load_out) + + +''' +class TestJitSaveLoadMultiMethods(unittest.TestCase): + def setUp(self): + # enable dygraph mode + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + #@test_with_dygraph_pir + def test_jit_save_load_inference(self): + model_path_inference = os.path.join( + self.temp_dir.name, "jit_save_load_multi_methods/model" + ) + IMAGE_SIZE = 224 + layer = LinearNetWithMultiStaticFunc(IMAGE_SIZE, 10) + layer = paddle.jit.to_static( + layer, + full_graph=True, + ) + inps = paddle.randn([1, IMAGE_SIZE]) + result_origin = {} + for func in dir(layer): + if func.startswith('forward'): + result_origin[func] = getattr(layer, func, None)(inps) + paddle.jit.save(layer, model_path_inference) + load_net = paddle.jit.load(model_path_inference) + for func, result in result_origin.items(): + self.assertTrue( + float( + (result - getattr(load_net, func, None)(inps)).abs().max() + ) + < 1e-5 + ) + + def test_jit_save_load_multi_methods_inputspec(self): + model_path = os.path.join( + self.temp_dir.name, 'jit_save_load_multi_methods/model' + ) + layer = LinearNetWithMultiStaticFunc(784, 1) + layer = paddle.jit.to_static( + layer, + full_graph=True, + ) + with self.assertRaises(ValueError): + paddle.jit.save( + layer, model_path, input_spec=[InputSpec(shape=[None, 784])] + ) + + def test_parse_name(self): + model_path_inference = os.path.join( + self.temp_dir.name, "jit_save_load_parse_name/model" + ) + IMAGE_SIZE = 224 + layer = LinearNet(IMAGE_SIZE, 1) + inps = paddle.randn([1, IMAGE_SIZE]) + layer(inps) + paddle.jit.save(layer, model_path_inference) + paddle.jit.save(layer, model_path_inference + '_v2') + load_net = paddle.jit.load(model_path_inference) + + self.assertFalse(hasattr(load_net, 'v2')) + + +class LayerSaved(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super().__init__() + self.hidden = 100 + self._linear_0 = Linear(in_size, self.hidden) + self._linear_1_0 = Linear(self.hidden, self.hidden) + self._linear_1_1 = Linear(self.hidden, self.hidden) + self._linear_2 = Linear(self.hidden, out_size) + self._scale = paddle.to_tensor([9.9]) + + def forward(self, x): + y = self._linear_0(x) + # Multiple blocks + if paddle.shape(x)[0] == 1: + y = self._linear_1_0(y) + else: + y += self._linear_1_1(y + self._scale) + return self._linear_2(y) + + +class Net(paddle.nn.Layer): + def __init__(self): + super().__init__() + self.fc1 = paddle.nn.Linear(4, 4) + self.fc2 = paddle.nn.Linear(4, 4) + self.bias = 0.4 + self.flag = paddle.ones([2], dtype="int32") + + @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) + def log_softmax(self, input): + return paddle.nn.functional.log_softmax(input, axis=-1) + + @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) + def forward(self, x): + out = self.fc1(x) + out = paddle.nn.functional.relu(out) + out = paddle.mean(out) + return out + + @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) + def infer(self, input): + out = self.fc2(input) + out = out + self.bias + out = paddle.mean(out) + return out + + # For extra Python float + @paddle.jit.to_static(property=True) + def fbias(self): + return self.bias + 1 + + @paddle.jit.to_static(property=True) + def down_sampling(self): + return 4 + + @paddle.jit.to_static(property=True) + def fstr(self): + return "save str property" + + @paddle.jit.to_static(property=True) + def ints(self): + return [10, 20] + + @paddle.jit.to_static(property=True) + def floats(self): + return [1.1, 2.2] + + @paddle.jit.to_static(property=True) + def strs(self): + return ["hello", "world"] + + +class NetTensor(paddle.nn.Layer): + def __init__(self): + super().__init__() + self.fc1 = paddle.nn.Linear(4, 4) + self.fc2 = paddle.nn.Linear(4, 4) + self.bias = 0.4 + self.flag = paddle.ones([2], dtype="int32") + + @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) + def forward(self, x): + out = self.fc1(x) + out = paddle.nn.functional.relu(out) + out = paddle.mean(out) + return out + + @paddle.jit.to_static(property=True) + def fflag(self): + return True + + +class TestJitSaveCombineProperty(unittest.TestCase): + def setUp(self): + # enable dygraph mode + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_jit_save_combine_property(self): + model_path = os.path.join( + self.temp_dir.name, "test_jit_save_combine/model" + ) + # Use new namespace + with unique_name.guard(): + net = Net() + # save + paddle.jit.save(net, model_path, combine_params=True) + + def test_jit_save_tensor_property(self): + model_path = os.path.join( + self.temp_dir.name, "test_jit_save_combine/model" + ) + # Use new namespace + with unique_name.guard(): + net = NetTensor() + + paddle.jit.save(net, model_path, combine_params=True) + + +class LayerLoadFinetune(paddle.nn.Layer): + def __init__(self, in_size, out_size, load_path): + super().__init__() + # Test duplicate name + self._linear_0 = Linear(in_size, in_size) + self._linear_1_0 = Linear(out_size, in_size) + self._linear_1_1 = Linear(out_size, in_size) + self._linear_2 = Linear(out_size, out_size) + self._scale = paddle.to_tensor([9.9]) + + # Load multiple times + self._load_l1 = paddle.jit.load(load_path) + self._load_l2 = paddle.jit.load(load_path) + + @paddle.jit.to_static + def forward(self, x): + y = self._linear_0(x) + y = self._load_l1(y) + # Multiple blocks + if paddle.shape(x)[0] == 1: + y = self._linear_1_0(y) + y = self._load_l1(y) + else: + y += self._linear_1_1(x + self._scale) + y = self._load_l2(y) + y = self._linear_1_0(y) + y = self._load_l1(y) + y = self._linear_1_0(y) + # Use the same layer multiple times. + y = self._load_l1(y) + return y + + +class TestJitSaveLoadSaveWithoutRunning(unittest.TestCase): + def setUp(self): + # enable dygraph mode + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_save_load_finetune_load(self): + model_path = os.path.join( + self.temp_dir.name, "test_jit_save_load_save_without_running/model" + ) + IMAGE_SIZE = 224 + inps0 = paddle.randn([1, IMAGE_SIZE]) + inps1 = paddle.randn([2, IMAGE_SIZE]) + # Use new namespace + with unique_name.guard(): + layer_save = LayerSaved(IMAGE_SIZE, IMAGE_SIZE) + layer_save = paddle.jit.to_static(layer_save, full_graph=True) + # save + paddle.jit.save( + layer_save, + model_path, + input_spec=[ + paddle.static.InputSpec( + shape=[None, IMAGE_SIZE], dtype='float32' + ) + ], + ) + result_00 = layer_save(inps0) + result_01 = layer_save(inps1) + # load and save without running + with unique_name.guard(): + layer_load = paddle.jit.load(model_path) + paddle.jit.save( + layer_load, + model_path, + input_spec=[ + paddle.static.InputSpec( + shape=[None, IMAGE_SIZE], dtype='float32' + ) + ], + ) + # reload + layer_reload = paddle.jit.load(model_path) + result_10 = layer_reload(inps0) + result_11 = layer_reload(inps1) + + self.assertTrue(float((result_00 - result_10).abs().max()) < 1e-5) + self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5) + + +class TestJitSaveLoadFinetuneLoad(unittest.TestCase): + def setUp(self): + # enable dygraph mode + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_save_load_finetune_load(self): + model_path = os.path.join( + self.temp_dir.name, "test_jit_save_load_finetune_load/model" + ) + IMAGE_SIZE = 224 + inps0 = paddle.randn([1, IMAGE_SIZE]) + inps1 = paddle.randn([2, IMAGE_SIZE]) + # Use new namespace + with unique_name.guard(): + layer_save = LayerSaved(IMAGE_SIZE, IMAGE_SIZE) + layer_save = paddle.jit.to_static(layer_save, full_graph=True) + layer_save(inps0) + # save + paddle.jit.save(layer_save, model_path) + # load + with unique_name.guard(): + layer_load = LayerLoadFinetune(IMAGE_SIZE, IMAGE_SIZE, model_path) + # train + train(layer_load, input_size=IMAGE_SIZE) + result_00 = layer_load(inps0) + result_01 = layer_load(inps1) + # save + paddle.jit.save(layer_load, model_path) + # load + layer_finetune = paddle.jit.load(model_path) + result_10 = layer_finetune(inps0) + result_11 = layer_finetune(inps1) + + self.assertTrue(float((result_00 - result_10).abs().max()) < 1e-5) + self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5) + + +# NOTE(weixin): When there are multiple test functions in an +# `unittest.TestCase`, functions will affect each other, +# and there is a risk of random failure. +# So divided into three TestCase: TestJitSaveLoadFunctionCase1, +# TestJitSaveLoadFunctionCase2, TestJitSaveLoadFunctionCase3. +class TestJitSaveLoadFunctionCase1(unittest.TestCase): + def setUp(self): + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_jit_save_load_static_function(self): + @paddle.jit.to_static + def fun(inputs): + return paddle.tanh(inputs) + + path = os.path.join( + self.temp_dir.name, 'test_jit_save_load_function_1/func' + ) + inps = paddle.rand([3, 6]) + origin = fun(inps) + + paddle.jit.save(fun, path) + load_func = paddle.jit.load(path) + + load_result = load_func(inps) + self.assertTrue((load_result - origin).abs().max() < 1e-10) + + +class TestJitSaveLoadFunctionCase2(unittest.TestCase): + def setUp(self): + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_jit_save_load_function_input_spec(self): + @paddle.jit.to_static( + input_spec=[ + InputSpec(shape=[None, 6], dtype='float32', name='x'), + ] + ) + def fun(inputs): + return paddle.nn.functional.relu(inputs) + + path = os.path.join( + self.temp_dir.name, 'test_jit_save_load_function_2/func' + ) + inps = paddle.rand([3, 6]) + origin = fun(inps) + + paddle.jit.save(fun, path) + load_func = paddle.jit.load(path) + load_result = load_func(inps) + self.assertTrue((load_result - origin).abs().max() < 1e-10) + + +class TestJitSaveLoadFunctionCase3(unittest.TestCase): + def setUp(self): + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_jit_save_load_function_function(self): + def fun(inputs): + return paddle.tanh(inputs) + + path = os.path.join( + self.temp_dir.name, 'test_jit_save_load_function_3/func' + ) + inps = paddle.rand([3, 6]) + origin = fun(inps) + + paddle.jit.save( + fun, + path, + input_spec=[ + InputSpec(shape=[None, 6], dtype='float32', name='x'), + ], + ) + load_func = paddle.jit.load(path) + + load_result = load_func(inps) + self.assertTrue((load_result - origin).abs().max() < 1e-10) + + +class TestJitSaveLoadFunctionWithParamCase1(unittest.TestCase): + def setUp(self): + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_jit_save_load_function(self): + class LinearNet(paddle.nn.Layer): + def __init__(self): + super().__init__() + self._linear = paddle.nn.Linear(5, 6) + + def forward(self, x): + return paddle.tanh(x) + + def anothor_forward(self, x): + return self._linear(x) + + layer = LinearNet() + + inps = paddle.rand([3, 5]) + origin = layer.anothor_forward(inps) + + func = paddle.jit.to_static( + layer.anothor_forward, [paddle.static.InputSpec(shape=[-1, 5])] + ) + path = os.path.join( + self.temp_dir.name, + 'test_jit_save_load_function_with_params_case1/func', + ) + paddle.jit.save(func, path) + load_func = paddle.jit.load(path) + + load_result = load_func(inps) + np.testing.assert_array_equal(load_result.numpy(), origin.numpy()) + + +class TestJitSaveLoadFunctionWithParamCase2(unittest.TestCase): + def setUp(self): + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_jit_save_load_function(self): + class LinearNet(paddle.nn.Layer): + def __init__(self): + super().__init__() + self._linear = paddle.nn.Linear(5, 6) + + def forward(self, x): + return paddle.tanh(x) + + @paddle.jit.to_static(input_spec=[InputSpec(shape=[-1, 5])]) + def anothor_forward(self, x): + return self._linear(x) + + layer = LinearNet() + + inps = paddle.rand([3, 5]) + + path = os.path.join( + self.temp_dir.name, + 'test_jit_save_load_function_with_params_case2/func', + ) + paddle.jit.save(layer.anothor_forward, path) + origin_result = layer.anothor_forward(inps) + load_func = paddle.jit.load(path) + + load_result = load_func(inps) + + np.testing.assert_array_equal( + origin_result.numpy(), load_result.numpy() + ) + + +class TestJitSaveLoadFunctionWithParamCase3(unittest.TestCase): + def setUp(self): + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_jit_save_load_function(self): + class LinearNet(paddle.nn.Layer): + def __init__(self): + super().__init__() + self._linear = paddle.nn.Linear(5, 6) + + def forward(self, x): + return paddle.tanh(x) + + @paddle.jit.to_static + def anothor_forward(self, x): + return self._linear(x) + + layer = LinearNet() + + inps = paddle.rand([3, 5]) + origin = layer.anothor_forward(inps) + + path = os.path.join( + self.temp_dir.name, + 'test_jit_save_load_function_with_params_case3/func', + ) + paddle.jit.save(layer.anothor_forward, path) + load_func = paddle.jit.load(path) + + load_result = load_func(inps) + np.testing.assert_array_equal(load_result.numpy(), origin.numpy()) + + +class TestJitSaveLoadDataParallel(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def verify_inference_correctness(self, layer, path): + layer.eval() + loaded_layer = paddle.jit.load(path) + loaded_layer.eval() + # inference & compare + x = paddle.to_tensor(np.random.random((1, 784)).astype('float32')) + pred = layer(x).numpy() + loaded_pred = loaded_layer(x).numpy() + np.testing.assert_array_equal( + pred, + loaded_pred, + err_msg=f'Result diff when load and inference:\nlayer result:\n{pred}\nloaded layer result:\n{loaded_pred}', + ) + + def test_jit_save_data_parallel_with_inputspec(self): + layer = LinearNetNotDeclarative(784, 1) + layer = paddle.DataParallel(layer) + path = os.path.join( + self.temp_dir.name, "jit_save_data_parallel_with_inputspec/model" + ) + paddle.jit.save( + layer=layer, path=path, input_spec=[InputSpec(shape=[None, 784])] + ) + + self.verify_inference_correctness(layer, path) + + def test_jit_save_data_parallel_with_to_static(self): + layer = LinearNetWithInputSpec(784, 1) + layer = paddle.DataParallel(layer) + + path = os.path.join( + self.temp_dir.name, "jit_save_data_parallel_with_to_static/model" + ) + paddle.jit.save(layer, path) + + self.verify_inference_correctness(layer, path) + + +class InputSepcLayer(paddle.nn.Layer): + #A layer with InputSpec to test InputSpec compatibility + + def forward(self, x, y): + return x, y + + +class TestInputSpecCompatibility(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def _assert_input_spec_layer_return(self, expect_layer, test_layer): + input_x = paddle.uniform([8, 8], dtype='float32') + input_y = paddle.uniform([8, 1], dtype='float64') + expected_result = expect_layer(input_x, input_y) + test_result = test_layer(input_x, input_y) + np.testing.assert_allclose( + expected_result[0].numpy(), test_result[0].numpy() + ) + np.testing.assert_allclose( + expected_result[1].numpy(), test_result[1].numpy() + ) + + def test_jit_save_compatible_input_sepc(self): + layer = InputSepcLayer() + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 8], dtype='float32', name='x'), + InputSpec(shape=[None, 1], dtype='float64', name='y'), + ], + full_graph=True, + ) + save_dir = os.path.join( + self.temp_dir.name, "jit_save_compatible_input_spec" + ) + path = save_dir + "/model" + + paddle.jit.save(layer=layer, path=path) + no_input_spec_layer = paddle.jit.load(path) + self._assert_input_spec_layer_return(layer, no_input_spec_layer) + shutil.rmtree(save_dir) + + paddle.jit.save( + layer=layer, + path=path, + input_spec=[ + InputSpec(shape=[None, 8], dtype='float32', name='x'), + InputSpec(shape=[None, 1], dtype='float64', name='y'), + ], + ) + same_input_spec_layer = paddle.jit.load(path) + self._assert_input_spec_layer_return(layer, same_input_spec_layer) + shutil.rmtree(save_dir) + + paddle.jit.save( + layer=layer, + path=path, + input_spec=[ + InputSpec(shape=[8, 8], dtype='float32'), + InputSpec(shape=[8, -1], dtype='float64'), + ], + ) + compatible_input_spec_layer = paddle.jit.load(path) + self._assert_input_spec_layer_return(layer, compatible_input_spec_layer) + shutil.rmtree(save_dir) + + def test_jit_save_incompatible_input_sepc(self): + layer = InputSepcLayer() + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 8], dtype='float32', name='x'), + InputSpec(shape=[None, 1], dtype='float64', name='y'), + ], + full_graph=True, + ) + save_dir = os.path.join( + self.temp_dir.name, "jit_save_compatible_input_spec" + ) + path = save_dir + "/model" + + with self.assertRaises(ValueError): + # type mismatch + paddle.jit.save( + layer=layer, + path=path, + input_spec=[ + InputSpec(shape=[None, 8], dtype='float64'), + InputSpec(shape=[None, 1], dtype='float64'), + ], + ) + + with self.assertRaises(ValueError): + # shape len mismatch + paddle.jit.save( + layer=layer, + path=path, + input_spec=[ + InputSpec(shape=[None, 8, 1], dtype='float32'), + InputSpec(shape=[None, 1], dtype='float64'), + ], + ) + + with self.assertRaises(ValueError): + # shape mismatch + paddle.jit.save( + layer=layer, + path=path, + input_spec=[ + InputSpec(shape=[None, 8], dtype='float32'), + InputSpec(shape=[None, 2], dtype='float64'), + ], + ) + if os.path.exists(save_dir): + shutil.rmtree(save_dir) + + +class NotJitForward(paddle.nn.Layer): + def __init__(self): + super().__init__() + + def forward(self, x, y): + return x + y + + +class TestNotJitForward(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_jit_not_save_forward(self): + layer = NotJitForward() + + save_dir = os.path.join(self.temp_dir.name, "jit_not_save_forward") + path = save_dir + "/model" + + paddle.jit.save(layer=layer, path=path, skip_forward=True) + + self.assertTrue(not os.path.exists(path + ".pdmodel")) + self.assertTrue(not os.path.exists(path + ".pdparam")) + + with self.assertRaises(ValueError): + paddle.jit.load(path=path) + + shutil.rmtree(save_dir) + +''' +if __name__ == '__main__': + unittest.main() From 1146d09e31177e0dc8647316b52a1f724ec6f3ae Mon Sep 17 00:00:00 2001 From: wangruting Date: Tue, 30 Apr 2024 01:38:02 +0000 Subject: [PATCH 08/15] change name all --- test/dygraph_to_static/test_pylayer.py | 2 +- test/legacy_test/CMakeLists.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/dygraph_to_static/test_pylayer.py b/test/dygraph_to_static/test_pylayer.py index bf09ba3db8a8cd..9724eb2749e40e 100644 --- a/test/dygraph_to_static/test_pylayer.py +++ b/test/dygraph_to_static/test_pylayer.py @@ -28,7 +28,7 @@ import unittest import numpy as np -from test_jit_save_load import train +from test_jit_save_load_rename import train import paddle from paddle.autograd.py_layer import PyLayer diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt index 254bb35b5f98fa..b1f637f61d2b68 100644 --- a/test/legacy_test/CMakeLists.txt +++ b/test/legacy_test/CMakeLists.txt @@ -492,7 +492,7 @@ foreach(TEST_OP ${TEST_OPS_WITH_GC}) endforeach() # Switch some dy2st UT to eager mode -set(TEST_EAGER_OPS test_jit_save_load test_translated_layer) +set(TEST_EAGER_OPS test_jit_save_load_rename test_translated_layer) foreach(TEST_OP ${TEST_EAGER_OPS}) list(REMOVE_ITEM TEST_OPS ${TEST_OP}) py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS FLAGS_enable_eager_mode=1) From 4e5d6e232da4997e2de528704fae4f7a6a480ff1 Mon Sep 17 00:00:00 2001 From: wangruting Date: Mon, 6 May 2024 01:23:32 +0000 Subject: [PATCH 09/15] modify timeout --- test/legacy_test/CMakeLists.txt | 1 + test/legacy_test/test_jit_save_load_rename.py | 95 ++++++++++++------- 2 files changed, 60 insertions(+), 36 deletions(-) diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt index b1f637f61d2b68..5549dd9426ee85 100644 --- a/test/legacy_test/CMakeLists.txt +++ b/test/legacy_test/CMakeLists.txt @@ -855,6 +855,7 @@ set_tests_properties(test_vision_models PROPERTIES TIMEOUT 120) set_tests_properties(test_dataset_uci_housing PROPERTIES TIMEOUT 120) set_tests_properties(test_dataset_imdb PROPERTIES TIMEOUT 300) set_tests_properties(test_callback_wandb PROPERTIES TIMEOUT 60) +set_tests_properties(test_jit_save_load_rename PROPERTIES TIMEOUT 50) if(WITH_COVERAGE) set_tests_properties(test_hapi_hub PROPERTIES TIMEOUT 300) endif() diff --git a/test/legacy_test/test_jit_save_load_rename.py b/test/legacy_test/test_jit_save_load_rename.py index 789d4277a7b318..9b4d4fb592fd71 100644 --- a/test/legacy_test/test_jit_save_load_rename.py +++ b/test/legacy_test/test_jit_save_load_rename.py @@ -16,6 +16,7 @@ import copy import os import pickle +import shutil import tempfile import unittest @@ -23,6 +24,7 @@ import paddle from paddle import base +from paddle.base import unique_name from paddle.jit.api import to_static from paddle.jit.translated_layer import INFER_PARAMS_INFO_SUFFIX from paddle.nn import Linear @@ -1532,6 +1534,7 @@ def test_jit_save_load_multi_methods_inputspec(self): layer, model_path, input_spec=[InputSpec(shape=[None, 784])] ) + @test_with_dygraph_pir def test_parse_name(self): model_path_inference = os.path.join( self.temp_dir.name, "jit_save_load_parse_name/model" @@ -1546,6 +1549,8 @@ def test_parse_name(self): self.assertFalse(hasattr(load_net, 'v2')) +''' + class LayerSaved(paddle.nn.Layer): def __init__(self, in_size, out_size): @@ -1627,7 +1632,6 @@ def __init__(self): self.bias = 0.4 self.flag = paddle.ones([2], dtype="int32") - @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) def forward(self, x): out = self.fc1(x) out = paddle.nn.functional.relu(out) @@ -1665,43 +1669,15 @@ def test_jit_save_tensor_property(self): # Use new namespace with unique_name.guard(): net = NetTensor() + net = paddle.jit.to_static( + net, + input_spec=[InputSpec([None, 4], dtype='float32')], + full_graph=True, + ) paddle.jit.save(net, model_path, combine_params=True) -class LayerLoadFinetune(paddle.nn.Layer): - def __init__(self, in_size, out_size, load_path): - super().__init__() - # Test duplicate name - self._linear_0 = Linear(in_size, in_size) - self._linear_1_0 = Linear(out_size, in_size) - self._linear_1_1 = Linear(out_size, in_size) - self._linear_2 = Linear(out_size, out_size) - self._scale = paddle.to_tensor([9.9]) - - # Load multiple times - self._load_l1 = paddle.jit.load(load_path) - self._load_l2 = paddle.jit.load(load_path) - - @paddle.jit.to_static - def forward(self, x): - y = self._linear_0(x) - y = self._load_l1(y) - # Multiple blocks - if paddle.shape(x)[0] == 1: - y = self._linear_1_0(y) - y = self._load_l1(y) - else: - y += self._linear_1_1(x + self._scale) - y = self._load_l2(y) - y = self._linear_1_0(y) - y = self._load_l1(y) - y = self._linear_1_0(y) - # Use the same layer multiple times. - y = self._load_l1(y) - return y - - class TestJitSaveLoadSaveWithoutRunning(unittest.TestCase): def setUp(self): # enable dygraph mode @@ -1755,6 +1731,38 @@ def test_save_load_finetune_load(self): self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5) +''' +class LayerLoadFinetune(paddle.nn.Layer): + def __init__(self, in_size, out_size, load_path): + super().__init__() + # Test duplicate name + self._linear_0 = Linear(in_size, in_size) + self._linear_1_0 = Linear(out_size, in_size) + self._linear_1_1 = Linear(out_size, in_size) + self._linear_2 = Linear(out_size, out_size) + self._scale = paddle.to_tensor([9.9]) + + # Load multiple times + self._load_l1 = paddle.jit.load(load_path) + self._load_l2 = paddle.jit.load(load_path) + + def forward(self, x): + y = self._linear_0(x) + y = self._load_l1(y) + # Multiple blocks + if paddle.shape(x)[0] == 1: + y = self._linear_1_0(y) + y = self._load_l1(y) + else: + y += self._linear_1_1(x + self._scale) + y = self._load_l2(y) + y = self._linear_1_0(y) + y = self._load_l1(y) + y = self._linear_1_0(y) + # Use the same layer multiple times. + y = self._load_l1(y) + return y + class TestJitSaveLoadFinetuneLoad(unittest.TestCase): def setUp(self): # enable dygraph mode @@ -1764,6 +1772,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + #@test_with_dygraph_pir def test_save_load_finetune_load(self): model_path = os.path.join( self.temp_dir.name, "test_jit_save_load_finetune_load/model" @@ -1781,6 +1790,7 @@ def test_save_load_finetune_load(self): # load with unique_name.guard(): layer_load = LayerLoadFinetune(IMAGE_SIZE, IMAGE_SIZE, model_path) + layer_load = paddle.jit.to_static(layer_load, full_graph=True) # train train(layer_load, input_size=IMAGE_SIZE) result_00 = layer_load(inps0) @@ -1795,6 +1805,8 @@ def test_save_load_finetune_load(self): self.assertTrue(float((result_00 - result_10).abs().max()) < 1e-5) self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5) +''' + # NOTE(weixin): When there are multiple test functions in an # `unittest.TestCase`, functions will affect each other, @@ -1809,6 +1821,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_jit_save_load_static_function(self): @paddle.jit.to_static def fun(inputs): @@ -1835,6 +1848,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_jit_save_load_function_input_spec(self): @paddle.jit.to_static( input_spec=[ @@ -1864,6 +1878,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_jit_save_load_function_function(self): def fun(inputs): return paddle.tanh(inputs) @@ -1895,6 +1910,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_jit_save_load_function(self): class LinearNet(paddle.nn.Layer): def __init__(self): @@ -1934,6 +1950,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_jit_save_load_function(self): class LinearNet(paddle.nn.Layer): def __init__(self): @@ -1974,6 +1991,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_jit_save_load_function(self): class LinearNet(paddle.nn.Layer): def __init__(self): @@ -2024,6 +2042,7 @@ def verify_inference_correctness(self, layer, path): err_msg=f'Result diff when load and inference:\nlayer result:\n{pred}\nloaded layer result:\n{loaded_pred}', ) + @test_with_dygraph_pir def test_jit_save_data_parallel_with_inputspec(self): layer = LinearNetNotDeclarative(784, 1) layer = paddle.DataParallel(layer) @@ -2036,6 +2055,7 @@ def test_jit_save_data_parallel_with_inputspec(self): self.verify_inference_correctness(layer, path) + @test_with_dygraph_pir def test_jit_save_data_parallel_with_to_static(self): layer = LinearNetWithInputSpec(784, 1) layer = paddle.DataParallel(layer) @@ -2049,7 +2069,7 @@ def test_jit_save_data_parallel_with_to_static(self): class InputSepcLayer(paddle.nn.Layer): - #A layer with InputSpec to test InputSpec compatibility + # A layer with InputSpec to test InputSpec compatibility def forward(self, x, y): return x, y @@ -2074,6 +2094,7 @@ def _assert_input_spec_layer_return(self, expect_layer, test_layer): expected_result[1].numpy(), test_result[1].numpy() ) + @test_with_dygraph_pir def test_jit_save_compatible_input_sepc(self): layer = InputSepcLayer() layer = paddle.jit.to_static( @@ -2118,6 +2139,7 @@ def test_jit_save_compatible_input_sepc(self): self._assert_input_spec_layer_return(layer, compatible_input_spec_layer) shutil.rmtree(save_dir) + @test_with_dygraph_pir def test_jit_save_incompatible_input_sepc(self): layer = InputSepcLayer() layer = paddle.jit.to_static( @@ -2184,6 +2206,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + @test_with_dygraph_pir def test_jit_not_save_forward(self): layer = NotJitForward() @@ -2200,6 +2223,6 @@ def test_jit_not_save_forward(self): shutil.rmtree(save_dir) -''' + if __name__ == '__main__': unittest.main() From aa8d7ae790588fff44a73d734ca816ed9f865b10 Mon Sep 17 00:00:00 2001 From: wangruting Date: Mon, 6 May 2024 09:36:27 +0000 Subject: [PATCH 10/15] modify new case --- python/paddle/jit/api.py | 77 +++++--- .../jit/dy2static/pir_partial_program.py | 3 + python/paddle/static/io_utils.py | 3 +- test/legacy_test/CMakeLists.txt | 2 +- test/legacy_test/test_jit_save_load_rename.py | 180 ++++++++++-------- 5 files changed, 156 insertions(+), 109 deletions(-) diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py index d7402c3ea5f3c7..a17eeb402febbe 100644 --- a/python/paddle/jit/api.py +++ b/python/paddle/jit/api.py @@ -1024,6 +1024,7 @@ def save(layer, path, input_spec=None, **configs): ] combine_vars = {} + combine_program = [] property_vals = [] # (value, key) concrete_program = None for attr_func in functions: @@ -1050,6 +1051,7 @@ def save(layer, path, input_spec=None, **configs): is_prim_infer=is_prim_infer, ) ) + elif 'forward' == attr_func: if configs.skip_forward: # do not jit.save forward function @@ -1067,6 +1069,7 @@ def save(layer, path, input_spec=None, **configs): input_spec=inner_input_spec, full_graph=True, ) + concrete_program = ( static_forward.concrete_program_specify_input_spec( with_hook=with_hook, is_prim_infer=is_prim_infer @@ -1132,7 +1135,6 @@ def save(layer, path, input_spec=None, **configs): for structured_name, var in dygraph_state_dict.items(): state_names_dict[var.name] = structured_name state_var_dict[var.name] = var - # 3. share parameters from Layer to scope & record var info with dygraph.guard(): if use_pir_api(): @@ -1216,18 +1218,8 @@ def save(layer, path, input_spec=None, **configs): ) file_prefix = file_prefix + '.' + attr_func file_prefix = os.path.join(model_path, file_prefix) - with scope_guard(scope): - if not use_pir_api(): - input_vars = [ - concrete_program.main_program.global_block().var(name) - for name in input_var_names - ] - clone_program = concrete_program.main_program.clone() - clone_input_vars = input_vars - clone_output_vars = output_vars - - else: + if use_pir_api(): value_map = paddle.pir.IrMapping() clone_program = concrete_program.main_program.clone(value_map) clone_input_vars = [] @@ -1245,6 +1237,15 @@ def save(layer, path, input_spec=None, **configs): clone_output_vars = [value_map.look_up(v) for v in output_vars] + else: + input_vars = [ + concrete_program.main_program.global_block().var(name) + for name in input_var_names + ] + clone_program = concrete_program.main_program.clone() + clone_input_vars = input_vars + clone_output_vars = output_vars + save_inference_model( path_prefix=file_prefix, feed_vars=clone_input_vars, @@ -1256,12 +1257,23 @@ def save(layer, path, input_spec=None, **configs): ) if combine_params: - clone_main_program = concrete_program.main_program.clone() - clone_main_program = clone_main_program._prune_with_input( - input_var_names, output_vars - ) - for block in clone_main_program.blocks: - combine_vars.update(block.vars) + if use_pir_api(): + # NOTE(Ruting): concrete_program has been pruned when init partialProgramLayer, + # so we do not neet to prune again. + + for var in concrete_program.main_program.list_vars(): + if var.persistable: + combine_vars[var.name] = var + # NOTE(Ruting): concrete_program will delete after this loop item, + # value delete at the same time, so we use list to Extend its lifecycle + combine_program.append(concrete_program.main_program) + else: + clone_main_program = concrete_program.main_program.clone() + clone_main_program = clone_main_program._prune_with_input( + input_var_names, output_vars + ) + for block in clone_main_program.blocks: + combine_vars.update(block.vars) # save shared params if combine_params: @@ -1273,16 +1285,25 @@ def save(layer, path, input_spec=None, **configs): params_filename = file_prefix + INFER_PARAMS_SUFFIX with scope_guard(scope): - paddle.static.save_vars( - Executor(_current_expected_place()), - dirname=model_path, - vars=list( - filter( - paddle.framework.io_utils.is_persistable, ordered_vars - ) - ), - filename=params_filename, - ) + if use_pir_api(): + paddle.static.save_vars( + Executor(_current_expected_place()), + dirname=model_path, + vars=ordered_vars, + filename=params_filename, + ) + else: + paddle.static.save_vars( + Executor(_current_expected_place()), + dirname=model_path, + vars=list( + filter( + paddle.framework.io_utils.is_persistable, + ordered_vars, + ) + ), + filename=params_filename, + ) # save property property_save_path = os.path.join( os.path.normpath(model_path), file_prefix + INFER_PROPERTY_SUFFIX diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py index 8a2eaee72fd01e..19a660470cd255 100644 --- a/python/paddle/jit/dy2static/pir_partial_program.py +++ b/python/paddle/jit/dy2static/pir_partial_program.py @@ -473,6 +473,9 @@ def __init__( assert isinstance(self._build_strategy, BuildStrategy) self._origin_main_program = self._verify_program(main_program) + if parameters is not None: + parameters[0][:] = self._params + parameters[1][:] = self._param_values with paddle.base.framework._dygraph_guard(paddle.base.dygraph.Tracer()): self._cuda_graph_vec = self._create_cuda_graph_vec() self._cuda_graph_capture_mode = "" diff --git a/python/paddle/static/io_utils.py b/python/paddle/static/io_utils.py index 946d978c8a867e..5c77b032e7c194 100644 --- a/python/paddle/static/io_utils.py +++ b/python/paddle/static/io_utils.py @@ -21,7 +21,6 @@ from paddle.base import ( CompiledProgram, Variable, - default_main_program, ) @@ -66,7 +65,7 @@ def _get_valid_program(program=None): return default main program if program is None. """ if program is None: - program = default_main_program() + program = paddle.static.default_main_program() elif isinstance(program, CompiledProgram): program = program._program if program is None: diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt index 5549dd9426ee85..62ab3e3d38740b 100644 --- a/test/legacy_test/CMakeLists.txt +++ b/test/legacy_test/CMakeLists.txt @@ -855,7 +855,7 @@ set_tests_properties(test_vision_models PROPERTIES TIMEOUT 120) set_tests_properties(test_dataset_uci_housing PROPERTIES TIMEOUT 120) set_tests_properties(test_dataset_imdb PROPERTIES TIMEOUT 300) set_tests_properties(test_callback_wandb PROPERTIES TIMEOUT 60) -set_tests_properties(test_jit_save_load_rename PROPERTIES TIMEOUT 50) +set_tests_properties(test_jit_save_load_rename PROPERTIES TIMEOUT 100) if(WITH_COVERAGE) set_tests_properties(test_hapi_hub PROPERTIES TIMEOUT 300) endif() diff --git a/test/legacy_test/test_jit_save_load_rename.py b/test/legacy_test/test_jit_save_load_rename.py index 9b4d4fb592fd71..d723f8bfd58385 100644 --- a/test/legacy_test/test_jit_save_load_rename.py +++ b/test/legacy_test/test_jit_save_load_rename.py @@ -1050,8 +1050,7 @@ def test_load_model_retransform_inference(self): ) state_dict = multi_loaded_layer.state_dict() name_set = set() - for name, var in state_dict.items(): - # print("name = ", name, " var.name = ", var.name) + for _, var in state_dict.items(): self.assertTrue(var.name not in name_set) name_set.add(var.name) @@ -1484,7 +1483,6 @@ def test_save_load_no_param_layer(self): np.testing.assert_array_equal(out, load_out) -''' class TestJitSaveLoadMultiMethods(unittest.TestCase): def setUp(self): # enable dygraph mode @@ -1494,7 +1492,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() - #@test_with_dygraph_pir + # @test_with_dygraph_pir def test_jit_save_load_inference(self): model_path_inference = os.path.join( self.temp_dir.name, "jit_save_load_multi_methods/model" @@ -1505,6 +1503,14 @@ def test_jit_save_load_inference(self): layer, full_graph=True, ) + layer.forward_no_param = paddle.jit.to_static( + layer.forward_no_param, + full_graph=True, + ) + layer.forward_general = paddle.jit.to_static( + layer.forward_general, + full_graph=True, + ) inps = paddle.randn([1, IMAGE_SIZE]) result_origin = {} for func in dir(layer): @@ -1520,6 +1526,7 @@ def test_jit_save_load_inference(self): < 1e-5 ) + @test_with_dygraph_pir def test_jit_save_load_multi_methods_inputspec(self): model_path = os.path.join( self.temp_dir.name, 'jit_save_load_multi_methods/model' @@ -1529,6 +1536,14 @@ def test_jit_save_load_multi_methods_inputspec(self): layer, full_graph=True, ) + layer.forward_no_param = paddle.jit.to_static( + layer.forward_no_param, + full_graph=True, + ) + layer.forward_general = paddle.jit.to_static( + layer.forward_general, + full_graph=True, + ) with self.assertRaises(ValueError): paddle.jit.save( layer, model_path, input_spec=[InputSpec(shape=[None, 784])] @@ -1549,8 +1564,6 @@ def test_parse_name(self): self.assertFalse(hasattr(load_net, 'v2')) -''' - class LayerSaved(paddle.nn.Layer): def __init__(self, in_size, out_size): @@ -1572,87 +1585,74 @@ def forward(self, x): return self._linear_2(y) -class Net(paddle.nn.Layer): - def __init__(self): - super().__init__() - self.fc1 = paddle.nn.Linear(4, 4) - self.fc2 = paddle.nn.Linear(4, 4) - self.bias = 0.4 - self.flag = paddle.ones([2], dtype="int32") - - @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) - def log_softmax(self, input): - return paddle.nn.functional.log_softmax(input, axis=-1) - - @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) - def forward(self, x): - out = self.fc1(x) - out = paddle.nn.functional.relu(out) - out = paddle.mean(out) - return out - - @paddle.jit.to_static(input_spec=[InputSpec([None, 4], dtype='float32')]) - def infer(self, input): - out = self.fc2(input) - out = out + self.bias - out = paddle.mean(out) - return out - - # For extra Python float - @paddle.jit.to_static(property=True) - def fbias(self): - return self.bias + 1 - - @paddle.jit.to_static(property=True) - def down_sampling(self): - return 4 +class TestJitSaveCombineProperty(unittest.TestCase): + def setUp(self): + # enable dygraph mode + paddle.disable_static() + self.temp_dir = tempfile.TemporaryDirectory() - @paddle.jit.to_static(property=True) - def fstr(self): - return "save str property" + def tearDown(self): + self.temp_dir.cleanup() - @paddle.jit.to_static(property=True) - def ints(self): - return [10, 20] + @test_with_dygraph_pir + def test_jit_save_combine_property(self): + class Net(paddle.nn.Layer): + def __init__(self): + super().__init__() + self.fc1 = paddle.nn.Linear(4, 4) + self.fc2 = paddle.nn.Linear(4, 4) + self.bias = 0.4 + self.flag = paddle.ones([2], dtype="int32") - @paddle.jit.to_static(property=True) - def floats(self): - return [1.1, 2.2] + @paddle.jit.to_static( + input_spec=[InputSpec([None, 4], dtype='float32')] + ) + def log_softmax(self, input): + return paddle.nn.functional.log_softmax(input, axis=-1) - @paddle.jit.to_static(property=True) - def strs(self): - return ["hello", "world"] + @paddle.jit.to_static( + input_spec=[InputSpec([None, 4], dtype='float32')] + ) + def forward(self, x): + out = self.fc1(x) + out = paddle.nn.functional.relu(out) + out = paddle.mean(out) + return out + @paddle.jit.to_static( + input_spec=[InputSpec([None, 4], dtype='float32')] + ) + def infer(self, input): + out = self.fc2(input) + out = out + self.bias + out = paddle.mean(out) + return out -class NetTensor(paddle.nn.Layer): - def __init__(self): - super().__init__() - self.fc1 = paddle.nn.Linear(4, 4) - self.fc2 = paddle.nn.Linear(4, 4) - self.bias = 0.4 - self.flag = paddle.ones([2], dtype="int32") + # For extra Python float + @paddle.jit.to_static(property=True, full_graph=True) + def fbias(self): + return self.bias + 1 - def forward(self, x): - out = self.fc1(x) - out = paddle.nn.functional.relu(out) - out = paddle.mean(out) - return out + @paddle.jit.to_static(property=True, full_graph=True) + def down_sampling(self): + return 4 - @paddle.jit.to_static(property=True) - def fflag(self): - return True + @paddle.jit.to_static(property=True, full_graph=True) + def fstr(self): + return "save str property" + @paddle.jit.to_static(property=True, full_graph=True) + def ints(self): + return [10, 20] -class TestJitSaveCombineProperty(unittest.TestCase): - def setUp(self): - # enable dygraph mode - paddle.disable_static() - self.temp_dir = tempfile.TemporaryDirectory() + @paddle.jit.to_static(property=True, full_graph=True) + def floats(self): + return [1.1, 2.2] - def tearDown(self): - self.temp_dir.cleanup() + @paddle.jit.to_static(property=True, full_graph=True) + def strs(self): + return ["hello", "world"] - def test_jit_save_combine_property(self): model_path = os.path.join( self.temp_dir.name, "test_jit_save_combine/model" ) @@ -1662,7 +1662,26 @@ def test_jit_save_combine_property(self): # save paddle.jit.save(net, model_path, combine_params=True) + @test_with_dygraph_pir def test_jit_save_tensor_property(self): + class NetTensor(paddle.nn.Layer): + def __init__(self): + super().__init__() + self.fc1 = paddle.nn.Linear(4, 4) + self.fc2 = paddle.nn.Linear(4, 4) + self.bias = 0.4 + self.flag = paddle.ones([2], dtype="int32") + + def forward(self, x): + out = self.fc1(x) + out = paddle.nn.functional.relu(out) + out = paddle.mean(out) + return out + + @paddle.jit.to_static(property=True, full_graph=True) + def fflag(self): + return True + model_path = os.path.join( self.temp_dir.name, "test_jit_save_combine/model" ) @@ -2105,9 +2124,7 @@ def test_jit_save_compatible_input_sepc(self): ], full_graph=True, ) - save_dir = os.path.join( - self.temp_dir.name, "jit_save_compatible_input_spec" - ) + save_dir = os.path.join(self.temp_dir.name, "jit_save_no_input_spec") path = save_dir + "/model" paddle.jit.save(layer=layer, path=path) @@ -2115,6 +2132,9 @@ def test_jit_save_compatible_input_sepc(self): self._assert_input_spec_layer_return(layer, no_input_spec_layer) shutil.rmtree(save_dir) + save_dir = os.path.join(self.temp_dir.name, "jit_save_same_input_spec") + path = save_dir + "/model" + paddle.jit.save( layer=layer, path=path, @@ -2127,6 +2147,10 @@ def test_jit_save_compatible_input_sepc(self): self._assert_input_spec_layer_return(layer, same_input_spec_layer) shutil.rmtree(save_dir) + save_dir = os.path.join( + self.temp_dir.name, "jit_save_compatible_input_spec" + ) + path = save_dir + "/model" paddle.jit.save( layer=layer, path=path, From c19926e06dd101e6695defa2126219378a6b4035 Mon Sep 17 00:00:00 2001 From: wangruting Date: Mon, 6 May 2024 11:03:35 +0000 Subject: [PATCH 11/15] modify TestJitSaveLoadMultiMethods --- python/paddle/jit/api.py | 6 ++-- python/paddle/jit/pir_translated_layer.py | 23 +++++++++++++++ test/legacy_test/test_jit_save_load_rename.py | 28 +++++++++++++++++-- 3 files changed, 52 insertions(+), 5 deletions(-) diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py index a17eeb402febbe..11f5559ac60528 100644 --- a/python/paddle/jit/api.py +++ b/python/paddle/jit/api.py @@ -1216,8 +1216,8 @@ def save(layer, path, input_spec=None, **configs): params_filename = ( file_prefix + '.' + attr_func + INFER_PARAMS_SUFFIX ) - file_prefix = file_prefix + '.' + attr_func - file_prefix = os.path.join(model_path, file_prefix) + path_prefix = file_prefix + '.' + attr_func + file_path = os.path.join(model_path, path_prefix) with scope_guard(scope): if use_pir_api(): value_map = paddle.pir.IrMapping() @@ -1247,7 +1247,7 @@ def save(layer, path, input_spec=None, **configs): clone_output_vars = output_vars save_inference_model( - path_prefix=file_prefix, + path_prefix=file_path, feed_vars=clone_input_vars, fetch_vars=clone_output_vars, executor=Executor(_current_expected_place()), diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py index fa005b2ba1f106..b80eeb2dbdd4e0 100644 --- a/python/paddle/jit/pir_translated_layer.py +++ b/python/paddle/jit/pir_translated_layer.py @@ -29,6 +29,7 @@ from .translated_layer import ( BUFFER_NAME_PREFIX, + INFER_PARAMS_SUFFIX, PARAMETER_NAME_PREFIX, ) @@ -274,6 +275,27 @@ def _construct_params_and_buffers( var_dict = _load_pir_persistable_vars( model_path, programs['forward'], params_filename ) + model_name = params_filename[: -len(INFER_PARAMS_SUFFIX)] + # Load every file that meets the requirements in the directory model_path. + for file_name in os.listdir(model_path): + if file_name.startswith(model_name) and file_name.endswith( + INFER_PARAMS_SUFFIX + ): + parsing_names = file_name[ + len(model_name) : -len(INFER_PARAMS_SUFFIX) + 1 + ].split('.') + if len(parsing_names) == 3 and len(parsing_names[1]) > 0: + func_name = parsing_names[1] + else: + continue + else: + continue + + var_dict.update( + _load_pir_persistable_vars( + model_path, programs[func_name], file_name + ) + ) return var_dict @@ -306,6 +328,7 @@ def _run_dygraph(instance, input, program_holder): input_tensors.append(tensor) persistable_tensors = [] + for var_name in program_holder.persistable_names: dy_var_name = instance._persistable_var_name_dict[var_name] if dy_var_name in instance._parameters: diff --git a/test/legacy_test/test_jit_save_load_rename.py b/test/legacy_test/test_jit_save_load_rename.py index d723f8bfd58385..9de9ffedae1231 100644 --- a/test/legacy_test/test_jit_save_load_rename.py +++ b/test/legacy_test/test_jit_save_load_rename.py @@ -1492,7 +1492,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() - # @test_with_dygraph_pir + @test_with_dygraph_pir def test_jit_save_load_inference(self): model_path_inference = os.path.join( self.temp_dir.name, "jit_save_load_multi_methods/model" @@ -2114,7 +2114,7 @@ def _assert_input_spec_layer_return(self, expect_layer, test_layer): ) @test_with_dygraph_pir - def test_jit_save_compatible_input_sepc(self): + def test_jit_save_no_input_sepc(self): layer = InputSepcLayer() layer = paddle.jit.to_static( layer, @@ -2132,6 +2132,18 @@ def test_jit_save_compatible_input_sepc(self): self._assert_input_spec_layer_return(layer, no_input_spec_layer) shutil.rmtree(save_dir) + @test_with_dygraph_pir + def test_jit_save_same_input_sepc(self): + layer = InputSepcLayer() + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 8], dtype='float32', name='x'), + InputSpec(shape=[None, 1], dtype='float64', name='y'), + ], + full_graph=True, + ) + save_dir = os.path.join(self.temp_dir.name, "jit_save_same_input_spec") path = save_dir + "/model" @@ -2147,6 +2159,18 @@ def test_jit_save_compatible_input_sepc(self): self._assert_input_spec_layer_return(layer, same_input_spec_layer) shutil.rmtree(save_dir) + @test_with_dygraph_pir + def test_jit_save_compatible_input_sepc(self): + layer = InputSepcLayer() + layer = paddle.jit.to_static( + layer, + input_spec=[ + InputSpec(shape=[None, 8], dtype='float32', name='x'), + InputSpec(shape=[None, 1], dtype='float64', name='y'), + ], + full_graph=True, + ) + save_dir = os.path.join( self.temp_dir.name, "jit_save_compatible_input_spec" ) From ef3022bfb15e830f91fc3265e80511d87a2484c6 Mon Sep 17 00:00:00 2001 From: wangruting Date: Tue, 7 May 2024 06:57:51 +0000 Subject: [PATCH 12/15] modify cpu tensor no holder bug --- paddle/fluid/pybind/pir.cc | 57 +++++++++++-------- test/legacy_test/test_jit_save_load_rename.py | 5 +- 2 files changed, 36 insertions(+), 26 deletions(-) diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc index 34602582b38b84..810f56736186e4 100644 --- a/paddle/fluid/pybind/pir.cc +++ b/paddle/fluid/pybind/pir.cc @@ -121,6 +121,16 @@ PyTypeObject *g_ir_value_pytype = nullptr; void BindOpsAPI(pybind11::module *module); +pir::Value FakeValue() { + // create a fake value to simplify `ForwardBackwardSplit`. + return pir::Value(nullptr); +} + +bool IsFakeValue(const pir::Value &value) { + // create a fake value to simplify `ForwardBackwardSplit`. + return value.impl() == nullptr || !value.type(); +} + inline int64_t GetProgramInt64Attr(const std::shared_ptr &program, const std::string &attr_name, int64_t default_value = 0) { @@ -225,6 +235,21 @@ void SetValueName(Value value, const std::string name) { } } +bool HasValueName(const Value &value) { + if (IsFakeValue(value)) { + return false; + } + if (value.defining_op()->isa<::pir::ParameterOp>() || + value.defining_op()->isa() || + value.isa() || + (value.first_use() && + (value.first_use().owner()->isa<::pir::ShadowOutputOp>()))) { + return true; + } else { + return false; + } +} + std::string GetValueName(Value value) { if (auto param_op = value.defining_op<::pir::ParameterOp>()) { return param_op.param_name(); @@ -1002,19 +1027,8 @@ void BindValue(py::module *m) { "name", [](Value self) { return GetValueName(self); }, [](Value self, const std::string &name) { SetValueName(self, name); }) - .def_property_readonly( - "has_name", - [](Value self) { - if (self.defining_op()->isa<::pir::ParameterOp>() || - self.defining_op()->isa() || - self.isa() || - (self.first_use() && - self.first_use().owner()->isa<::pir::ShadowOutputOp>())) { - return true; - } else { - return false; - } - }) + .def_property_readonly("has_name", + [](Value self) { return HasValueName(self); }) .def_property( "shape", [](Value self) { return phi::vectorize(GetValueDims(self)); }, @@ -1508,16 +1522,6 @@ using SplitedProgram = std::vector>; using SplitedAttribute = std::map>; using SplitedResult = std::pair; -pir::Value FakeValue() { - // create a fake value to simplify `ForwardBackwardSplit`. - return pir::Value(nullptr); -} - -bool IsFakeValue(const pir::Value &value) { - // create a fake value to simplify `ForwardBackwardSplit`. - return value.impl() == nullptr || !value.type(); -} - static auto GetNoNeedBufferValue( const ::pir::Block *whole_block, std::vector range, @@ -1626,10 +1630,12 @@ int AppendShadowOutputs(Program *forward_program, std::string name_prefix) { int counter = 0; std::unordered_set added_value; - for (const auto &value : outputs) { if (!added_value.count(value) || IsFakeValue(value)) { std::string shadow_output_name = name_prefix + std::to_string(counter); + if (HasValueName(value)) { + shadow_output_name = GetValueName(value); + } AppendShadowOutput( forward_program, value, shadow_output_name, start_point + counter); counter += 1; @@ -1759,6 +1765,9 @@ SplitedResult SplitForwardBackward( } std::string shadow_output_name = std::string("output_") + std::to_string(counter); + if (HasValueName(v)) { + shadow_output_name = GetValueName(v); + } auto op_info = ctx->GetRegisteredOpInfo(pir::ShadowOutputOp::name()); pir::AttributeMap attribute_map = { {"output_name", pir::StrAttribute::get(ctx, shadow_output_name)}, diff --git a/test/legacy_test/test_jit_save_load_rename.py b/test/legacy_test/test_jit_save_load_rename.py index 9de9ffedae1231..aa0d0290131fbe 100644 --- a/test/legacy_test/test_jit_save_load_rename.py +++ b/test/legacy_test/test_jit_save_load_rename.py @@ -1516,6 +1516,7 @@ def test_jit_save_load_inference(self): for func in dir(layer): if func.startswith('forward'): result_origin[func] = getattr(layer, func, None)(inps) + paddle.jit.save(layer, model_path_inference) load_net = paddle.jit.load(model_path_inference) for func, result in result_origin.items(): @@ -1750,7 +1751,6 @@ def test_save_load_finetune_load(self): self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5) -''' class LayerLoadFinetune(paddle.nn.Layer): def __init__(self, in_size, out_size, load_path): super().__init__() @@ -1782,6 +1782,8 @@ def forward(self, x): y = self._load_l1(y) return y + +''' class TestJitSaveLoadFinetuneLoad(unittest.TestCase): def setUp(self): # enable dygraph mode @@ -1823,7 +1825,6 @@ def test_save_load_finetune_load(self): self.assertTrue(float((result_00 - result_10).abs().max()) < 1e-5) self.assertTrue(float((result_01 - result_11).abs().max()) < 1e-5) - ''' From 6547db0859d79321c98d67dce71221ecb55cc3fa Mon Sep 17 00:00:00 2001 From: wangruting Date: Wed, 8 May 2024 01:41:06 +0000 Subject: [PATCH 13/15] modify jit_save_load name --- test/dygraph_to_static/test_pylayer.py | 2 +- test/legacy_test/CMakeLists.txt | 4 ++-- test/{deprecated => }/legacy_test/test_io_save_load.py | 4 ++-- .../{test_jit_save_load_rename.py => test_jit_save_load.py} | 1 - 4 files changed, 5 insertions(+), 6 deletions(-) rename test/{deprecated => }/legacy_test/test_io_save_load.py (97%) rename test/legacy_test/{test_jit_save_load_rename.py => test_jit_save_load.py} (99%) diff --git a/test/dygraph_to_static/test_pylayer.py b/test/dygraph_to_static/test_pylayer.py index 9724eb2749e40e..bf09ba3db8a8cd 100644 --- a/test/dygraph_to_static/test_pylayer.py +++ b/test/dygraph_to_static/test_pylayer.py @@ -28,7 +28,7 @@ import unittest import numpy as np -from test_jit_save_load_rename import train +from test_jit_save_load import train import paddle from paddle.autograd.py_layer import PyLayer diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt index 6ac86fa7897dc2..46c21cd0351bb2 100644 --- a/test/legacy_test/CMakeLists.txt +++ b/test/legacy_test/CMakeLists.txt @@ -492,7 +492,7 @@ foreach(TEST_OP ${TEST_OPS_WITH_GC}) endforeach() # Switch some dy2st UT to eager mode -set(TEST_EAGER_OPS test_jit_save_load_rename test_translated_layer) +set(TEST_EAGER_OPS test_jit_save_load test_translated_layer) foreach(TEST_OP ${TEST_EAGER_OPS}) list(REMOVE_ITEM TEST_OPS ${TEST_OP}) py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS FLAGS_enable_eager_mode=1) @@ -855,7 +855,7 @@ set_tests_properties(test_vision_models PROPERTIES TIMEOUT 120) set_tests_properties(test_dataset_uci_housing PROPERTIES TIMEOUT 120) set_tests_properties(test_dataset_imdb PROPERTIES TIMEOUT 300) set_tests_properties(test_callback_wandb PROPERTIES TIMEOUT 60) -set_tests_properties(test_jit_save_load_rename PROPERTIES TIMEOUT 100) +set_tests_properties(test_jit_save_load PROPERTIES TIMEOUT 100) if(WITH_COVERAGE) set_tests_properties(test_hapi_hub PROPERTIES TIMEOUT 300) endif() diff --git a/test/deprecated/legacy_test/test_io_save_load.py b/test/legacy_test/test_io_save_load.py similarity index 97% rename from test/deprecated/legacy_test/test_io_save_load.py rename to test/legacy_test/test_io_save_load.py index a9eee37b281ded..7c2201423f9a7e 100644 --- a/test/deprecated/legacy_test/test_io_save_load.py +++ b/test/legacy_test/test_io_save_load.py @@ -19,7 +19,7 @@ import paddle from paddle import base, static from paddle.base import core -from paddle.pir_utils import test_with_pir_api +from paddle.pir_utils import test_with_dygraph_pir, test_with_pir_api class TestSaveLoadAPIError(unittest.TestCase): @@ -97,7 +97,7 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() - @test_with_pir_api + @test_with_dygraph_pir def test_when_train_with_no_grad(self): paddle.disable_static() net = paddle.nn.Linear(1024, 1) diff --git a/test/legacy_test/test_jit_save_load_rename.py b/test/legacy_test/test_jit_save_load.py similarity index 99% rename from test/legacy_test/test_jit_save_load_rename.py rename to test/legacy_test/test_jit_save_load.py index aa0d0290131fbe..14229e4d1120ac 100644 --- a/test/legacy_test/test_jit_save_load_rename.py +++ b/test/legacy_test/test_jit_save_load.py @@ -1793,7 +1793,6 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() - #@test_with_dygraph_pir def test_save_load_finetune_load(self): model_path = os.path.join( self.temp_dir.name, "test_jit_save_load_finetune_load/model" From 113854ca12818650170c09a68fecd56757970aa3 Mon Sep 17 00:00:00 2001 From: wangruting Date: Wed, 8 May 2024 01:53:40 +0000 Subject: [PATCH 14/15] move test_jit_save_load --- test/dygraph_to_static/test_pylayer.py | 2 +- .../{test_jit_save_load_rename.py => test_jit_save_load.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename test/legacy_test/{test_jit_save_load_rename.py => test_jit_save_load.py} (100%) diff --git a/test/dygraph_to_static/test_pylayer.py b/test/dygraph_to_static/test_pylayer.py index 9724eb2749e40e..bf09ba3db8a8cd 100644 --- a/test/dygraph_to_static/test_pylayer.py +++ b/test/dygraph_to_static/test_pylayer.py @@ -28,7 +28,7 @@ import unittest import numpy as np -from test_jit_save_load_rename import train +from test_jit_save_load import train import paddle from paddle.autograd.py_layer import PyLayer diff --git a/test/legacy_test/test_jit_save_load_rename.py b/test/legacy_test/test_jit_save_load.py similarity index 100% rename from test/legacy_test/test_jit_save_load_rename.py rename to test/legacy_test/test_jit_save_load.py From e25cabe9e97f3d835b5c28bd4fcd1213e97137c5 Mon Sep 17 00:00:00 2001 From: wangruting Date: Wed, 8 May 2024 05:47:05 +0000 Subject: [PATCH 15/15] modify cmake --- test/legacy_test/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt index 81e36703fafec2..fcea1be457055c 100644 --- a/test/legacy_test/CMakeLists.txt +++ b/test/legacy_test/CMakeLists.txt @@ -492,7 +492,7 @@ foreach(TEST_OP ${TEST_OPS_WITH_GC}) endforeach() # Switch some dy2st UT to eager mode -set(TEST_EAGER_OPS test_jit_save_load_rename test_translated_layer) +set(TEST_EAGER_OPS test_jit_save_load test_translated_layer) foreach(TEST_OP ${TEST_EAGER_OPS}) list(REMOVE_ITEM TEST_OPS ${TEST_OP}) py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS FLAGS_enable_eager_mode=1)