diff --git a/python/paddle/framework/io.py b/python/paddle/framework/io.py index 493574c5bef47a50c633a95488d33fdd7b7dd29d..de2116cd4382d603e381bc4b0c24f72f93df0325 100644 --- a/python/paddle/framework/io.py +++ b/python/paddle/framework/io.py @@ -496,7 +496,7 @@ def save(obj, path, protocol=2, **configs): Save an object to the specified path. .. note:: - Now supports saving ``state_dict`` of Layer/Optimizer, Layer, Tensor and nested structure containing Tensor. + Now supports saving ``state_dict`` of Layer/Optimizer, Layer, Tensor and nested structure containing Tensor, Program. .. note:: Different from ``paddle.jit.save``, since the save result of ``paddle.save`` is a single file, @@ -544,7 +544,18 @@ def save(obj, path, protocol=2, **configs): # save weight of emb paddle.save(emb.weight, "emb.weight.pdtensor") - # example 2: static graph + # example 2: Save multiple state_dict at the same time + from paddle import nn + from paddle.optimizer import Adam + + layer = paddle.nn.Linear(3, 4) + adam = Adam(learning_rate=0.001, parameters=layer.parameters()) + obj = {'model': layer.state_dict(), 'opt': adam.state_dict(), 'epoch': 100} + path = 'example/model.pdparams' + paddle.save(obj, path) + + + # example 3: static graph import paddle import paddle.static as static @@ -570,6 +581,18 @@ def save(obj, path, protocol=2, **configs): # save/load state_dict path_state_dict = 'temp/model.pdparams' paddle.save(prog.state_dict("param"), path_tensor) + + # example 4: save program + import paddle + + paddle.enable_static() + + data = paddle.static.data( + name='x_static_save', shape=(None, 224), dtype='float32') + y_static = z = paddle.static.nn.fc(data, 10) + main_program = paddle.static.default_main_program() + path = "example/main_program.pdmodel" + paddle.save(main_program, path) ''' # 1. input check filename = os.path.basename(path) @@ -667,7 +690,7 @@ def load(path, **configs): Load an object can be used in paddle from specified path. .. note:: - Now supports loading ``state_dict`` of Layer/Optimizer, Layer, Tensor and nested structure containing Tensor. + Now supports loading ``state_dict`` of Layer/Optimizer, Layer, Tensor and nested structure containing Tensor, Program. .. note:: In order to use the model parameters saved by paddle more efficiently, @@ -714,8 +737,6 @@ def load(path, **configs): Examples: .. code-block:: python - import paddle - # example 1: dynamic graph import paddle emb = paddle.nn.Embedding(10, 10) @@ -744,7 +765,19 @@ def load(path, **configs): load_weight = paddle.load("emb.weight.pdtensor") - # example 2: static graph + # example 2: Load multiple state_dict at the same time + from paddle import nn + from paddle.optimizer import Adam + + layer = paddle.nn.Linear(3, 4) + adam = Adam(learning_rate=0.001, parameters=layer.parameters()) + obj = {'model': layer.state_dict(), 'opt': adam.state_dict(), 'epoch': 100} + path = 'example/model.pdparams' + paddle.save(obj, path) + obj_load = paddle.load(path) + + + # example 3: static graph import paddle import paddle.static as static @@ -773,6 +806,22 @@ def load(path, **configs): paddle.save(prog.state_dict("param"), path_tensor) load_state_dict = paddle.load(path_tensor) + + # example 4: load program + import paddle + + paddle.enable_static() + + data = paddle.static.data( + name='x_static_save', shape=(None, 224), dtype='float32') + y_static = z = paddle.static.nn.fc(data, 10) + main_program = paddle.static.default_main_program() + path = "example/main_program.pdmodel" + paddle.save(main_program, path) + load_main = paddle.load(path) + print(load_main) + + ''' if os.path.isfile(path):