diff --git a/python/paddle/fluid/dygraph/jit.py b/python/paddle/fluid/dygraph/jit.py index 90b0085fe330e86aabbb3ea37c6742563e0b7417..4b35d77845970378f0acdd21a3ce9c5d55c1ae27 100644 --- a/python/paddle/fluid/dygraph/jit.py +++ b/python/paddle/fluid/dygraph/jit.py @@ -25,7 +25,7 @@ import paddle from paddle.fluid import core from paddle.fluid.compiler import BuildStrategy, CompiledProgram, ExecutionStrategy from paddle.fluid.data_feeder import check_type -from paddle.fluid.layers.utils import flatten +from paddle.fluid.layers.utils import flatten, pack_sequence_as from paddle.fluid.dygraph.base import program_desc_tracing_guard, switch_to_static_graph from paddle.fluid.dygraph.dygraph_to_static import logging_utils from paddle.fluid.dygraph.dygraph_to_static.convert_call_func import ConversionOptions, CONVERSION_OPTIONS @@ -681,6 +681,11 @@ def save(layer, path, input_spec=None, **configs): inner_input_spec) elif 'forward' == attr_func: # transform in jit.save, if input_spec is incomplete, declarative will throw error + # inner_input_spec is list[InputSpec], it should be packed with same sturcture + # as original input_spec here. + if inner_input_spec: + inner_input_spec = pack_sequence_as(input_spec, + inner_input_spec) static_forward = declarative( inner_layer.forward, input_spec=inner_input_spec) concrete_program = static_forward.concrete_program diff --git a/python/paddle/fluid/tests/unittests/test_jit_save_load.py b/python/paddle/fluid/tests/unittests/test_jit_save_load.py index a43918765d44f02868910a8288277a302fa71e02..bf9912c89cb8736c17516d4498535f157fb2c914 100644 --- a/python/paddle/fluid/tests/unittests/test_jit_save_load.py +++ b/python/paddle/fluid/tests/unittests/test_jit_save_load.py @@ -222,6 +222,16 @@ class LinearNetWithDictInput(paddle.nn.Layer): return out +class LinearNetWithDictInputNoPrune(paddle.nn.Layer): + def __init__(self, in_size, out_size): + super(LinearNetWithDictInputNoPrune, self).__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(EmptyLayer, self).__init__() @@ -443,6 +453,30 @@ class TestSaveLoadWithDictInput(unittest.TestCase): self.assertEqual(len(loaded_net._input_spec()), 1) +class TestSaveLoadWithDictInputNoPrune(unittest.TestCase): + def test_dict_input(self): + net = LinearNetWithDictInputNoPrune(8, 8) + + path = "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) + + class TestSaveLoadWithInputSpec(unittest.TestCase): def setUp(self): # enable dygraph mode