[PIR save/load]Add json_to_pdmodel (#68909)

changeyoung98 · web-flow · commit c74c3a92373c · 2024-10-25T16:33:41.000+08:00
* add json_to_pdmodel

* add json_to_pdmodel refine

* remove from public api
diff --git a/python/paddle/jit/__init__.py b/python/paddle/jit/__init__.py
@@ -13,7 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .api import ignore_module, load, not_to_static, save, to_static
+from .api import (
+    ignore_module,
+    json_to_pdmodel,  # noqa: F401
+    load,
+    not_to_static,
+    save,
+    to_static,
+)
 from .dy2static.logging_utils import set_code_level, set_verbosity
 from .dy2static.program_translator import enable_to_static
 from .translated_layer import TranslatedLayer
diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py
@@ -1801,3 +1801,19 @@ def get_ast_static_function(function):
             )
             return ast_static_function
     return function
+
+
+def json_to_pdmodel(net, input_spec, load_path, save_path):
+    net1 = paddle.jit.load(load_path)
+    state_dict = {}
+    for val in net1.state_dict().values():
+        name = val.name[: val.name.rfind('_')]
+        state_dict[name] = val
+
+    name_state_dict = {}
+    for name, var in net.to_static_state_dict().items():
+        name_state_dict[name] = state_dict[var.name]
+
+    net.set_state_dict(name_state_dict)
+    with paddle.pir_utils.OldIrGuard():
+        paddle.jit.save(net, save_path, input_spec)
diff --git a/test/ir/pir/test_ir_save_load.py b/test/ir/pir/test_ir_save_load.py
@@ -18,6 +18,8 @@
 
 import paddle
 from paddle import base
+from paddle.vision.models import ResNet
+from paddle.vision.models.resnet import BottleneckBlock
 
 
 class TestSaveModuleWithCommonOp(unittest.TestCase):
@@ -367,5 +369,32 @@ def internal_body(j, init, sums):
         )
 
 
+class TestJsonToPdmodel(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        paddle.disable_static()
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_json_to_pdmodel(self):
+        net = ResNet(BottleneckBlock, 50)
+        net = paddle.jit.to_static(net, full_graph=True)
+        save_json = os.path.join(self.temp_dir.name, 'save1')
+        save_model = os.path.join(self.temp_dir.name, 'save2')
+        input_spec = [
+            paddle.static.InputSpec(shape=[1, 3, 224, 224], dtype='float32')
+        ]
+        paddle.jit.save(net, save_json, input_spec)
+
+        # load and save to pdmodel
+        with paddle.pir_utils.OldIrGuard():
+            input_spec = [
+                paddle.static.InputSpec(shape=[1, 3, 224, 224], dtype='float32')
+            ]
+        paddle.jit.json_to_pdmodel(net, input_spec, save_json, save_model)
+        self.assertTrue(os.path.exists(save_model + '.pdmodel'))
+
+
 if __name__ == '__main__':
     unittest.main()
