io.py 93.8 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# 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
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13 14
# 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.

15
import os
T
bug fix  
tangwei12 已提交
16
import errno
D
dzhwinter 已提交
17
import warnings
18
import logging
Y
Yang Zhang 已提交
19
import pickle
H
hong 已提交
20
import contextlib
21
from functools import reduce
22
import sys
23
from io import BytesIO
24

H
hong 已提交
25
import numpy as np
26
import math
27
import paddle
28
from paddle.fluid import layers
H
hong 已提交
29
from paddle.fluid.executor import Executor, global_scope
30
from paddle.fluid.evaluator import Evaluator
31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
from paddle.fluid.framework import (
    Program,
    Parameter,
    default_main_program,
    default_startup_program,
    Variable,
    program_guard,
    dygraph_not_support,
    static_only,
)
from paddle.reader import (
    cache,
    map_readers,
    buffered,
    compose,
    chain,
    shuffle,
    ComposeNotAligned,
    firstn,
    xmap_readers,
    multiprocess_reader,
)
53
from .wrapped_decorator import signature_safe_contextmanager
T
tangwei12 已提交
54
from paddle.fluid.compiler import CompiledProgram
55
from paddle.fluid.log_helper import get_logger
S
sneaxiy 已提交
56
from . import reader
57
from . import unique_name
S
sneaxiy 已提交
58
from .reader import *
59 60
from . import dataloader
from .dataloader import *
K
fix bug  
Kexin Zhao 已提交
61
from . import core
62 63
from paddle.utils import deprecated
from paddle.fluid.framework import static_only
64

65 66
batch = paddle.batch

67
__all__ = [
68 69 70 71 72 73 74 75 76 77 78 79 80
    'save_vars',
    'save_params',
    'save_persistables',
    'load_vars',
    'load_params',
    'load_persistables',
    'save_inference_model',
    'load_inference_model',
    'batch',
    'save',
    'load',
    'load_program_state',
    'set_program_state',
H
hong 已提交
81 82
    'get_program_parameter',
    'get_program_persistable_vars',
83
] + reader.__all__
84

85 86 87
_logger = get_logger(
    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
)
88

89

90
class _open_buffer:
91 92 93 94 95 96 97 98 99
    def __init__(self, buffer):
        self.buffer = buffer

    def __enter__(self):
        return self.buffer


class _buffer_reader(_open_buffer):
    def __init__(self, buffer):
100
        super().__init__(buffer)
101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127
        self.initial_tell = self.buffer.tell()

    def __exit__(self, *args):
        # `args[0]` is type of exception. When the `read` is abnormal, the file pointer returns to the initial position.
        if args[0] is not None:
            self.buffer.seek(self.initial_tell)


class _buffer_writer(_open_buffer):
    def __exit__(self, *args):
        self.buffer.flush()


def _is_file_path(path):
    return isinstance(path, str)


def _open_file_buffer(path_or_buffer, mode):

    if _is_file_path(path_or_buffer):
        return open(path_or_buffer, mode)
    else:
        if 'w' in mode:
            return _buffer_writer(path_or_buffer)
        elif 'r' in mode:
            return _buffer_reader(path_or_buffer)
        else:
128
            raise ValueError(
129 130
                "Expected 'r' or 'w' in mode but got {}".format(mode)
            )
131 132 133 134 135 136


def _is_memory_buffer(buffer):
    return isinstance(buffer, BytesIO)


137
def is_parameter(var):
F
fengjiayi 已提交
138 139
    """
    Check whether the given variable is an instance of Parameter.
140 141

    Args:
F
fengjiayi 已提交
142
        var(Variable): The variable to be checked.
143 144

    Returns:
F
fengjiayi 已提交
145 146 147 148 149 150
        bool: True if the given `var` is an instance of Parameter,
        False if not.

    Examples:
        .. code-block:: python

151
            import paddle
152
            import paddle.fluid as fluid
153 154

            paddle.enable_static()
F
fengjiayi 已提交
155 156
            param = fluid.default_main_program().global_block().var('fc.w')
            res = fluid.io.is_parameter(param)
157
    """
158 159 160 161
    return isinstance(var, Parameter)


def is_persistable(var):
F
fengjiayi 已提交
162 163 164 165 166 167 168 169 170 171 172 173 174
    """
    Check whether the given variable is persistable.

    Args:
        var(Variable): The variable to be checked.

    Returns:
        bool: True if the given `var` is persistable
        False if not.

    Examples:
        .. code-block:: python

175
            import paddle
176
            import paddle.fluid as fluid
177 178

            paddle.enable_static()
179
            param = fluid.default_main_program().global_block().var('fc.b')
F
fengjiayi 已提交
180 181
            res = fluid.io.is_persistable(param)
    """
182 183 184 185 186
    if (
        var.desc.type() == core.VarDesc.VarType.FEED_MINIBATCH
        or var.desc.type() == core.VarDesc.VarType.FETCH_LIST
        or var.desc.type() == core.VarDesc.VarType.READER
    ):
187
        return False
188 189 190
    return var.persistable


H
hong 已提交
191
def is_belong_to_optimizer(var):
192
    if not (isinstance(var, Parameter) or var.desc.need_check_feed()):
193 194 195
        return is_persistable(var)

    return False
H
hong 已提交
196 197


198
@dygraph_not_support
H
hong 已提交
199 200 201 202 203 204 205 206 207 208 209 210 211
def get_program_parameter(program):
    """
    Get all the parameters from Program.

    Args:
        var(Program): The Program to get parameters

    Returns:
        list: The list contains all parameters in the program

    Examples:
        .. code-block:: python

212
            import paddle
H
hong 已提交
213
            import paddle.fluid as fluid
214 215

            paddle.enable_static()
H
hong 已提交
216
            data = fluid.data(name="img", shape=[64, 784])
217 218
            w = paddle.create_parameter(shape=[784, 200], dtype='float32', name='fc_w')
            b = paddle.create_parameter(shape=[200], dtype='float32', name='fc_b')
H
hong 已提交
219 220 221 222 223
            list_para  = fluid.io.get_program_parameter(  fluid.default_main_program() )
    """
    return list(filter(is_parameter, program.list_vars()))


224
@dygraph_not_support
H
hong 已提交
225 226 227 228 229 230 231 232 233 234 235 236 237
def get_program_persistable_vars(program):
    """
    Get all the persistable vars from Program.

    Args:
        var(Program): The Program to get persistable vars

    Returns:
        list: The list contains all persistable vars in the program

    Examples:
        .. code-block:: python

238
            import paddle
H
hong 已提交
239
            import paddle.fluid as fluid
240 241

            paddle.enable_static()
H
hong 已提交
242
            data = fluid.data(name="img", shape=[64, 784])
243 244
            w = paddle.create_parameter(shape=[784, 200], dtype='float32', name='fc_w')
            b = paddle.create_parameter(shape=[200], dtype='float32', name='fc_b')
H
hong 已提交
245 246 247 248 249
            list_para  = fluid.io.get_program_persistable_vars(  fluid.default_main_program() )
    """
    return list(filter(is_persistable, program.list_vars()))


250 251
def _clone_var_in_block_(block, var):
    assert isinstance(var, Variable)
252
    if var.desc.type() == core.VarDesc.VarType.LOD_TENSOR:
253 254 255 256 257 258 259 260
        return block.create_var(
            name=var.name,
            shape=var.shape,
            dtype=var.dtype,
            type=var.type,
            lod_level=var.lod_level,
            persistable=True,
        )
261
    else:
262 263 264 265 266 267 268
        return block.create_var(
            name=var.name,
            shape=var.shape,
            dtype=var.dtype,
            type=var.type,
            persistable=True,
        )
269 270


271
@signature_safe_contextmanager
H
hong 已提交
272 273 274 275 276 277 278
def _load_program_scope(main=None, startup=None, scope=None):
    prog = main if main else paddle.fluid.Program()
    startup_prog = startup if startup else paddle.fluid.Program()
    scope = scope if scope else paddle.fluid.core.Scope()
    with paddle.fluid.scope_guard(scope):
        with paddle.fluid.program_guard(prog, startup_prog):
            with paddle.fluid.unique_name.guard():
279 280
                with paddle.fluid.framework._dygraph_guard(None):
                    yield
H
hong 已提交
281 282


283
def _get_valid_program(main_program=None):
C
chengduo 已提交
284 285 286 287 288
    if main_program is None:
        main_program = default_main_program()
    elif isinstance(main_program, CompiledProgram):
        main_program = main_program._program
        if main_program is None:
289 290 291
            raise TypeError(
                "The type of input main_program is invalid, expected tyep is Program, but received None"
            )
C
chengduo 已提交
292
        warnings.warn(
293 294
            "The input is a CompiledProgram, this is not recommended."
        )
C
chengduo 已提交
295
    if not isinstance(main_program, Program):
296 297
        raise TypeError(
            "The type of input main_program is invalid, expected type is fluid.Program, but received %s"
298 299
            % type(main_program)
        )
C
chengduo 已提交
300 301 302
    return main_program


303
@dygraph_not_support
304 305 306 307 308 309 310 311
def save_vars(
    executor,
    dirname,
    main_program=None,
    vars=None,
    predicate=None,
    filename=None,
):
312
    """
313
    Save specific variables in the `Program` to files.
F
fengjiayi 已提交
314

315
    There are two ways to specify the variables to be saved: set variables in
316 317
    a list and assign it to the `vars`, or use the `predicate` function to select
    variables that make `predicate(variable) == True`. The first way has a higher priority.
318

319
    The `dirname` is used to specify the folder where to save variables.
T
tianshuo78520a 已提交
320
    If you prefer to save variables in separate files in the `dirname` folder,
321
    do not set `filename`. If you prefer to save all variables in a single file,
F
fengjiayi 已提交
322
    use `filename` to specify it.
323

F
fengjiayi 已提交
324 325
    Args:
        executor(Executor): The executor to run for saving variables.
326 327
        dirname(str, optional): The folder where to save variables.
                            When you need to save the parameter to the memory, set it to None.
328
        main_program(Program, optional): The program whose variables will be saved.
329
                                    If it is None, the default main program will
F
fengjiayi 已提交
330 331
                                    be used automatically.
                                    Default: None
332 333 334
        vars(list[Variable], optional): The list contains all variables to be saved.
                                        Default: None
        predicate(function, optional): The function selects the variables that make
335
                                       `predicate(variable) == True`.
336 337
                                       Default: None
        filename(str, optional): If you prefer to save all variables in a single file,
338
                                 use `filename` to specify it. Otherwise, let `filename` be None.
339
                                 Default: None
F
fengjiayi 已提交
340 341

    Returns:
342 343
        str: When saving parameters to a file, returns None.
             When saving parameters to memory, returns a binary string containing parameters.
F
fengjiayi 已提交
344 345 346 347 348 349 350

    Raises:
        TypeError: If `main_program` is not an instance of Program nor None.

    Examples:
        .. code-block:: python

351
            import paddle
352
            import paddle.fluid as fluid
353

354
            paddle.enable_static()
355 356 357
            main_prog = fluid.Program()
            startup_prog = fluid.Program()
            with fluid.program_guard(main_prog, startup_prog):
G
GGBond8488 已提交
358
                data = paddle.static.data(name="img", shape=[64, 784])
359 360
                w = paddle.create_parameter(shape=[784, 200], dtype='float32', name='fc_w')
                b = paddle.create_parameter(shape=[200], dtype='float32', name='fc_b')
361
                hidden_w = paddle.matmul(x=data, y=w)
H
HongyuJia 已提交
362
                hidden_b = paddle.add(hidden_w, b)
363 364 365
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            exe.run(startup_prog)
F
fengjiayi 已提交
366

367
            # The first usage: use `vars` to set the saved variables.
368 369
            var_list = [w, b]
            path = "./my_paddle_vars"
370
            fluid.io.save_vars(executor=exe, dirname=path, vars=var_list,
371 372 373 374 375 376 377 378 379 380
                            filename="vars_file")
            # w and b will be save in a file named "var_file".

            # The second usage: use `predicate` to select the saved variable.
            def name_has_fc(var):
                res = "fc" in var.name
                return res
            param_path = "./my_paddle_model"
            fluid.io.save_vars(executor=exe, dirname=param_path, main_program=main_prog, vars=None, predicate = name_has_fc)
            # all variables whose names contain "fc " are saved.
381
    """
382 383 384 385
    save_to_memory = False
    if dirname is None and filename is None:
        save_to_memory = True

C
chengduo 已提交
386
    main_program = _get_valid_program(main_program)
T
tangwei12 已提交
387

388
    if vars is None:
389 390 391 392 393 394 395
        return save_vars(
            executor,
            main_program=main_program,
            dirname=dirname,
            vars=list(filter(predicate, main_program.list_vars())),
            filename=filename,
        )
396
    else:
石晓伟 已提交
397
        params_var_name = "saved_params"
398 399 400 401 402 403 404
        # give warning when there is no var in model
        if len(list(vars)) == 0:
            warnings.warn(
                "no variable in your model, please ensure there are any variables in your model to save"
            )
            return None

405 406
        save_program = Program()
        save_block = save_program.global_block()
407 408

        save_var_map = {}
409
        for each_var in vars:
410 411 412
            # NOTE: don't save the variable which type is RAW
            if each_var.type == core.VarDesc.VarType.RAW:
                continue
413
            new_var = _clone_var_in_block_(save_block, each_var)
414
            if filename is None and save_to_memory is False:
415 416 417
                save_file_path = os.path.join(
                    os.path.normpath(dirname), new_var.name
                )
418 419 420 421
                save_block.append_op(
                    type='save',
                    inputs={'X': [new_var]},
                    outputs={},
422 423
                    attrs={'file_path': os.path.normpath(save_file_path)},
                )
424 425 426
            else:
                save_var_map[new_var.name] = new_var

427
        if filename is not None or save_to_memory:
428 429 430 431
            save_var_list = []
            for name in sorted(save_var_map.keys()):
                save_var_list.append(save_var_map[name])

432 433 434 435
            save_path = str()
            if save_to_memory is False:
                save_path = os.path.join(os.path.normpath(dirname), filename)

436 437 438
            saved_params = save_block.create_var(
                type=core.VarDesc.VarType.RAW, name=params_var_name
            )
439
            saved_params.desc.set_persistable(True)
440 441 442 443 444 445 446 447 448
            save_block.append_op(
                type='save_combine',
                inputs={'X': save_var_list},
                outputs={'Y': saved_params},
                attrs={
                    'file_path': save_path,
                    'save_to_memory': save_to_memory,
                },
            )
449

450
        # NOTE(zhiqiu): save op will add variable kLookupTablePath in save_program.desc,
451 452 453
        # which leads to diff on save_program and its desc. Call _sync_with_cpp
        # to keep consistency.
        save_program._sync_with_cpp()
454
        executor.run(save_program)
455 456
        if save_to_memory:
            return global_scope().find_var(params_var_name).get_bytes()
457 458


459
@dygraph_not_support
460
def save_params(executor, dirname, main_program=None, filename=None):
461
    """
462
    Save all parameters from the :code:`main_program` to
463
    the folder :code:`dirname` or file :code:`filename`. You can refer to
G
guofei 已提交
464
    :ref:`api_guide_model_save_reader_en` for more details.
F
fengjiayi 已提交
465

G
guofei 已提交
466 467 468
    Use the :code:`dirname` to specify the saving folder. If you would like to
    save parameters in separate files, set :code:`filename` None; if you would
    like to save all parameters in a single file, use :code:`filename` to specify
F
fengjiayi 已提交
469 470
    the file name.

471
    Note:
G
guofei 已提交
472
        Some variables are not Parameter while they are necessary for
473
        training, such as learning rate, global step, etc. So you can NOT save
G
guofei 已提交
474 475
        and continue your training just by :ref:`api_fluid_io_save_params`
        and :ref:`api_fluid_io_load_params`. Please use :ref:`api_fluid_io_save_persistables`
476 477 478
        and :ref:`api_fluid_io_load_persistables` instead.

        If you want to save your model for the inference, please use the
G
guofei 已提交
479 480
        :ref:`api_fluid_io_save_inference_model`. You can refer to
        :ref:`api_guide_model_save_reader_en` for more details.
F
fengjiayi 已提交
481 482

    Args:
483
        executor(Executor): The executor to run for saving parameters, You can
G
guofei 已提交
484
                            refer to :ref:`api_guide_executor_en`.
485 486
        dirname(str, optional): The saving directory path.
                            When you need to save the parameter to the memory, set it to None.
G
guofei 已提交
487
        main_program(Program, optional): The program whose parameters will be
488 489
                                         saved. You can refer to
                                         :ref:`api_guide_Program_en` for more
G
guofei 已提交
490 491 492 493 494 495 496
                                         details. If it is None, the default main
                                         program will be used.
                                         Default: None
        filename(str, optional): The file to save all parameters. If you prefer
                                 to save parameters in different files, set it
                                 to None.
                                 Default: None
F
fengjiayi 已提交
497 498

    Returns:
499 500
        str: When saving parameters to a file, returns None.
             When saving parameters to memory, returns a binary string containing parameters.
F
fengjiayi 已提交
501 502 503 504

    Examples:
        .. code-block:: python

505
            import paddle
H
Huihuang Zheng 已提交
506
            import paddle.fluid as fluid
507

508 509

            paddle.enable_static()
G
guofei 已提交
510 511 512 513
            params_path = "./my_paddle_model"
            image = fluid.data(name='img', shape=[None, 28, 28], dtype='float32')
            label = fluid.data(name='label', shape=[None, 1], dtype='int64')
            feeder = fluid.DataFeeder(feed_list=[image, label], place=fluid.CPUPlace())
C
Charles-hit 已提交
514
            predict = paddle.static.nn.fc(x=image, size=10, activation='softmax')
515

516 517 518 519
            loss = paddle.nn.functional.cross_entropy(
                input=predict, label=label,
                reduction='none', use_softmax=False
            )
520
            avg_loss = paddle.mean(loss)
521

F
fengjiayi 已提交
522
            exe = fluid.Executor(fluid.CPUPlace())
G
guofei 已提交
523 524
            exe.run(fluid.default_startup_program())
            fluid.io.save_params(executor=exe, dirname=params_path)
525 526
            # The parameters weights and bias of the fc layer in the network are going to
            # be saved in different files in the path "./my_paddle_model"
527
    """
528 529 530 531 532 533 534 535
    return save_vars(
        executor,
        dirname=dirname,
        main_program=main_program,
        vars=None,
        predicate=is_parameter,
        filename=filename,
    )
536 537


538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559
def _save_distributed_persistables(executor, dirname, main_program):
    """
    save_persistables for distributed training.
    the method will do things listed below:
    1.save part of persistable variables on trainer.
    2.receive "remote prefetch variables" from parameter servers and merge them.
    3.save "distributed lookup table" on parameter servers.
    4.receive "optimizer variables" from parameter servers and merge them.

    Args:
        executor(Executor): The executor to run for saving parameters.
        dirname(str): The saving directory path.
        main_program(Program): The program whose parameters will be
                            saved. the main_program must be the trainer_program
                            get after transpiler.

    Returns:
        None

    Examples:
        .. code-block:: python

560
            import paddle
561
            import paddle.fluid as fluid
562 563

            paddle.enable_static()
564 565 566 567 568 569 570 571 572 573
            exe = fluid.Executor(fluid.CPUPlace())
            param_path = "./my_paddle_model"
            t = distribute_transpiler.DistributeTranspiler()
            t.transpile(...)
            train_program = t.get_trainer_program()
            _save_distributed_persistables(executor=exe, dirname=param_path, main_program=train_program)
    """

    def __save_remote_params(executor, dirname, remote_params_map):
        """
T
tianshuo78520a 已提交
574
        receive params on pserver through rpc.
575 576 577 578 579 580 581 582 583 584
        if the params are be sliced, will concat them to one, then save it.
        """
        if not remote_params_map:
            return

        prog = Program()
        block = prog.global_block()

        # recv optimize vars from pserver
        for name, remote_params in remote_params_map.items():
T
tangwei12 已提交
585 586 587 588 589 590 591
            origin = remote_params[0].origin
            is_slice = remote_params[0].is_slice

            slices = [None] * len(remote_params)
            slice_varnames = [None] * len(remote_params)
            remote_varnames = [None] * len(remote_params)
            endpoints = [None] * len(remote_params)
592 593 594

            for idx, optimizer in enumerate(remote_params):
                block_id = optimizer.block_id
T
tangwei12 已提交
595
                slice = optimizer.slice
596 597 598
                endpoint = optimizer.endpoint

                index = block_id if is_slice else idx
T
tangwei12 已提交
599 600 601
                slices[index] = slice
                slice_varnames[index] = "{}.slice.{}".format(slice.name, idx)
                remote_varnames[index] = slice.name
602 603
                endpoints[index] = endpoint

T
tangwei12 已提交
604 605 606 607 608
            slice_shapes = []
            for slice in slices:
                tmp = [str(dim) for dim in slice.shape]
                slice_shapes.append(",".join(tmp))

609 610 611 612 613 614 615 616 617 618 619 620
            block.append_op(
                type='recv_save',
                attrs={
                    "trainer_id": 0,
                    "shape": origin.shape,
                    "slice_shapes": slice_shapes,
                    "slice_varnames": slice_varnames,
                    "remote_varnames": remote_varnames,
                    "endpoints": endpoints,
                    "file_path": os.path.join(dirname, origin.name),
                },
            )
T
tangwei12 已提交
621

622 623
        executor.run(prog)

624 625 626
    def __save_distributed_lookup_tables(
        executor, dirname, distributed_lookup_table, endpoints
    ):
627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642
        """
        because the distributed lookup table may too huge to merge and save at one place,
        it will be saved at parameter server independent respectively.

        the save directory is dirname/"__lookup_table__".

        """
        prog = Program()
        block = prog.global_block()

        # if there is lookup table, the trainer 0 will notify all pserver to save.
        lookup_table_filename = os.path.join(dirname, "__lookup_table__")
        attrs = {}
        attrs['epmap'] = endpoints
        attrs['dir'] = lookup_table_filename
        attrs['lookup_table'] = distributed_lookup_table
643 644 645
        block.append_op(
            type='checkpoint_notify', inputs={}, outputs={}, attrs=attrs
        )
646 647 648 649 650 651
        executor.run(prog)

    def __exclude_vars(exclude_var_names=[]):
        def is_valid(var):
            if var.name in exclude_var_names:
                return False
652 653 654 655 656
            if (
                var.desc.type() == core.VarDesc.VarType.FEED_MINIBATCH
                or var.desc.type() == core.VarDesc.VarType.FETCH_LIST
                or var.desc.type() == core.VarDesc.VarType.READER
            ):
657 658 659 660 661 662
                return False
            return var.persistable

        return is_valid

    if not isinstance(main_program, Program):
T
tangwei12 已提交
663
        raise TypeError("'main_program' should be an instance of Program.")
664 665 666 667 668 669

    if not main_program._is_distributed:
        raise ValueError(
            "'_save_distributed_persistables' just be designed for distributed training."
        )

670 671 672 673 674
    remote_params_map = (
        main_program._parameters_on_pservers.get_distributed_vars_by_vtypes(
            ["Optimizer", "RemotePrefetch"], groupby=True
        )
    )
675 676 677 678 679 680 681 682 683 684 685 686

    exclude_var_names = []
    if remote_params_map:
        exclude_var_names.extend(remote_params_map.keys())

    if main_program._distributed_lookup_table:
        if isinstance(main_program._distributed_lookup_table, list):
            exclude_var_names.extend(main_program._distributed_lookup_table)
        else:
            exclude_var_names.append(main_program._distributed_lookup_table)

    local_vars = list(
687 688 689 690 691
        filter(__exclude_vars(exclude_var_names), main_program.list_vars())
    )
    save_vars(
        executor, main_program=main_program, dirname=dirname, vars=local_vars
    )
692 693 694 695 696 697

    if main_program._is_chief:
        if remote_params_map:
            __save_remote_params(executor, dirname, remote_params_map)
        if main_program._distributed_lookup_table:
            __save_distributed_lookup_tables(
698 699 700 701 702
                executor,
                dirname,
                main_program._distributed_lookup_table,
                main_program._endpoints,
            )
703 704


705
@dygraph_not_support
706
def save_persistables(executor, dirname, main_program=None, filename=None):
707
    """
708 709
    Save all persistable variables from :code:`main_program` to
    the folder :code:`dirname` or file :code:`filename`. You can refer to
G
guofei 已提交
710
    :ref:`api_guide_model_save_reader_en` for more details. And then
711 712
    saves these persistables variables to the folder :code:`dirname` or file
    :code:`filename`.
F
fengjiayi 已提交
713

G
guofei 已提交
714
    The :code:`dirname` is used to specify the folder where persistable variables
715
    are going to be saved. If you would like to save variables in separate
G
guofei 已提交
716 717
    files, set :code:`filename` None; if you would like to save all variables in a
    single file, use :code:`filename` to specify the file name.
F
fengjiayi 已提交
718 719 720

    Args:
        executor(Executor): The executor to run for saving persistable variables.
721
                            You can refer to :ref:`api_guide_executor_en` for
G
guofei 已提交
722
                            more details.
723

724 725 726
        dirname(str, optional): The saving directory path.
                            When you need to save the parameter to the memory, set it to None.
        main_program(Program, optional): The program whose persistbale variables will
727
                                         be saved. You can refer to
G
guofei 已提交
728
                                         :ref:`api_guide_Program_en` for more details.
729
                                         If it is None, the default main program will
G
guofei 已提交
730 731 732 733 734
                                         be used.
                                         Default: None.
        filename(str, optional): The file to save all variables. If you prefer to
                                 save variables in different files, set it to None.
                                 Default: None.
F
fengjiayi 已提交
735 736

    Returns:
737 738
        str: When saving parameters to a file, returns None.
             When saving parameters to memory, returns a binary string containing parameters.
F
fengjiayi 已提交
739 740 741 742

    Examples:
        .. code-block:: python

743
            import paddle
H
Huihuang Zheng 已提交
744
            import paddle.fluid as fluid
745

746
            paddle.enable_static()
G
guofei 已提交
747 748 749 750 751
            dir_path = "./my_paddle_model"
            file_name = "persistables"
            image = fluid.data(name='img', shape=[None, 28, 28], dtype='float32')
            label = fluid.data(name='label', shape=[None, 1], dtype='int64')
            feeder = fluid.DataFeeder(feed_list=[image, label], place=fluid.CPUPlace())
752

C
Charles-hit 已提交
753
            predict = paddle.static.nn.fc(x=image, size=10, activation='softmax')
754 755 756 757
            loss = paddle.nn.functional.cross_entropy(
                input=predict, label=label,
                reduction='none', use_softmax=False
            )
758
            avg_loss = paddle.mean(loss)
F
fengjiayi 已提交
759
            exe = fluid.Executor(fluid.CPUPlace())
G
guofei 已提交
760 761
            exe.run(fluid.default_startup_program())
            fluid.io.save_persistables(executor=exe, dirname=dir_path, filename=file_name)
762
            # The persistables variables weights and bias in the fc layer of the network
G
guofei 已提交
763 764
            # are going to be saved in the same file named "persistables" in the path
            # "./my_paddle_model"
765
    """
766
    if main_program and main_program._is_distributed:
767 768 769
        return _save_distributed_persistables(
            executor, dirname=dirname, main_program=main_program
        )
770
    else:
771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788
        return save_vars(
            executor,
            dirname=dirname,
            main_program=main_program,
            vars=None,
            predicate=is_persistable,
            filename=filename,
        )


def load_vars(
    executor,
    dirname,
    main_program=None,
    vars=None,
    predicate=None,
    filename=None,
):
789
    """
790 791
    :api_attr: Static Graph

792
    This API loads variables from files by executor.
F
fengjiayi 已提交
793

794
    There are two ways to specify the variables to be loaded: the first way, set
795 796
    variables in a list and assign it to the `vars`; the second way, use the
    `predicate` function to select variables that make `predicate(variable) == True`.
797
    The first way has a higher priority.
F
fengjiayi 已提交
798

799
    The `dirname` is used to specify the folder where to load variables.
800
    If variables were saved in separate files in the folder `dirname`,
801
    set `filename` None. If all variables were saved in a single file,
F
fengjiayi 已提交
802
    use `filename` to specify it.
803

F
fengjiayi 已提交
804 805
    Args:
        executor(Executor): The executor to run for loading variables.
806 807
        dirname(str): The folder where to load the variables.
        main_program(Program, optional): The program whose variables will be loaded.
808
                                    If it is None, the default main program will
F
fengjiayi 已提交
809 810
                                    be used automatically.
                                    Default: None
811
        vars(list[Variable], optional): The list that contains all variables to be loaded.
F
fengjiayi 已提交
812
                                   Default: None
813
        predicate(function, optional): The function selects variables that make
814 815 816 817 818
                                        `predicate(variable) == True`.
                                        Default: None
        filename(str, optional): The file which saved all required variables. If variables
                                were saved in separate files, set it to be None.
                                Default: None
F
fengjiayi 已提交
819 820 821 822 823 824 825

    Returns:
        None

    Examples:
        .. code-block:: python

826
            import paddle
827
            import paddle.fluid as fluid
828

829
            paddle.enable_static()
830 831 832
            main_prog = fluid.Program()
            startup_prog = fluid.Program()
            with fluid.program_guard(main_prog, startup_prog):
G
GGBond8488 已提交
833
                data = paddle.static.data(name="img", shape=[64, 784])
834 835
                w = paddle.create_parameter(shape=[784, 200], dtype='float32', name='fc_w')
                b = paddle.create_parameter(shape=[200], dtype='float32', name='fc_b')
836
                hidden_w = paddle.matmul(x=data, y=w)
H
HongyuJia 已提交
837
                hidden_b = paddle.add(hidden_w, b)
838 839 840
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            exe.run(startup_prog)
F
fengjiayi 已提交
841

842 843 844 845 846 847 848 849 850 851 852
            # The first usage: using `vars` to specify the variables.
            path = "./my_paddle_vars"
            var_list = [w, b]
            fluid.io.save_vars(executor=exe, dirname=path, vars=var_list,
                               filename="vars_file")
            fluid.io.load_vars(executor=exe, dirname=path, vars=var_list,
                               filename="vars_file")
            # w and b will be loaded, and they are supposed to
            # be saved in the same file named 'var_file' in the path "./my_paddle_vars".

            # The second usage: using the `predicate` function to select variables
853
            param_path = "./my_paddle_model"
F
fengjiayi 已提交
854 855 856
            def name_has_fc(var):
                res = "fc" in var.name
                return res
857 858 859
            fluid.io.save_vars(executor=exe, dirname=param_path, main_program=main_prog,
                              vars=None, predicate=name_has_fc)
            fluid.io.load_vars(executor=exe, dirname=param_path, main_program=main_prog,
C
chengduo 已提交
860
                               vars=None, predicate=name_has_fc)
861 862
            # Load All variables in the `main_program` whose name includes "fc".
            # And all the variables are supposed to be saved in separate files.
F
fengjiayi 已提交
863

864
    """
865 866 867 868 869
    vars_from_memory = False
    if dirname is not None:
        dirname = os.path.normpath(dirname)
    else:
        vars_from_memory = True
T
tangwei12 已提交
870

871
    if vars is None:
872
        if main_program is None:
Y
Yu Yang 已提交
873
            main_program = default_main_program()
874
        if not isinstance(main_program, Program):
875 876
            raise TypeError(
                "The type of input main_program is invalid, expected type is fluid.Program, but received %s"
877 878
                % type(main_program)
            )
879

880 881 882 883 884 885 886
        load_vars(
            executor,
            dirname=dirname,
            main_program=main_program,
            vars=list(filter(predicate, main_program.list_vars())),
            filename=filename,
        )
887 888 889
    else:
        load_prog = Program()
        load_block = load_prog.global_block()
890

891 892
        if main_program is None:
            main_program = default_main_program()
T
tangwei12 已提交
893

894
        if not isinstance(main_program, Program):
895 896
            raise TypeError(
                "The type of input main_program is invalid, expected type is fluid.Program, but received %s"
897 898
                % type(main_program)
            )
899

T
tangwei12 已提交
900
        # save origin param shape
H
hong 已提交
901
        orig_para_shape = {}
902
        load_var_map = {}
903 904 905 906

        check_vars = []
        sparse_vars = []

907 908
        for each_var in vars:
            assert isinstance(each_var, Variable)
909

T
tangwei12 已提交
910 911
            if each_var.type == core.VarDesc.VarType.RAW:
                continue
H
hong 已提交
912 913

            if isinstance(each_var, Parameter):
914
                orig_para_shape[each_var.name] = tuple(
915 916
                    each_var.desc.get_shape()
                )
917 918 919 920 921

            if each_var.type == core.VarDesc.VarType.SELECTED_ROWS:
                sparse_vars.append(each_var)
                continue

922
            new_var = _clone_var_in_block_(load_block, each_var)
923 924
            check_vars.append(each_var)

925
            if filename is None:
926 927 928 929
                if dirname is None:
                    raise ValueError(
                        "The directory path and params cannot be None at the same time."
                    )
930 931 932 933
                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [new_var]},
934 935
                    attrs={'file_path': os.path.join(dirname, new_var.name)},
                )
936 937 938
            else:
                load_var_map[new_var.name] = new_var

939 940 941 942 943
        for each_var in sparse_vars:
            assert isinstance(each_var, Variable)

            if filename is not None:
                raise ValueError(
944 945
                    "SelectedRows can not be load with load_combine"
                )
946 947 948 949 950

            new_var = _clone_var_in_block_(load_block, each_var)

            var_path = os.path.join(dirname, new_var.name)
            if not os.path.exists(var_path):
951 952
                raise ValueError(
                    "SelectedRows var {} can not find at {}".format(
953 954 955
                        new_var.name, var_path
                    )
                )
956 957 958 959 960 961

            if os.path.isfile(var_path):
                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [new_var]},
962 963
                    attrs={'file_path': os.path.join(dirname, new_var.name)},
                )
964 965 966 967 968 969 970 971 972 973
            else:
                blocks = []
                block_paths = os.listdir(var_path)

                for block in block_paths:
                    if block.startswith(new_var.name):
                        blocks.append(block)

                slices = []
                for block in blocks:
974 975 976 977 978 979 980
                    slice = load_block.create_var(
                        name=block,
                        type=new_var.type,
                        shape=new_var.shape,
                        dtype=new_var.dtype,
                        persistable=False,
                    )
981 982 983
                    slices.append(slice)

                    file_path = os.path.join(var_path, block, "Param")
984 985 986 987 988 989
                    load_block.append_op(
                        type='load',
                        inputs={},
                        outputs={'Out': [slice]},
                        attrs={'file_path': file_path},
                    )
990

991 992 993 994 995 996
                load_block.append_op(
                    type='lookup_sparse_table_merge',
                    inputs={'X': slices},
                    outputs={'Out': new_var},
                    attrs={},
                )
997

998
        if filename is not None:
999 1000 1001 1002
            load_var_list = []
            for name in sorted(load_var_map.keys()):
                load_var_list.append(load_var_map[name])

1003 1004 1005
            if vars_from_memory is False:
                filename = os.path.join(dirname, filename)

1006 1007 1008 1009 1010 1011 1012 1013 1014
            load_block.append_op(
                type='load_combine',
                inputs={},
                outputs={"Out": load_var_list},
                attrs={
                    'file_path': filename,
                    'model_from_memory': vars_from_memory,
                },
            )
1015 1016
        executor.run(load_prog)

T
tangwei12 已提交
1017
        # check var shape
1018
        for each_var in check_vars:
H
hong 已提交
1019 1020 1021
            if not isinstance(each_var, Parameter):
                continue
            var_temp = paddle.fluid.global_scope().find_var(each_var.name)
1022
            assert var_temp is not None, "can't not find var: " + each_var.name
H
hong 已提交
1023
            new_shape = (np.array(var_temp.get_tensor())).shape
1024 1025 1026
            assert each_var.name in orig_para_shape, (
                each_var.name + "MUST in var list"
            )
H
hong 已提交
1027 1028 1029
            orig_shape = orig_para_shape.get(each_var.name)
            if new_shape != orig_shape:
                raise RuntimeError(
1030
                    "Variable's shape does not match, the Program requires a parameter with the shape of ({}), "
1031 1032 1033 1034
                    "while the loaded parameter (namely [ {} ]) has a shape of  ({}).".format(
                        orig_shape, each_var.name, new_shape
                    )
                )
H
hong 已提交
1035

1036

1037
@dygraph_not_support
1038
def load_params(executor, dirname, main_program=None, filename=None):
1039
    """
1040 1041
    :api_attr: Static Graph

1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060
    This API filters out all parameters from the give ``main_program``
    and then tries to load these parameters from the directory ``dirname`` or
    the file ``filename``.

    Use the ``dirname`` to specify the directory where parameters were saved. If
    parameters were saved in separate files under the directory `dirname`, set
    ``filename`` as None; if all parameters were saved in a single file, use
    ``filename`` to specify the file name.

    **Note**:
        Some variables are not Parameter while they are necessary for
        training, such as learning rate, global step, etc. So you cannot save and
        continue your training just by using :ref:`api_fluid_io_save_params` and
        :ref:`api_fluid_io_load_params`. Please use :ref:`api_fluid_io_save_persistables`
        and :ref:`api_fluid_io_load_persistables` instead.

        If you want to load the pre-trained model structure and parameters
        for the inference, please use the :ref:`api_fluid_io_load_inference_model` API. You can
        refer to :ref:`api_guide_model_save_reader_en` for more details.
F
fengjiayi 已提交
1061 1062

    Args:
1063 1064
        executor(Executor): The executor used for loading parameters.
                            See :ref:`api_guide_executor_en` for more details about it.
F
fengjiayi 已提交
1065
        dirname(str): The directory path.
1066 1067 1068 1069 1070 1071 1072 1073
        main_program(Program, optional): The program whose parameters will be
                                    loaded. If it is None, the ``default_main_program``
                                    will be used automatically. See :ref:`api_guide_Program_en`
                                    for more about ``Program``.
                                    Default: None.
        filename(str, optional): The file which saved all parameters. If parameters
                            were saved in separated files, set it to None.
                            Default: None.
F
fengjiayi 已提交
1074 1075 1076 1077 1078 1079 1080

    Returns:
        None

    Examples:
        .. code-block:: python

1081
            import paddle
1082
            import paddle.fluid as fluid
1083

1084
            paddle.enable_static()
F
fengjiayi 已提交
1085 1086 1087
            exe = fluid.Executor(fluid.CPUPlace())
            param_path = "./my_paddle_model"
            prog = fluid.default_main_program()
1088
            fluid.io.load_params(executor=exe, dirname=param_path,
F
fengjiayi 已提交
1089
                                main_program=None)
1090
    """
1091 1092 1093 1094 1095 1096 1097
    load_vars(
        executor,
        dirname=dirname,
        main_program=main_program,
        predicate=is_parameter,
        filename=filename,
    )
1098 1099


1100
@dygraph_not_support
1101
def load_persistables(executor, dirname, main_program=None, filename=None):
1102
    """
1103
    :api_attr: Static Graph
1104

1105 1106
    This API filters out all variables with ``persistable==True`` from the
    given ``main_program`` and then tries to load these variables from the
T
tianshuo78520a 已提交
1107
    directory ``dirname`` or the file ``filename``.
F
fengjiayi 已提交
1108

1109 1110 1111 1112
    Use the ``dirname`` to specify the directory where persistable variables
    (refer to :ref:`api_guide_model_save_reader_en`) were saved. If variables
    were saved in separate files, set ``filename`` as None; if all variables
    were saved in a single file, use ``filename`` to specify the file name.
F
fengjiayi 已提交
1113 1114

    Args:
1115 1116
        executor(Executor): The executor used for loading persistable variables.
                            See :ref:`api_guide_executor_en` for more details about it.
F
fengjiayi 已提交
1117
        dirname(str): The directory path.
T
tianshuo78520a 已提交
1118
        main_program(Program, optional): The program whose persistable variables will
1119 1120 1121 1122 1123 1124 1125
                                    be loaded. If it is None, the ``default_main_program``
                                    will be used automatically. See :ref:`api_guide_Program_en`
                                    for more about ``Program``.
                                    Default: None.
        filename(str, optional): The file which saved all persistable variables. If variables
                                 were saved in separated files, set it to None.
                                 Default: None.
F
fengjiayi 已提交
1126 1127 1128 1129 1130 1131 1132

    Returns:
        None

    Examples:
        .. code-block:: python

1133
            import paddle
1134
            import paddle.fluid as fluid
1135

1136
            paddle.enable_static()
F
fengjiayi 已提交
1137 1138 1139
            exe = fluid.Executor(fluid.CPUPlace())
            param_path = "./my_paddle_model"
            prog = fluid.default_main_program()
1140
            fluid.io.load_persistables(executor=exe, dirname=param_path,
F
fengjiayi 已提交
1141
                                       main_program=None)
1142
    """
1143 1144

    if main_program and main_program._is_distributed:
1145 1146 1147
        _load_distributed_persistables(
            executor, dirname=dirname, main_program=main_program
        )
1148
    else:
1149 1150 1151 1152 1153 1154 1155
        load_vars(
            executor,
            dirname=dirname,
            main_program=main_program,
            predicate=is_persistable,
            filename=filename,
        )
1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175


def _load_distributed_persistables(executor, dirname, main_program=None):
    """
    customized load_persistables for distributed training.
    it should be used on parameter server,

    Args:
        executor(Executor): The executor to run for saving parameters.
        dirname(str): The load directory path.
        main_program(Program): The program whose parameters will be
                            loaded. the main_program must be the pserver_program
                            get after transpiler.

    Returns:
        None

    Examples:
        .. code-block:: python

1176
            import paddle
1177
            import paddle.fluid as fluid
1178 1179

            paddle.enable_static()
1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204
            exe = fluid.Executor(fluid.CPUPlace())
            param_path = "./my_paddle_model"
            t = distribute_transpiler.DistributeTranspiler()
            t.transpile(...)
            pserver_prog = t.get_pserver_program(...)
            _load_distributed_persistables(executor=exe, dirname=param_path, main_program=pserver_prog)
    """

    def __is_distributed_part_var(varname):
        trainer_idx = varname.find(".trainer_")
        block_idx = varname.find(".block")
        return trainer_idx or block_idx

    def __load_persistable_vars(executor, dirname, need_load_vars):
        load_prog = Program()
        load_block = load_prog.global_block()
        need_delete_vars = []

        for param in need_load_vars:
            origin_var = param.origin
            slice_var = param.slice
            is_slice = param.is_slice
            offset = param.offset

            if is_slice:
1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222
                slice = load_block.create_var(
                    name=slice_var.name,
                    type=slice_var.type,
                    shape=slice_var.shape,
                    dtype=slice_var.dtype,
                    persistable=True,
                )

                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [slice]},
                    attrs={
                        'file_path': os.path.join(dirname, origin_var.name),
                        'seek': offset,
                        'shape': slice.shape,
                    },
                )
1223
            else:
1224 1225 1226 1227 1228 1229 1230
                origin = load_block.create_var(
                    name="{}".format(origin_var.name),
                    type=origin_var.type,
                    shape=origin_var.shape,
                    dtype=origin_var.dtype,
                    persistable=True,
                )
1231 1232 1233 1234
                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [origin]},
1235 1236
                    attrs={'file_path': os.path.join(dirname, origin_var.name)},
                )
1237 1238 1239

        load_block.append_op(
            type='delete_var',
1240 1241
            inputs={'X': need_delete_vars},
        )
1242 1243 1244 1245

        executor.run(load_prog)

    if not isinstance(main_program, Program):
T
tangwei12 已提交
1246
        raise TypeError("'main_program' should be an instance of Program.")
1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257

    if not main_program._is_distributed:
        raise ValueError(
            "'_load_distributed_persistables' just be designed for distributed training."
        )

    if not main_program._ps_endpoint:
        raise ValueError(
            "'_load_distributed_persistables' need current_endpoint set in DistributeTranspiler.transpile"
        )

1258 1259 1260 1261 1262
    need_load_vars = (
        main_program._parameters_on_pservers.get_distributed_vars_by_ep(
            main_program._ps_endpoint
        )
    )
1263
    __load_persistable_vars(executor, dirname, need_load_vars)
1264 1265


1266 1267 1268
def prepend_feed_ops(
    inference_program, feed_target_names, feed_holder_name='feed'
):
Q
Qiao Longfei 已提交
1269 1270 1271
    if len(feed_target_names) == 0:
        return

K
Kexin Zhao 已提交
1272
    global_block = inference_program.global_block()
1273 1274 1275 1276 1277
    feed_var = global_block.create_var(
        name=feed_holder_name,
        type=core.VarDesc.VarType.FEED_MINIBATCH,
        persistable=True,
    )
K
Kexin Zhao 已提交
1278

1279
    for i, name in enumerate(feed_target_names):
1280 1281 1282 1283
        if not global_block.has_var(name):
            raise ValueError(
                "The feeded_var_names[{i}]: '{name}' doesn't exist in pruned inference program. "
                "Please check whether '{name}' is a valid feed_var name, or remove it from feeded_var_names "
1284 1285 1286 1287
                "if '{name}' is not involved in the target_vars calculation.".format(
                    i=i, name=name
                )
            )
K
fix bug  
Kexin Zhao 已提交
1288
        out = global_block.var(name)
1289 1290 1291 1292 1293 1294
        global_block._prepend_op(
            type='feed',
            inputs={'X': [feed_var]},
            outputs={'Out': [out]},
            attrs={'col': i},
        )
K
Kexin Zhao 已提交
1295 1296


1297 1298 1299
def append_fetch_ops(
    inference_program, fetch_target_names, fetch_holder_name='fetch'
):
K
Kexin Zhao 已提交
1300
    global_block = inference_program.global_block()
1301 1302 1303 1304 1305
    fetch_var = global_block.create_var(
        name=fetch_holder_name,
        type=core.VarDesc.VarType.FETCH_LIST,
        persistable=True,
    )
K
Kexin Zhao 已提交
1306

1307
    for i, name in enumerate(fetch_target_names):
1308 1309 1310 1311 1312 1313
        global_block.append_op(
            type='fetch',
            inputs={'X': [name]},
            outputs={'Out': [fetch_var]},
            attrs={'col': i},
        )
K
Kexin Zhao 已提交
1314 1315


1316 1317
@static_only
@deprecated(since="2.0.0", update_to="paddle.static.save_inference_model")
1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329
def save_inference_model(
    dirname,
    feeded_var_names,
    target_vars,
    executor,
    main_program=None,
    model_filename=None,
    params_filename=None,
    export_for_deployment=True,
    program_only=False,
    clip_extra=True,
):
1330
    """
F
fengjiayi 已提交
1331
    Prune the given `main_program` to build a new program especially for inference,
G
guofei 已提交
1332
    and then save it and all related parameters to given `dirname` .
1333
    If you just want to save parameters of your trained model, please use the
G
guofei 已提交
1334 1335
    :ref:`api_fluid_io_save_params` . You can refer to :ref:`api_guide_model_save_reader_en`
    for more details.
1336

G
guofei 已提交
1337
    Note:
1338
        The :code:`dirname` is used to specify the folder where inference model
G
guofei 已提交
1339
        structure and parameters are going to be saved. If you would like to save params of
1340
        Program in separate files, set `params_filename` None; if you would like to save all
G
guofei 已提交
1341
        params of Program in a single file, use `params_filename` to specify the file name.
F
fengjiayi 已提交
1342 1343 1344

    Args:
        dirname(str): The directory path to save the inference model.
T
tianshuo78520a 已提交
1345
        feeded_var_names(list[str]): list of string. Names of variables that need to be fed
G
guofei 已提交
1346
                                     data during inference.
1347
        target_vars(list[Variable]): list of Variable. Variables from which we can get
G
guofei 已提交
1348
                                     inference results.
1349
        executor(Executor): The executor that saves the inference model. You can refer
G
guofei 已提交
1350 1351
                            to :ref:`api_guide_executor_en` for more details.
        main_program(Program, optional): The original program, which will be pruned to
T
tianshuo78520a 已提交
1352
                                         build the inference model. If is set None,
G
guofei 已提交
1353 1354 1355
                                         the global default :code:`_main_program_` will be used.
                                         Default: None.
        model_filename(str, optional): The name of file to save the inference program
T
tianshuo78520a 已提交
1356
                                       itself. If is set None, a default filename
G
guofei 已提交
1357 1358
                                       :code:`__model__` will be used.
        params_filename(str, optional): The name of file to save all related parameters.
T
tianshuo78520a 已提交
1359
                                        If it is set None, parameters will be saved
G
guofei 已提交
1360
                                        in separate files .
1361
        export_for_deployment(bool, optional): If True, programs are modified to only support
X
Xin Pan 已提交
1362 1363 1364 1365
                                     direct inference deployment. Otherwise,
                                     more information will be stored for flexible
                                     optimization and re-training. Currently, only
                                     True is supported.
G
guofei 已提交
1366
                                     Default: True.
1367
        program_only(bool, optional): If True, It will save inference program only, and do not
G
guofei 已提交
1368 1369
                                      save params of Program.
                                      Default: False.
1370

F
fengjiayi 已提交
1371
    Returns:
1372
        list, The fetch variables' name list.
F
fengjiayi 已提交
1373 1374 1375

    Examples:
        .. code-block:: python
F
fengjiayi 已提交
1376

1377
            import paddle
1378 1379
            import paddle.fluid as fluid

1380
            paddle.enable_static()
F
fengjiayi 已提交
1381 1382
            path = "./infer_model"

T
tianshuo78520a 已提交
1383
            # User defined network, here a softmax regession example
G
guofei 已提交
1384 1385
            image = fluid.data(name='img', shape=[None, 28, 28], dtype='float32')
            label = fluid.data(name='label', shape=[None, 1], dtype='int64')
1386
            feeder = fluid.DataFeeder(feed_list=[image, label], place=fluid.CPUPlace())
C
Charles-hit 已提交
1387
            predict = paddle.static.nn.fc(x=image, size=10, activation='softmax')
1388

1389 1390 1391 1392
            loss = paddle.nn.functional.cross_entropy(
                input=predict, label=label,
                reduction='none', use_softmax=False
            )
1393
            avg_loss = paddle.mean(loss)
1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405

            exe = fluid.Executor(fluid.CPUPlace())
            exe.run(fluid.default_startup_program())

            # Feed data and train process

            # Save inference model. Note we don't save label and loss in this example
            fluid.io.save_inference_model(dirname=path,
                                          feeded_var_names=['img'],
                                          target_vars=[predict],
                                          executor=exe)

G
guofei 已提交
1406
            # In this example, the save_inference_mode inference will prune the default
1407
            # main program according to the network's input node (img) and output node(predict).
G
guofei 已提交
1408
            # The pruned inference program is going to be saved in the "./infer_model/__model__"
F
fengjiayi 已提交
1409
            # and parameters are going to be saved in separate files under folder
1410
            # "./infer_model".
1411 1412

    """
1413
    if isinstance(feeded_var_names, str):
F
fengjiayi 已提交
1414
        feeded_var_names = [feeded_var_names]
X
Xin Pan 已提交
1415
    elif export_for_deployment:
Q
Qiao Longfei 已提交
1416
        if len(feeded_var_names) > 0:
1417
            # TODO(paddle-dev): polish these code blocks
1418 1419 1420 1421
            if not (
                bool(feeded_var_names)
                and all(isinstance(name, str) for name in feeded_var_names)
            ):
M
minqiyang 已提交
1422
                raise ValueError("'feed_var_names' should be a list of str.")
F
fengjiayi 已提交
1423 1424

    if isinstance(target_vars, Variable):
F
fengjiayi 已提交
1425
        target_vars = [target_vars]
X
Xin Pan 已提交
1426
    elif export_for_deployment:
1427 1428 1429 1430
        if not (
            bool(target_vars)
            and all(isinstance(var, Variable) for var in target_vars)
        ):
F
fengjiayi 已提交
1431 1432
            raise ValueError("'target_vars' should be a list of Variable.")

C
chengduo 已提交
1433
    main_program = _get_valid_program(main_program)
T
tangwei12 已提交
1434

1435
    # remind user to set auc_states to zeros if the program contains auc op
1436 1437
    all_ops = main_program.global_block().ops
    for op in all_ops:
1438 1439 1440
        # clear device of Op
        device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName()
        op._set_attr(device_attr_name, "")
1441 1442 1443 1444 1445 1446
        if op.type == 'auc':
            warnings.warn(
                "please ensure that you have set the auc states to zeros before saving inference model"
            )
            break

1447 1448
    with program_guard(main_program):
        uniq_target_vars = []
F
flame 已提交
1449 1450
        for i, var in enumerate(target_vars):
            uniq_target_vars.append(var)
1451
        target_vars = uniq_target_vars
F
flame 已提交
1452
    target_var_name_list = [var.name for var in target_vars]
1453

1454
    # when a pserver and a trainer running on the same machine, mkdir may conflict
L
lujun 已提交
1455
    save_dirname = dirname
1456
    try:
L
lujun 已提交
1457 1458
        save_dirname = os.path.normpath(dirname)
        os.makedirs(save_dirname)
1459 1460 1461 1462
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise

X
Xin Pan 已提交
1463 1464 1465 1466
    if model_filename is not None:
        model_basename = os.path.basename(model_filename)
    else:
        model_basename = "__model__"
L
lujun 已提交
1467
    model_basename = os.path.join(save_dirname, model_basename)
1468

X
Xin Pan 已提交
1469 1470 1471 1472
    # When export_for_deployment is true, we modify the program online so that
    # it can only be loaded for inference directly. If it's false, the whole
    # original program and related meta are saved so that future usage can be
    # more flexible.
1473 1474 1475

    origin_program = main_program.clone()

X
Xin Pan 已提交
1476
    if export_for_deployment:
X
Xin Pan 已提交
1477 1478
        main_program = main_program.clone()
        global_block = main_program.global_block()
1479
        need_to_remove_op_index = []
X
Xin Pan 已提交
1480 1481 1482
        for i, op in enumerate(global_block.ops):
            op.desc.set_is_target(False)
            if op.type == "feed" or op.type == "fetch":
1483 1484 1485 1486 1487
                need_to_remove_op_index.append(i)

        for index in need_to_remove_op_index[::-1]:
            global_block._remove_op(index)

X
Xin Pan 已提交
1488
        main_program.desc.flush()
X
Xin Pan 已提交
1489

1490
        main_program = main_program._prune_with_input(
1491 1492
            feeded_var_names=feeded_var_names, targets=target_vars
        )
X
Xin Pan 已提交
1493
        main_program = main_program._inference_optimize(prune_read_op=True)
X
Xin Pan 已提交
1494 1495
        fetch_var_names = [v.name for v in target_vars]

1496 1497 1498 1499 1500
        for target_v in target_vars:
            if not main_program.global_block().has_var(target_v.name):
                main_program.global_block().create_var(
                    name=target_v.name,
                    shape=target_v.shape,
1501
                    dtype=target_v.dtype,
1502 1503
                    persistable=target_v.persistable,
                )
1504

X
Xin Pan 已提交
1505 1506 1507
        prepend_feed_ops(main_program, feeded_var_names)
        append_fetch_ops(main_program, fetch_var_names)

1508
        main_program.desc._set_version()
1509
        paddle.fluid.core.save_op_version_info(main_program.desc)
X
Xin Pan 已提交
1510
        with open(model_basename, "wb") as f:
1511
            f.write(
1512
                main_program._remove_training_info(
1513 1514 1515
                    clip_extra=clip_extra
                ).desc.serialize_to_string()
            )
X
Xin Pan 已提交
1516 1517 1518
    else:
        # TODO(panyx0718): Save more information so that it can also be used
        # for training and more flexible post-processing.
X
Xin Pan 已提交
1519
        with open(model_basename + ".main_program", "wb") as f:
1520
            f.write(
1521
                main_program._remove_training_info(
1522 1523 1524
                    clip_extra=clip_extra
                ).desc.serialize_to_string()
            )
T
tangwei12 已提交
1525

T
tangwei12 已提交
1526 1527 1528 1529 1530 1531
    if program_only:
        warnings.warn(
            "save_inference_model specified the param `program_only` to True, It will not save params of Program."
        )
        return target_var_name_list

1532 1533
    main_program._copy_dist_param_info_from(origin_program)

X
fix  
Xin Pan 已提交
1534 1535
    if params_filename is not None:
        params_filename = os.path.basename(params_filename)
1536

L
lujun 已提交
1537
    save_persistables(executor, save_dirname, main_program, params_filename)
F
flame 已提交
1538
    return target_var_name_list
X
fix  
Xin Pan 已提交
1539

1540

1541 1542
@static_only
@deprecated(since="2.0.0", update_to="paddle.static.load_inference_model")
1543 1544 1545 1546 1547 1548 1549
def load_inference_model(
    dirname,
    executor,
    model_filename=None,
    params_filename=None,
    pserver_endpoints=None,
):
1550
    """
1551 1552 1553
    Load the inference model from a given directory. By this API, you can get the model
    structure(Inference Program) and model parameters. If you just want to load
    parameters of the pre-trained model, please use the :ref:`api_fluid_io_load_params` API.
1554
    You can refer to :ref:`api_guide_model_save_reader_en` for more details.
1555

F
fengjiayi 已提交
1556
    Args:
1557 1558 1559
        dirname(str): One of the following:
          - The given directory path.
          - Set to None when reading the model from memory.
F
fengjiayi 已提交
1560
        executor(Executor): The executor to run for loading inference model.
1561
                            See :ref:`api_guide_executor_en` for more details about it.
1562 1563 1564 1565 1566 1567 1568
        model_filename(str, optional): One of the following:
          - The name of file to load the inference program.
          - If it is None, the default filename ``__model__`` will be used.
          - When ``dirname`` is ``None``, it must be set to a string containing model.
          Default: ``None``.
        params_filename(str, optional): It is only used for the case that all
            parameters were saved in a single binary file. One of the following:
1569
          - The name of file to load all parameters.
1570 1571 1572
          - When ``dirname`` is ``None``, it must be set to a string containing all the parameters.
          - If parameters were saved in separate files, set it as ``None``.
            Default: ``None``.
1573 1574 1575 1576

        pserver_endpoints(list, optional): It is only needed by the distributed inference.
                                    If using a distributed look up table during the training,
                                    this table is also needed by the inference process. Its value is
1577
                                    a list of pserver endpoints.
F
fengjiayi 已提交
1578 1579

    Returns:
1580
        list: The return of this API is a list with three elements:
1581
        (program, feed_target_names, fetch_targets). The `program` is a
1582 1583 1584 1585 1586
        ``Program`` (refer to :ref:`api_guide_Program_en`), which is used for inference.
        The `feed_target_names` is a list of ``str``, which contains names of variables
        that need to feed data in the inference program. The `fetch_targets` is a list of
        ``Variable`` (refer to :ref:`api_guide_Program_en`). It contains variables from which
        we can get inference results.
F
fengjiayi 已提交
1587 1588 1589 1590 1591


    Examples:
        .. code-block:: python

1592
            import paddle
1593 1594
            import paddle.fluid as fluid
            import numpy as np
1595

1596
            paddle.enable_static()
1597
            # Build the model
1598 1599 1600
            main_prog = fluid.Program()
            startup_prog = fluid.Program()
            with fluid.program_guard(main_prog, startup_prog):
G
GGBond8488 已提交
1601
                data = paddle.static.data(name="img", shape=[-1, 64, 784])
1602 1603
                w = paddle.create_parameter(shape=[784, 200], dtype='float32')
                b = paddle.create_parameter(shape=[200], dtype='float32')
1604
                hidden_w = paddle.matmul(x=data, y=w)
H
HongyuJia 已提交
1605
                hidden_b = paddle.add(hidden_w, b)
1606 1607 1608
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            exe.run(startup_prog)
1609 1610

            # Save the inference model
F
fengjiayi 已提交
1611
            path = "./infer_model"
1612 1613
            fluid.io.save_inference_model(dirname=path, feeded_var_names=['img'],
                         target_vars=[hidden_b], executor=exe, main_program=main_prog)
1614 1615 1616

            # Demo one. Not need to set the distributed look up table, because the
            # training doesn't use a distributed look up table.
1617 1618
            [inference_program, feed_target_names, fetch_targets] = (
                fluid.io.load_inference_model(dirname=path, executor=exe))
1619
            tensor_img = np.array(np.random.random((1, 64, 784)), dtype=np.float32)
F
fengjiayi 已提交
1620 1621 1622 1623
            results = exe.run(inference_program,
                          feed={feed_target_names[0]: tensor_img},
                          fetch_list=fetch_targets)

1624 1625 1626
            # Demo two. If the training uses a distributed look up table, the pserver
            # endpoints list should be supported when loading the inference model.
            # The below is just an example.
1627
            endpoints = ["127.0.0.1:2023","127.0.0.1:2024"]
1628
            [dist_inference_program, dist_feed_target_names, dist_fetch_targets] = (
1629 1630
                fluid.io.load_inference_model(dirname=path,
                                              executor=exe,
1631
                                              pserver_endpoints=endpoints))
1632

1633
            # In this example, the inference program was saved in the file
1634
            # "./infer_model/__model__" and parameters were saved in
1635 1636 1637 1638
            # separate files under the directory "./infer_model".
            # By the inference program, feed_target_names and
            # fetch_targets, we can use an executor to run the inference
            # program for getting the inference result.
1639
    """
1640 1641 1642 1643
    load_from_memory = False
    if dirname is not None:
        load_dirname = os.path.normpath(dirname)
        if not os.path.isdir(load_dirname):
1644
            raise ValueError("There is no directory named '%s'" % dirname)
1645

1646 1647
        if model_filename is None:
            model_filename = '__model__'
1648

1649 1650 1651
        model_filename = os.path.join(
            load_dirname, os.path.basename(model_filename)
        )
1652

1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666
        if params_filename is not None:
            params_filename = os.path.basename(params_filename)

        with open(model_filename, "rb") as f:
            program_desc_str = f.read()
    else:
        load_from_memory = True
        if params_filename is None:
            raise ValueError(
                "The path of params cannot be None when the directory path is None."
            )
        load_dirname = dirname
        program_desc_str = model_filename
        params_filename = params_filename
1667

1668
    program = Program.parse_from_string(program_desc_str)
X
Xin Pan 已提交
1669
    if not core._is_program_version_supported(program._version()):
1670 1671 1672
        raise ValueError(
            "Unsupported program version: %d\n" % program._version()
        )
X
version  
Xin Pan 已提交
1673
    # Binary data also need versioning.
L
lujun 已提交
1674
    load_persistables(executor, load_dirname, program, params_filename)
1675

T
tangwei12 已提交
1676
    if pserver_endpoints:
T
tangwei12 已提交
1677
        program = _endpoints_replacement(program, pserver_endpoints)
T
tangwei12 已提交
1678

1679 1680
    feed_target_names = program.desc.get_feed_target_names()
    fetch_target_names = program.desc.get_fetch_target_names()
1681 1682 1683 1684 1685
    fetch_targets = [
        program.global_block().var(name) for name in fetch_target_names
    ]

    return [program, feed_target_names, fetch_targets]
X
xuwei06 已提交
1686 1687


T
tangwei12 已提交
1688 1689 1690
def _endpoints_replacement(program, endpoints):
    ENDPOINT_MAP = "epmap"
    for op in program.global_block().ops:
T
tangwei12 已提交
1691 1692
        if op.has_attr(ENDPOINT_MAP):
            op.set_attr(ENDPOINT_MAP, endpoints)
T
fix  
tangwei12 已提交
1693
    program._sync_with_cpp()
T
tangwei12 已提交
1694
    return program
T
tangwei12 已提交
1695 1696


X
xuwei06 已提交
1697 1698
def get_parameter_value(para, executor):
    """
F
fengjiayi 已提交
1699 1700 1701 1702 1703 1704 1705 1706 1707 1708 1709
    Get the LoDTensor value of the given parameter.

    Args:
        para(Parameter): The parameter to get value from.
        executor(Executor): The executor to run for retrieving the value.

    Returns:
        numpy.array: The given parameter's values.

    Raises:
        AssertionError: If the `para` is not an instance of Parameter.
X
xuwei06 已提交
1710

F
fengjiayi 已提交
1711 1712
    Examples:
        .. code-block:: python
X
xuwei06 已提交
1713

1714
            import paddle
1715
            import paddle.fluid as fluid
1716 1717

            paddle.enable_static()
F
fengjiayi 已提交
1718 1719 1720
            exe = fluid.Executor(fluid.CPUPlace())
            param = fluid.default_main_program().global_block().var('fc.w')
            p = fluid.io.get_parameter_value(param, exe)
1721

X
xuwei06 已提交
1722
    """
1723
    assert is_parameter(para), "The input variable is not parameter."
X
xuwei06 已提交
1724

X
xuwei06 已提交
1725 1726 1727 1728 1729 1730 1731 1732
    get_program = Program()
    block = get_program.global_block()
    new_var = _clone_var_in_block_(block, para)
    return executor.run(get_program, feed={}, fetch_list=[new_var])[0]


def get_parameter_value_by_name(name, executor, program=None):
    """
F
fengjiayi 已提交
1733
    Get the LoDTensor value of a certain parameter by its name.
X
xuwei06 已提交
1734

F
fengjiayi 已提交
1735 1736 1737 1738 1739 1740 1741
    Args:
        name(str): The parameter's name.
        executor(Executor): The executor to run for retrieving the value.
        program(Program | None): The program where to find the parameter.
                               If it's set to be None, the function will
                               try to find the parameter in the default
                               main program.
X
xuwei06 已提交
1742

F
fengjiayi 已提交
1743 1744
    Returns:
        numpy.array: The parameter's values.
1745

F
fengjiayi 已提交
1746 1747 1748
    Examples:
        .. code-block:: python

1749
            import paddle
1750
            import paddle.fluid as fluid
1751 1752

            paddle.enable_static()
F
fengjiayi 已提交
1753 1754
            exe = fluid.Executor(fluid.CPUPlace())
            p = fluid.io.get_parameter_value('fc.w', exe)
X
xuwei06 已提交
1755 1756
    """
    if program is None:
Y
Yu Yang 已提交
1757
        program = default_main_program()
X
xuwei06 已提交
1758 1759
    var = program.global_block().var(name)
    return get_parameter_value(var, executor)
1760 1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772 1773 1774


def _save_persistable_nodes(executor, dirname, graph):
    """
    Save persistable nodes to the given directory by the executor.

    Args:
        executor(Executor): The executor to run for saving node values.
        dirname(str): The directory path.
        graph(IrGraph): All the required persistable nodes in the graph will be saved.
    """
    persistable_node_names = set()
    persistable_nodes = []
    all_persistable_nodes = graph.all_persistable_nodes()
    for node in all_persistable_nodes:
1775
        name = node.name()
1776 1777 1778 1779 1780 1781 1782
        if name not in persistable_node_names:
            persistable_node_names.add(name)
            persistable_nodes.append(node)
    program = Program()
    var_list = []
    for node in persistable_nodes:
        var_desc = node.var()
1783 1784 1785 1786
        if (
            var_desc.type() == core.VarDesc.VarType.RAW
            or var_desc.type() == core.VarDesc.VarType.READER
        ):
1787 1788 1789 1790 1791 1792 1793
            continue
        var = program.global_block().create_var(
            name=var_desc.name(),
            shape=var_desc.shape(),
            dtype=var_desc.dtype(),
            type=var_desc.type(),
            lod_level=var_desc.lod_level(),
1794 1795
            persistable=var_desc.persistable(),
        )
1796 1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811 1812
        var_list.append(var)
    save_vars(executor=executor, dirname=dirname, vars=var_list)


def _load_persistable_nodes(executor, dirname, graph):
    """
    Load persistable node values from the given directory by the executor.

    Args:
        executor(Executor): The executor to run for loading node values.
        dirname(str): The directory path.
        graph(IrGraph): All the required persistable nodes in the graph will be loaded.
    """
    persistable_node_names = set()
    persistable_nodes = []
    all_persistable_nodes = graph.all_persistable_nodes()
    for node in all_persistable_nodes:
1813
        name = node.name()
1814 1815 1816 1817 1818 1819 1820 1821 1822 1823 1824
        if name not in persistable_node_names:
            persistable_node_names.add(name)
            persistable_nodes.append(node)
    program = Program()
    var_list = []

    def _exist(var):
        return os.path.exists(os.path.join(dirname, var.name))

    for node in persistable_nodes:
        var_desc = node.var()
1825 1826 1827 1828
        if (
            var_desc.type() == core.VarDesc.VarType.RAW
            or var_desc.type() == core.VarDesc.VarType.READER
        ):
1829 1830 1831 1832 1833 1834 1835
            continue
        var = program.global_block().create_var(
            name=var_desc.name(),
            shape=var_desc.shape(),
            dtype=var_desc.dtype(),
            type=var_desc.type(),
            lod_level=var_desc.lod_level(),
1836 1837
            persistable=var_desc.persistable(),
        )
1838 1839 1840 1841 1842
        if _exist(var):
            var_list.append(var)
        else:
            _logger.warn("Cannot find the var %s!!!" % (node.name()))
    load_vars(executor=executor, dirname=dirname, vars=var_list)
H
hong 已提交
1843 1844


W
WeiXin 已提交
1845
def _unpack_saved_dict(saved_obj, protocol):
1846 1847
    temp_saved_obj = {}
    unpack_infor = {}
W
WeiXin 已提交
1848 1849 1850 1851 1852 1853
    # When pickle protocol=2 or protocol=3 the serialized object cannot be larger than 4G.
    if 1 < protocol < 4:
        if isinstance(saved_obj, dict):
            for key, value in saved_obj.items():
                if isinstance(value, np.ndarray):
                    MAX_NUMBER_OF_ELEMENT = int(
1854 1855
                        (2**30 - 1) / value.dtype.itemsize
                    )
W
WeiXin 已提交
1856 1857 1858 1859 1860 1861 1862
                    num_element = np.prod(value.shape)
                    if num_element > MAX_NUMBER_OF_ELEMENT:
                        unpack_infor[key] = {}
                        unpack_infor[key]["OriginShape"] = value.shape
                        unpack_infor[key]["slices"] = []
                        value = value.flatten()
                        for i in range(
1863 1864 1865 1866 1867 1868
                            int(
                                math.ceil(
                                    num_element * 1.0 / MAX_NUMBER_OF_ELEMENT
                                )
                            )
                        ):
W
WeiXin 已提交
1869 1870 1871
                            part_name = key + "@@." + str(i)
                            unpack_infor[key]["slices"].append(part_name)
                            temp_saved_obj[part_name] = value[
1872 1873 1874 1875
                                i
                                * MAX_NUMBER_OF_ELEMENT : MAX_NUMBER_OF_ELEMENT
                                * (i + 1)
                            ]
1876 1877 1878 1879 1880 1881 1882 1883 1884 1885 1886 1887

    if unpack_infor:
        for key, value in unpack_infor.items():
            if key in saved_obj:
                saved_obj.pop(key)
                for part in value['slices']:
                    saved_obj[part] = temp_saved_obj[part]
        saved_obj['UnpackBigParamInfor@@'] = unpack_infor
    return saved_obj


def _pack_loaded_dict(load_obj):
W
WeiXin 已提交
1888 1889 1890 1891 1892 1893
    if isinstance(load_obj, dict):
        unpack_info = 'UnpackBigParamInfor@@'
        if unpack_info in load_obj:
            removes = []
            for key, value in load_obj[unpack_info].items():
                slices = [load_obj[part] for part in value["slices"]]
1894
                load_obj[key] = np.concatenate(slices).reshape(
1895 1896
                    value["OriginShape"]
                )
W
WeiXin 已提交
1897 1898 1899 1900 1901
                removes += value["slices"]
            for key in removes:
                load_obj.pop(key)
            load_obj.pop(unpack_info)

1902 1903 1904
    return load_obj


1905
@static_only
1906 1907
def _legacy_save(param_dict, model_path, protocol=2):
    def get_tensor(var):
J
Jiabin Yang 已提交
1908
        if isinstance(var, (core.VarBase, core.eager.Tensor)):
1909 1910 1911 1912 1913 1914 1915 1916
            return var.numpy()
        elif isinstance(var, core.LoDTensor):
            return np.array(var)
        return var

    param_dict = {name: get_tensor(param_dict[name]) for name in param_dict}

    # When value of dict is lager than 4GB ,there is a Bug on 'MAC python3'
1917 1918 1919 1920 1921
    if (
        _is_file_path(model_path)
        and sys.platform == 'darwin'
        and sys.version_info.major == 3
    ):
1922 1923 1924 1925
        pickle_bytes = pickle.dumps(param_dict, protocol=protocol)
        with open(model_path, 'wb') as f:
            max_bytes = 2**30
            for i in range(0, len(pickle_bytes), max_bytes):
1926
                f.write(pickle_bytes[i : i + max_bytes])
1927
    else:
1928
        with _open_file_buffer(model_path, 'wb') as f:
1929 1930 1931 1932
            pickle.dump(param_dict, f, protocol=protocol)


@static_only
1933
def save(program, model_path, protocol=4, **configs):
H
hong 已提交
1934
    """
1935

1936
    This function save parameters, optimizer information and network description to model_path.
H
hong 已提交
1937

1938 1939
    The parameters contains all the trainable Tensor, will save to a file with suffix ".pdparams".
    The optimizer information contains all the Tensor used by optimizer. For Adam optimizer, contains beta1, beta2, momentum etc. All the information will save to a file with suffix ".pdopt". (If the optimizer have no Tensor need to save (like SGD), the fill will not generated).
H
hong 已提交
1940
    The network description is the description of the program. It's only used for deployment. The description  will save to a file with a suffix ".pdmodel".
1941

H
hong 已提交
1942 1943 1944
    Args:
        program(Program) : The program to saved.
        model_path(str): the file prefix to save the program. The format is "dirname/file_prefix". If file_prefix is empty str. A exception will be raised
1945
        protocol(int, optional): The protocol version of pickle module must be greater than 1 and less than 5.
1946
                                 Default: 4
1947
        configs(dict, optional) : optional keyword arguments.
H
hong 已提交
1948 1949 1950 1951 1952 1953 1954

    Returns:
        None

    Examples:
        .. code-block:: python

1955
            import paddle
1956
            import paddle.static as static
H
hong 已提交
1957

1958
            paddle.enable_static()
H
hong 已提交
1959

1960 1961 1962 1963 1964 1965 1966 1967 1968 1969
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)

            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()

            static.save(prog, "./temp")
H
hong 已提交
1970 1971 1972
    """

    base_name = os.path.basename(model_path)
1973 1974 1975
    assert (
        base_name != ""
    ), "The input model_path MUST be format of dirname/filename [dirname\\filename in Windows system], but received model_path is empty string."
1976 1977 1978 1979 1980
    if 'pickle_protocol' in configs:
        protocol = configs['pickle_protocol']
        warnings.warn(
            "'pickle_protocol' is a deprecated argument. Please use 'protocol' instead."
        )
H
hong 已提交
1981

1982
    if not isinstance(protocol, int):
1983 1984 1985 1986 1987
        raise ValueError(
            "The 'protocol' MUST be `int`, but received {}".format(
                type(protocol)
            )
        )
W
WeiXin 已提交
1988

1989
    if protocol < 2 or protocol > 4:
1990
        raise ValueError(
1991 1992
            "Expected 1<'protocol'<5, but received protocol={}".format(protocol)
        )
W
WeiXin 已提交
1993

1994 1995 1996 1997
    dir_name = os.path.dirname(model_path)
    if dir_name and not os.path.exists(dir_name):
        os.makedirs(dir_name)

Y
Yang Zhang 已提交
1998 1999 2000 2001
    def get_tensor(var):
        t = global_scope().find_var(var.name).get_tensor()
        return np.array(t)

H
hong 已提交
2002
    parameter_list = list(filter(is_parameter, program.list_vars()))
Y
Yang Zhang 已提交
2003
    param_dict = {p.name: get_tensor(p) for p in parameter_list}
W
WeiXin 已提交
2004

2005
    param_dict = _unpack_saved_dict(param_dict, protocol)
2006

2007 2008 2009
    # When value of dict is lager than 4GB ,there is a Bug on 'MAC python3'
    if sys.platform == 'darwin' and sys.version_info.major == 3:
        pickle_bytes = pickle.dumps(param_dict, protocol=protocol)
2010 2011 2012
        with open(model_path + ".pdparams", 'wb') as f:
            max_bytes = 2**30
            for i in range(0, len(pickle_bytes), max_bytes):
2013
                f.write(pickle_bytes[i : i + max_bytes])
2014 2015
    else:
        with open(model_path + ".pdparams", 'wb') as f:
2016
            pickle.dump(param_dict, f, protocol=protocol)
H
hong 已提交
2017 2018

    optimizer_var_list = list(
2019 2020
        filter(is_belong_to_optimizer, program.list_vars())
    )
H
hong 已提交
2021

Y
Yang Zhang 已提交
2022 2023
    opt_dict = {p.name: get_tensor(p) for p in optimizer_var_list}
    with open(model_path + ".pdopt", 'wb') as f:
2024
        pickle.dump(opt_dict, f, protocol=protocol)
H
hong 已提交
2025 2026 2027 2028

    main_program = program.clone()
    program.desc.flush()
    main_program.desc._set_version()
2029
    paddle.fluid.core.save_op_version_info(program.desc)
H
hong 已提交
2030 2031 2032 2033 2034

    with open(model_path + ".pdmodel", "wb") as f:
        f.write(program.desc.serialize_to_string())


2035 2036 2037 2038 2039 2040
def _pickle_loads_mac(path, f):
    pickle_bytes = bytearray(0)
    file_size = os.path.getsize(path)
    max_bytes = 2**30
    for _ in range(0, file_size, max_bytes):
        pickle_bytes += f.read(max_bytes)
T
tianshuo78520a 已提交
2041
    load_result = pickle.loads(pickle_bytes, encoding='latin1')
2042 2043 2044
    return load_result


2045
@static_only
H
hong 已提交
2046
def load(program, model_path, executor=None, var_list=None):
H
hong 已提交
2047
    """
2048 2049
    :api_attr: Static Graph

H
hong 已提交
2050
    This function get parameters and optimizer information from program, and then get corresponding value from file.
2051
    An exception will throw if shape or dtype of the parameters is not match.
H
hong 已提交
2052

2053 2054
    This function can also load model file saved with [ save_params, save_persistables, save_vars ].
    var_list can not be None  when load single model file
H
hong 已提交
2055 2056
    ( filename is not None When save_params, save_persistables or save_vars is called ).

2057
    Args:
2058 2059
        program(Program): The program will be loaded
        model_path(str): The file prefix store the program
2060
        executor(Executor, optional): The executor used for initialize the parameter
2061
                                      When startup program is not run.
2062
        var_list(list|tuple, optional): The Tensor list/tuple to load single model file saved with
2063
                                  [ save_params, save_persistables, save_vars ].
H
hong 已提交
2064
                                  Default: None
H
hong 已提交
2065 2066 2067

    Returns:
        None
2068

H
hong 已提交
2069 2070 2071
     Examples:
        .. code-block:: python

2072
            import paddle
2073
            import paddle.static as static
H
hong 已提交
2074

2075
            paddle.enable_static()
H
hong 已提交
2076

2077 2078 2079
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
H
hong 已提交
2080

2081 2082 2083 2084 2085 2086 2087
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()

            static.save(prog, "./temp")
            static.load(prog, "./temp")
H
hong 已提交
2088 2089
    """

2090 2091
    assert executor is None or isinstance(executor, Executor)

H
hong 已提交
2092 2093 2094 2095 2096 2097 2098 2099 2100 2101 2102 2103 2104
    model_prefix = model_path
    if model_prefix.endswith(".pdparams"):
        model_prefix = model_prefix[:-9]
    elif model_prefix.endswith(".pdopt"):
        model_prefix = model_prefix[:-6]
    elif model_prefix.endswith(".pdmodel"):
        model_prefix = model_prefix[:-8]

    parameter_file_name = model_prefix + ".pdparams"

    if not os.path.exists(parameter_file_name):
        # model file save by fluid.save not found, try to load model file saved with
        # [save_vars, save_params, save_persistables]
2105
        _logger.debug(
2106 2107 2108 2109
            "{} not found, try to load model file saved with [ save_params, save_persistables, save_vars ]".format(
                parameter_file_name
            )
        )
H
hong 已提交
2110 2111 2112 2113
        if executor is None:
            raise ValueError(
                "executor is required when loading model file saved with [ save_params, save_persistables, save_vars ]"
            )
2114 2115 2116 2117 2118 2119

        if var_list is not None:
            var_list_names = [var.name for var in var_list]
        else:
            var_list_names = None

H
hong 已提交
2120 2121 2122 2123 2124
        if os.path.isdir(model_path):
            binary_file_set = set()
            for root, dirs, files in os.walk(model_path, topdown=False):
                for f in files:
                    binary_file_set.add(
2125 2126
                        os.path.join(root, f).replace("\\", "/")
                    )
H
hong 已提交
2127 2128 2129 2130
            program_var_list = list(program.list_vars())
            loaded_var_list = []
            for var in program_var_list:
                var_path = os.path.join(model_path, var.name).replace("\\", "/")
2131 2132 2133
                load_condition = (
                    var_list_names is None or var.name in var_list_names
                )
2134
                if var_path in binary_file_set and load_condition:
H
hong 已提交
2135 2136 2137 2138
                    loaded_var_list.append(var)
                    binary_file_set.remove(var_path)
            if len(binary_file_set) > 0:
                unused_var_list = " ".join(list(binary_file_set))
2139 2140 2141 2142
                _logger.warning(
                    "variable file [ %s ] not used"
                    % (" ".join(list(binary_file_set)))
                )
H
hong 已提交
2143
            try:
2144 2145 2146
                load_vars(
                    executor=executor, dirname=model_path, vars=loaded_var_list
                )
H
hong 已提交
2147 2148 2149 2150 2151
            except RuntimeError as e:
                _logger.error(e)
                raise e
            except:
                raise RuntimeError(
2152
                    "Failed to load model file, please make sure model file is saved with the "
2153 2154
                    "following APIs: save_params, save_persistables, save_vars"
                )
H
hong 已提交
2155 2156 2157

            return
        elif os.path.isfile(model_path):
2158
            if var_list is None:
H
hong 已提交
2159 2160 2161 2162 2163 2164 2165 2166 2167 2168
                raise ValueError(
                    "var_list is required when loading model file saved with [ save_params, save_persistables, save_vars ]"
                )
            program_var_list = program.list_vars()
            program_var_name_set = set([var.name for var in program_var_list])

            # check all the variable inlcuded in program
            for var in var_list:
                if var.name not in program_var_name_set:
                    raise LookupError(
2169 2170
                        "loaded var [{}] is not in program variable list"
                    )
H
hong 已提交
2171 2172 2173

            dir_name, file_name = os.path.split(model_path)
            try:
2174 2175 2176 2177 2178 2179
                load_vars(
                    executor=executor,
                    dirname=dir_name,
                    vars=var_list,
                    filename=file_name,
                )
H
hong 已提交
2180 2181 2182 2183
            except RuntimeError as e:
                _logger.error(e)
                raise e
            except:
2184 2185 2186 2187 2188
                raise RuntimeError(
                    "Failed to load model file , please make sure model file is saved with the "
                    "the following APIs: [ save_params, save_persistables, save_vars ]. "
                    "When these API called, filename CANNOT be None"
                )
H
hong 已提交
2189 2190

            return
Y
Yang Zhang 已提交
2191 2192 2193 2194 2195 2196 2197 2198

    def set_var(var, ndarray):
        t = global_scope().find_var(var.name).get_tensor()
        p = t._place()
        if p.is_cpu_place():
            place = paddle.fluid.CPUPlace()
        elif p.is_cuda_pinned_place():
            place = paddle.fluid.CUDAPinnedPlace()
2199 2200 2201 2202
        elif p.is_xpu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.XPUPlace(p.xpu_device_id())
2203 2204 2205 2206
        elif p.is_npu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.NPUPlace(p.npu_device_id())
2207 2208 2209 2210
        elif p.is_mlu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.MLUPlace(p.mlu_device_id())
Y
Yang Zhang 已提交
2211 2212 2213 2214 2215 2216
        else:
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.CUDAPlace(p.gpu_device_id())

        t.set(ndarray, place)
H
hong 已提交
2217 2218

    parameter_list = list(filter(is_parameter, program.list_vars()))
2219 2220

    if executor:
2221 2222 2223
        paddle.fluid.core._create_loaded_parameter(
            parameter_list, global_scope(), executor._default_executor
        )
Y
Yang Zhang 已提交
2224
    with open(parameter_file_name, 'rb') as f:
2225 2226 2227 2228 2229

        # When value of dict is lager than 4GB ,there is a Bug on 'MAC python3'
        if sys.platform == 'darwin' and sys.version_info.major == 3:
            load_dict = _pickle_loads_mac(parameter_file_name, f)
        else:
T
tianshuo78520a 已提交
2230
            load_dict = pickle.load(f, encoding='latin1')
2231
        load_dict = _pack_loaded_dict(load_dict)
Y
Yang Zhang 已提交
2232
    for v in parameter_list:
2233 2234 2235 2236 2237
        assert (
            v.name in load_dict
        ), "Can not find [{}] in model file [{}]".format(
            v.name, parameter_file_name
        )
Y
Yang Zhang 已提交
2238
        set_var(v, load_dict[v.name])
H
hong 已提交
2239 2240

    optimizer_var_list = list(
2241 2242
        filter(is_belong_to_optimizer, program.list_vars())
    )
H
hong 已提交
2243 2244

    if len(optimizer_var_list) > 0:
H
hong 已提交
2245
        opt_file_name = model_prefix + ".pdopt"
2246 2247 2248
        assert os.path.exists(
            opt_file_name
        ), "Optimizer file [{}] not exits".format(opt_file_name)
2249 2250 2251

        if executor:
            paddle.fluid.core._create_loaded_parameter(
2252 2253
                optimizer_var_list, global_scope(), executor._default_executor
            )
Y
Yang Zhang 已提交
2254 2255

        with open(opt_file_name, 'rb') as f:
T
tianshuo78520a 已提交
2256
            load_dict = pickle.load(f, encoding='latin1')
Y
Yang Zhang 已提交
2257
        for v in optimizer_var_list:
2258 2259 2260 2261 2262
            assert (
                v.name in load_dict
            ), "Can not find [{}] in model file [{}]".format(
                v.name, opt_file_name
            )
Y
Yang Zhang 已提交
2263
            set_var(v, load_dict[v.name])
2264 2265


H
hong 已提交
2266
def load_program_state(model_path, var_list=None):
2267
    """
2268

2269
    Load program state from local file
2270

2271 2272
    Args:
        model_path(str): The file prefix store the program
2273
        var_list(list|tuple, optional): The Tensor list/tuple to load saved with
2274
                                  [ save_params, save_persistables, save_vars ].
H
hong 已提交
2275
                                  Default: None.
2276
                                  The var_list is only used to get name,
H
hong 已提交
2277
                                  will not be modified.
2278 2279 2280 2281
    Returns:
        state_dict(dict): the dict store Parameter and optimizer information

    Examples:
2282

2283 2284
        .. code-block:: python

2285
            import paddle
2286
            import paddle.static as static
2287 2288

            paddle.enable_static()
2289

2290 2291 2292
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
2293

2294 2295 2296 2297
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()
2298

2299 2300
            static.save(prog, "./temp")
            program_state = static.load_program_state("./temp")
2301
    """
H
hong 已提交
2302 2303 2304 2305 2306 2307 2308 2309 2310 2311 2312 2313
    model_prefix = model_path
    if model_prefix.endswith(".pdparams"):
        model_prefix = model_prefix[:-9]
    elif model_prefix.endswith(".pdopt"):
        model_prefix = model_prefix[:-6]
    elif model_prefix.endswith(".pdmodel"):
        model_prefix = model_prefix[:-8]

    parameter_file_name = model_prefix + ".pdparams"
    if not os.path.exists(parameter_file_name):
        # model file saved with fluid.save is not found, try to load model file saved with
        # [save_vars, save_params, save_persistables]
2314
        _logger.debug(
2315 2316 2317 2318
            "{} not found, try to load model file saved with [ save_params, save_persistables, save_vars ]".format(
                parameter_file_name
            )
        )
H
hong 已提交
2319 2320 2321 2322

        var_name_list = []
        if var_list is None and os.path.isfile(model_path):
            raise ValueError(
2323 2324
                "var_list can not be None when model_path is a file type"
            )
H
hong 已提交
2325 2326 2327 2328 2329 2330 2331 2332 2333 2334 2335 2336 2337 2338 2339 2340 2341 2342 2343 2344

        for root, dirs, files in os.walk(model_path, topdown=False):
            for f in files:
                file_path = os.path.join(root, f)
                var_temp_name = os.path.relpath(file_path, model_path)
                var_temp_name = var_temp_name.replace("\\", "/")
                var_name_list.append(var_temp_name)

        with _load_program_scope():
            load_prog = Program()
            load_block = load_prog.global_block()

            def clone_var_to_block(block, var):
                if not isinstance(var, Variable):
                    raise TypeError("value in var_list must be variable")
                return block.create_var(
                    name=var.name,
                    shape=var.shape,
                    dtype=var.dtype,
                    type=var.type,
2345 2346 2347 2348 2349 2350 2351 2352 2353
                    lod_level=var.lod_level
                    if var.desc.type() == core.VarDesc.VarType.LOD_TENSOR
                    else None,
                    persistable=True,
                )

            def _load_vars_with_try_catch(
                exe, dirname, vars, filename, raise_error=True
            ):
2354
                try:
2355 2356 2357 2358 2359 2360
                    load_vars(
                        executor=exe,
                        dirname=dirname,
                        vars=vars,
                        filename=filename,
                    )
2361 2362
                    return True
                except:
2363 2364 2365 2366 2367 2368 2369 2370 2371 2372
                    error_str = (
                        "Failed to load model/variables `%s`, please make sure "
                        "model/variables file is saved with the following APIs: "
                        "save_params, save_persistables, save_vars."
                    )
                    filenames = (
                        [var.name for var in vars]
                        if filename is None
                        else filename
                    )
2373 2374 2375 2376 2377 2378 2379 2380 2381
                    if raise_error:
                        raise RuntimeError(error_str % filenames)
                    else:
                        warnings.warn(error_str % filenames, RuntimeWarning)
                return False

            place = paddle.fluid.CPUPlace()
            exe = paddle.fluid.Executor(place)

H
hong 已提交
2382 2383
            loaded_var_list = []

2384 2385 2386
            if os.path.isfile(model_path):
                # when model_path is file, var_list cannot be None
                dir_name, file_name = os.path.split(model_path)
H
hong 已提交
2387 2388
                for var in var_list:
                    loaded_var_list.append(clone_var_to_block(load_block, var))
2389 2390 2391
                _load_vars_with_try_catch(
                    exe, dir_name, loaded_var_list, file_name
                )
H
hong 已提交
2392
            else:
2393 2394 2395 2396
                # var_list can be None or not None
                if var_list is not None:
                    for var in var_list:
                        loaded_var_list.append(
2397 2398 2399 2400 2401
                            clone_var_to_block(load_block, var)
                        )
                    _load_vars_with_try_catch(
                        exe, model_path, loaded_var_list, None
                    )
H
hong 已提交
2402
                else:
2403
                    for var_name in var_name_list:
2404 2405 2406 2407
                        # NOTE(chenweihang): If identify which files the user wants
                        # to load from the disk, we load these variables one by one.
                        # If a file does not exist, we only warn the user that the
                        # file may be an irrelevant file, but does not throw an error
2408
                        # to ensure that other legal variables can be loaded.
2409 2410 2411 2412 2413 2414
                        temp_var = load_block.create_var(
                            name=var_name, persistable=True
                        )
                        if _load_vars_with_try_catch(
                            exe, model_path, [temp_var], None, False
                        ):
2415 2416
                            loaded_var_list.append(temp_var)

H
hong 已提交
2417 2418
            res_dict = {}
            for var in loaded_var_list:
2419
                res_dict[var.name] = np.asarray(
2420 2421
                    paddle.fluid.global_scope().find_var(var.name).get_tensor()
                )
H
hong 已提交
2422 2423 2424

            return res_dict

2425 2426 2427
    assert os.path.exists(
        parameter_file_name
    ), "Parameter file [{}] not exits".format(parameter_file_name)
2428 2429

    with open(parameter_file_name, 'rb') as f:
2430 2431 2432 2433
        # When value of dict is lager than 4GB ,there is a Bug on 'MAC python3'
        if sys.platform == 'darwin' and sys.version_info.major == 3:
            para_dict = _pickle_loads_mac(parameter_file_name, f)
        else:
T
tianshuo78520a 已提交
2434
            para_dict = pickle.load(f, encoding='latin1')
2435
    para_dict = _pack_loaded_dict(para_dict)
2436

H
hong 已提交
2437
    opt_file_name = model_prefix + ".pdopt"
2438 2439
    if os.path.exists(opt_file_name):
        with open(opt_file_name, 'rb') as f:
T
tianshuo78520a 已提交
2440
            opti_dict = pickle.load(f, encoding='latin1')
2441 2442 2443 2444 2445 2446

        para_dict.update(opti_dict)

    return para_dict


2447
@static_only
2448 2449 2450 2451
def set_program_state(program, state_dict):
    """
    Set program parameter from state_dict

2452
    An exception will throw if shape or dtype of the parameters is not match.
2453 2454 2455 2456 2457 2458

    NOTICE: This function MUST called after run start_up_program

    Args:
        program(Program): The program to be set
        state_dict(dict): the dict store Parameter and optimizer information
2459
    Returns:
2460
        None
2461

2462 2463
    Examples:
        .. code-block:: python
2464

2465
            import paddle
2466
            import paddle.static as static
2467 2468

            paddle.enable_static()
2469

2470 2471 2472
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
2473

2474 2475 2476 2477
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()
2478

2479 2480
            static.save(prog, "./temp")
            program_state = static.load_program_state("./temp")
H
hong 已提交
2481

2482
            static.set_program_state(prog, program_state)
2483
    """
2484
    state_dict = _pack_loaded_dict(state_dict)
2485 2486 2487 2488 2489
    parameter_list = list(filter(is_persistable, program.list_vars()))

    used_para_list = {}
    for para in parameter_list:
        var_temp = paddle.fluid.global_scope().find_var(para.name)
2490
        assert (
2491
            var_temp is not None
2492 2493 2494
        ), "Variable [ {} ] Not found, Please make sure run startup program".format(
            para.name
        )
2495 2496 2497 2498
        if para.name in state_dict:
            # set value from state dict
            orig_para_np = np.array(var_temp.get_tensor())
            new_para_np = state_dict[para.name]
2499 2500 2501 2502 2503 2504 2505 2506 2507 2508 2509 2510
            assert orig_para_np.shape == new_para_np.shape, (
                "Parameter's shape does not match, the Program requires a parameter with the shape of ({}), "
                "while the loaded parameter (namely [ {} ]) has a shape of  ({}).".format(
                    orig_para_np.shape, para.name, new_para_np.shape
                )
            )
            assert orig_para_np.dtype == new_para_np.dtype, (
                "Parameter's data type does not match, the Program requires a parameter with a dtype of ({}), "
                "while the loaded parameter (namely [ {} ]) has a dtype of  ({}).".format(
                    orig_para_np.dtype, para.name, new_para_np.dtype
                )
            )
2511 2512 2513 2514

            ten = var_temp.get_tensor()
            ten_place = ten._place()

2515
            # assert ten_place.is_gpu_place() or ten_place.is_cpu_place(), \
Q
QingshuChen 已提交
2516
            #    "Place not support, only support CPUPlace and GPUPlace, now is {}".format(str(ten_place))
2517 2518 2519 2520 2521 2522 2523
            py_place = paddle.fluid.CPUPlace()
            if ten_place.is_cuda_pinned_place():
                place = paddle.fluid.CUDAPinnedPlace()
            elif ten_place.is_gpu_place():
                p = paddle.fluid.core.Place()
                p.set_place(ten_place)
                py_place = paddle.fluid.CUDAPlace(p.gpu_device_id())
Q
QingshuChen 已提交
2524 2525 2526 2527
            elif ten_place.is_xpu_place():
                p = paddle.fluid.core.Place()
                p.set_place(ten_place)
                py_place = paddle.fluid.XPUPlace(p.xpu_device_id())
2528 2529 2530 2531
            elif ten_place.is_npu_place():
                p = paddle.fluid.core.Place()
                p.set_place(ten_place)
                py_place = paddle.fluid.NPUPlace(p.npu_device_id())
2532 2533 2534 2535
            elif ten_place.is_mlu_place():
                p = paddle.fluid.core.Place()
                p.set_place(ten_place)
                py_place = paddle.fluid.MLUPlace(p.mlu_device_id())
2536 2537 2538 2539 2540 2541 2542 2543 2544 2545 2546

            ten.set(new_para_np, py_place)

            used_para_list[para.name] = 1

    unused_para_list = []
    for k, v in state_dict.items():
        if k not in used_para_list:
            unused_para_list.append(k)
    if len(unused_para_list) > 0:
        warnings.warn(
2547 2548 2549 2550
            "This list is not set, Because of Paramerter not found in program. There are: {}".format(
                " ".join(unused_para_list)
            )
        )