io.py 93.7 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 217 218 219 220 221 222 223
            data = fluid.data(name="img", shape=[64, 784])
            w = fluid.layers.create_parameter(shape=[784, 200], dtype='float32', name='fc_w')
            b = fluid.layers.create_parameter(shape=[200], dtype='float32', name='fc_b')
            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 243 244 245 246 247 248 249
            data = fluid.data(name="img", shape=[64, 784])
            w = fluid.layers.create_parameter(shape=[784, 200], dtype='float32', name='fc_w')
            b = fluid.layers.create_parameter(shape=[200], dtype='float32', name='fc_b')
            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 358 359 360 361 362 363 364 365
            main_prog = fluid.Program()
            startup_prog = fluid.Program()
            with fluid.program_guard(main_prog, startup_prog):
                data = fluid.layers.data(name="img", shape=[64, 784], append_batch_size=False)
                w = fluid.layers.create_parameter(shape=[784, 200], dtype='float32', name='fc_w')
                b = fluid.layers.create_parameter(shape=[200], dtype='float32', name='fc_b')
                hidden_w = fluid.layers.matmul(x=data, y=w)
                hidden_b = fluid.layers.elementwise_add(hidden_w, b)
            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 514
            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())
            predict = fluid.layers.fc(input=image, size=10, act='softmax')
515

G
guofei 已提交
516
            loss = fluid.layers.cross_entropy(input=predict, label=label)
517
            avg_loss = paddle.mean(loss)
518

F
fengjiayi 已提交
519
            exe = fluid.Executor(fluid.CPUPlace())
G
guofei 已提交
520 521
            exe.run(fluid.default_startup_program())
            fluid.io.save_params(executor=exe, dirname=params_path)
522 523
            # 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"
524
    """
525 526 527 528 529 530 531 532
    return save_vars(
        executor,
        dirname=dirname,
        main_program=main_program,
        vars=None,
        predicate=is_parameter,
        filename=filename,
    )
533 534


535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556
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

557
            import paddle
558
            import paddle.fluid as fluid
559 560

            paddle.enable_static()
561 562 563 564 565 566 567 568 569 570
            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 已提交
571
        receive params on pserver through rpc.
572 573 574 575 576 577 578 579 580 581
        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 已提交
582 583 584 585 586 587 588
            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)
589 590 591

            for idx, optimizer in enumerate(remote_params):
                block_id = optimizer.block_id
T
tangwei12 已提交
592
                slice = optimizer.slice
593 594 595
                endpoint = optimizer.endpoint

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

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

606 607 608 609 610 611 612 613 614 615 616 617
            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 已提交
618

619 620
        executor.run(prog)

621 622 623
    def __save_distributed_lookup_tables(
        executor, dirname, distributed_lookup_table, endpoints
    ):
624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639
        """
        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
640 641 642
        block.append_op(
            type='checkpoint_notify', inputs={}, outputs={}, attrs=attrs
        )
643 644 645 646 647 648
        executor.run(prog)

    def __exclude_vars(exclude_var_names=[]):
        def is_valid(var):
            if var.name in exclude_var_names:
                return False
649 650 651 652 653
            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
            ):
654 655 656 657 658 659
                return False
            return var.persistable

        return is_valid

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

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

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

    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(
684 685 686 687 688
        filter(__exclude_vars(exclude_var_names), main_program.list_vars())
    )
    save_vars(
        executor, main_program=main_program, dirname=dirname, vars=local_vars
    )
689 690 691 692 693 694

    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(
695 696 697 698 699
                executor,
                dirname,
                main_program._distributed_lookup_table,
                main_program._endpoints,
            )
700 701


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

G
guofei 已提交
711
    The :code:`dirname` is used to specify the folder where persistable variables
712
    are going to be saved. If you would like to save variables in separate
G
guofei 已提交
713 714
    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 已提交
715 716 717

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

721 722 723
        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
724
                                         be saved. You can refer to
G
guofei 已提交
725
                                         :ref:`api_guide_Program_en` for more details.
726
                                         If it is None, the default main program will
G
guofei 已提交
727 728 729 730 731
                                         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 已提交
732 733

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

    Examples:
        .. code-block:: python

740
            import paddle
H
Huihuang Zheng 已提交
741
            import paddle.fluid as fluid
742

743
            paddle.enable_static()
G
guofei 已提交
744 745 746 747 748
            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())
749

G
guofei 已提交
750 751
            predict = fluid.layers.fc(input=image, size=10, act='softmax')
            loss = fluid.layers.cross_entropy(input=predict, label=label)
752
            avg_loss = paddle.mean(loss)
F
fengjiayi 已提交
753
            exe = fluid.Executor(fluid.CPUPlace())
G
guofei 已提交
754 755
            exe.run(fluid.default_startup_program())
            fluid.io.save_persistables(executor=exe, dirname=dir_path, filename=file_name)
756
            # The persistables variables weights and bias in the fc layer of the network
G
guofei 已提交
757 758
            # are going to be saved in the same file named "persistables" in the path
            # "./my_paddle_model"
759
    """
760
    if main_program and main_program._is_distributed:
761 762 763
        return _save_distributed_persistables(
            executor, dirname=dirname, main_program=main_program
        )
764
    else:
765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782
        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,
):
783
    """
784 785
    :api_attr: Static Graph

786
    This API loads variables from files by executor.
F
fengjiayi 已提交
787

788
    There are two ways to specify the variables to be loaded: the first way, set
789 790
    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`.
791
    The first way has a higher priority.
F
fengjiayi 已提交
792

793
    The `dirname` is used to specify the folder where to load variables.
794
    If variables were saved in separate files in the folder `dirname`,
795
    set `filename` None. If all variables were saved in a single file,
F
fengjiayi 已提交
796
    use `filename` to specify it.
797

F
fengjiayi 已提交
798 799
    Args:
        executor(Executor): The executor to run for loading variables.
800 801
        dirname(str): The folder where to load the variables.
        main_program(Program, optional): The program whose variables will be loaded.
802
                                    If it is None, the default main program will
F
fengjiayi 已提交
803 804
                                    be used automatically.
                                    Default: None
805
        vars(list[Variable], optional): The list that contains all variables to be loaded.
F
fengjiayi 已提交
806
                                   Default: None
807
        predicate(function, optional): The function selects variables that make
808 809 810 811 812
                                        `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 已提交
813 814 815 816 817 818 819

    Returns:
        None

    Examples:
        .. code-block:: python

820
            import paddle
821
            import paddle.fluid as fluid
822

823
            paddle.enable_static()
824 825 826 827 828 829 830 831 832 833 834
            main_prog = fluid.Program()
            startup_prog = fluid.Program()
            with fluid.program_guard(main_prog, startup_prog):
                data = fluid.layers.data(name="img", shape=[64, 784], append_batch_size=False)
                w = fluid.layers.create_parameter(shape=[784, 200], dtype='float32', name='fc_w')
                b = fluid.layers.create_parameter(shape=[200], dtype='float32', name='fc_b')
                hidden_w = fluid.layers.matmul(x=data, y=w)
                hidden_b = fluid.layers.elementwise_add(hidden_w, b)
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            exe.run(startup_prog)
F
fengjiayi 已提交
835

836 837 838 839 840 841 842 843 844 845 846
            # 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
847
            param_path = "./my_paddle_model"
F
fengjiayi 已提交
848 849 850
            def name_has_fc(var):
                res = "fc" in var.name
                return res
851 852 853
            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 已提交
854
                               vars=None, predicate=name_has_fc)
855 856
            # 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 已提交
857

858
    """
859 860 861 862 863
    vars_from_memory = False
    if dirname is not None:
        dirname = os.path.normpath(dirname)
    else:
        vars_from_memory = True
T
tangwei12 已提交
864

865
    if vars is None:
866
        if main_program is None:
Y
Yu Yang 已提交
867
            main_program = default_main_program()
868
        if not isinstance(main_program, Program):
869 870
            raise TypeError(
                "The type of input main_program is invalid, expected type is fluid.Program, but received %s"
871 872
                % type(main_program)
            )
873

874 875 876 877 878 879 880
        load_vars(
            executor,
            dirname=dirname,
            main_program=main_program,
            vars=list(filter(predicate, main_program.list_vars())),
            filename=filename,
        )
881 882 883
    else:
        load_prog = Program()
        load_block = load_prog.global_block()
884

885 886
        if main_program is None:
            main_program = default_main_program()
T
tangwei12 已提交
887

888
        if not isinstance(main_program, Program):
889 890
            raise TypeError(
                "The type of input main_program is invalid, expected type is fluid.Program, but received %s"
891 892
                % type(main_program)
            )
893

T
tangwei12 已提交
894
        # save origin param shape
H
hong 已提交
895
        orig_para_shape = {}
896
        load_var_map = {}
897 898 899 900

        check_vars = []
        sparse_vars = []

901 902
        for each_var in vars:
            assert isinstance(each_var, Variable)
903

T
tangwei12 已提交
904 905
            if each_var.type == core.VarDesc.VarType.RAW:
                continue
H
hong 已提交
906 907

            if isinstance(each_var, Parameter):
908
                orig_para_shape[each_var.name] = tuple(
909 910
                    each_var.desc.get_shape()
                )
911 912 913 914 915

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

916
            new_var = _clone_var_in_block_(load_block, each_var)
917 918
            check_vars.append(each_var)

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

933 934 935 936 937
        for each_var in sparse_vars:
            assert isinstance(each_var, Variable)

            if filename is not None:
                raise ValueError(
938 939
                    "SelectedRows can not be load with load_combine"
                )
940 941 942 943 944

            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):
945 946
                raise ValueError(
                    "SelectedRows var {} can not find at {}".format(
947 948 949
                        new_var.name, var_path
                    )
                )
950 951 952 953 954 955

            if os.path.isfile(var_path):
                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [new_var]},
956 957
                    attrs={'file_path': os.path.join(dirname, new_var.name)},
                )
958 959 960 961 962 963 964 965 966 967
            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:
968 969 970 971 972 973 974
                    slice = load_block.create_var(
                        name=block,
                        type=new_var.type,
                        shape=new_var.shape,
                        dtype=new_var.dtype,
                        persistable=False,
                    )
975 976 977
                    slices.append(slice)

                    file_path = os.path.join(var_path, block, "Param")
978 979 980 981 982 983
                    load_block.append_op(
                        type='load',
                        inputs={},
                        outputs={'Out': [slice]},
                        attrs={'file_path': file_path},
                    )
984

985 986 987 988 989 990
                load_block.append_op(
                    type='lookup_sparse_table_merge',
                    inputs={'X': slices},
                    outputs={'Out': new_var},
                    attrs={},
                )
991

992
        if filename is not None:
993 994 995 996
            load_var_list = []
            for name in sorted(load_var_map.keys()):
                load_var_list.append(load_var_map[name])

997 998 999
            if vars_from_memory is False:
                filename = os.path.join(dirname, filename)

1000 1001 1002 1003 1004 1005 1006 1007 1008
            load_block.append_op(
                type='load_combine',
                inputs={},
                outputs={"Out": load_var_list},
                attrs={
                    'file_path': filename,
                    'model_from_memory': vars_from_memory,
                },
            )
1009 1010
        executor.run(load_prog)

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

1030

1031
@dygraph_not_support
1032
def load_params(executor, dirname, main_program=None, filename=None):
1033
    """
1034 1035
    :api_attr: Static Graph

1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054
    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 已提交
1055 1056

    Args:
1057 1058
        executor(Executor): The executor used for loading parameters.
                            See :ref:`api_guide_executor_en` for more details about it.
F
fengjiayi 已提交
1059
        dirname(str): The directory path.
1060 1061 1062 1063 1064 1065 1066 1067
        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 已提交
1068 1069 1070 1071 1072 1073 1074

    Returns:
        None

    Examples:
        .. code-block:: python

1075
            import paddle
1076
            import paddle.fluid as fluid
1077

1078
            paddle.enable_static()
F
fengjiayi 已提交
1079 1080 1081
            exe = fluid.Executor(fluid.CPUPlace())
            param_path = "./my_paddle_model"
            prog = fluid.default_main_program()
1082
            fluid.io.load_params(executor=exe, dirname=param_path,
F
fengjiayi 已提交
1083
                                main_program=None)
1084
    """
1085 1086 1087 1088 1089 1090 1091
    load_vars(
        executor,
        dirname=dirname,
        main_program=main_program,
        predicate=is_parameter,
        filename=filename,
    )
1092 1093


1094
@dygraph_not_support
1095
def load_persistables(executor, dirname, main_program=None, filename=None):
1096
    """
1097
    :api_attr: Static Graph
1098

1099 1100
    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 已提交
1101
    directory ``dirname`` or the file ``filename``.
F
fengjiayi 已提交
1102

1103 1104 1105 1106
    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 已提交
1107 1108

    Args:
1109 1110
        executor(Executor): The executor used for loading persistable variables.
                            See :ref:`api_guide_executor_en` for more details about it.
F
fengjiayi 已提交
1111
        dirname(str): The directory path.
T
tianshuo78520a 已提交
1112
        main_program(Program, optional): The program whose persistable variables will
1113 1114 1115 1116 1117 1118 1119
                                    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 已提交
1120 1121 1122 1123 1124 1125 1126

    Returns:
        None

    Examples:
        .. code-block:: python

1127
            import paddle
1128
            import paddle.fluid as fluid
1129

1130
            paddle.enable_static()
F
fengjiayi 已提交
1131 1132 1133
            exe = fluid.Executor(fluid.CPUPlace())
            param_path = "./my_paddle_model"
            prog = fluid.default_main_program()
1134
            fluid.io.load_persistables(executor=exe, dirname=param_path,
F
fengjiayi 已提交
1135
                                       main_program=None)
1136
    """
1137 1138

    if main_program and main_program._is_distributed:
1139 1140 1141
        _load_distributed_persistables(
            executor, dirname=dirname, main_program=main_program
        )
1142
    else:
1143 1144 1145 1146 1147 1148 1149
        load_vars(
            executor,
            dirname=dirname,
            main_program=main_program,
            predicate=is_persistable,
            filename=filename,
        )
1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169


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

1170
            import paddle
1171
            import paddle.fluid as fluid
1172 1173

            paddle.enable_static()
1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198
            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:
1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216
                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,
                    },
                )
1217
            else:
1218 1219 1220 1221 1222 1223 1224
                origin = load_block.create_var(
                    name="{}".format(origin_var.name),
                    type=origin_var.type,
                    shape=origin_var.shape,
                    dtype=origin_var.dtype,
                    persistable=True,
                )
1225 1226 1227 1228
                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [origin]},
1229 1230
                    attrs={'file_path': os.path.join(dirname, origin_var.name)},
                )
1231 1232 1233

        load_block.append_op(
            type='delete_var',
1234 1235
            inputs={'X': need_delete_vars},
        )
1236 1237 1238 1239

        executor.run(load_prog)

    if not isinstance(main_program, Program):
T
tangwei12 已提交
1240
        raise TypeError("'main_program' should be an instance of Program.")
1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251

    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"
        )

1252 1253 1254 1255 1256
    need_load_vars = (
        main_program._parameters_on_pservers.get_distributed_vars_by_ep(
            main_program._ps_endpoint
        )
    )
1257
    __load_persistable_vars(executor, dirname, need_load_vars)
1258 1259


1260 1261 1262
def prepend_feed_ops(
    inference_program, feed_target_names, feed_holder_name='feed'
):
Q
Qiao Longfei 已提交
1263 1264 1265
    if len(feed_target_names) == 0:
        return

K
Kexin Zhao 已提交
1266
    global_block = inference_program.global_block()
1267 1268 1269 1270 1271
    feed_var = global_block.create_var(
        name=feed_holder_name,
        type=core.VarDesc.VarType.FEED_MINIBATCH,
        persistable=True,
    )
K
Kexin Zhao 已提交
1272

1273
    for i, name in enumerate(feed_target_names):
1274 1275 1276 1277
        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 "
1278 1279 1280 1281
                "if '{name}' is not involved in the target_vars calculation.".format(
                    i=i, name=name
                )
            )
K
fix bug  
Kexin Zhao 已提交
1282
        out = global_block.var(name)
1283 1284 1285 1286 1287 1288
        global_block._prepend_op(
            type='feed',
            inputs={'X': [feed_var]},
            outputs={'Out': [out]},
            attrs={'col': i},
        )
K
Kexin Zhao 已提交
1289 1290


1291 1292 1293
def append_fetch_ops(
    inference_program, fetch_target_names, fetch_holder_name='fetch'
):
K
Kexin Zhao 已提交
1294
    global_block = inference_program.global_block()
1295 1296 1297 1298 1299
    fetch_var = global_block.create_var(
        name=fetch_holder_name,
        type=core.VarDesc.VarType.FETCH_LIST,
        persistable=True,
    )
K
Kexin Zhao 已提交
1300

1301
    for i, name in enumerate(fetch_target_names):
1302 1303 1304 1305 1306 1307
        global_block.append_op(
            type='fetch',
            inputs={'X': [name]},
            outputs={'Out': [fetch_var]},
            attrs={'col': i},
        )
K
Kexin Zhao 已提交
1308 1309


1310 1311
@static_only
@deprecated(since="2.0.0", update_to="paddle.static.save_inference_model")
1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323
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,
):
1324
    """
F
fengjiayi 已提交
1325
    Prune the given `main_program` to build a new program especially for inference,
G
guofei 已提交
1326
    and then save it and all related parameters to given `dirname` .
1327
    If you just want to save parameters of your trained model, please use the
G
guofei 已提交
1328 1329
    :ref:`api_fluid_io_save_params` . You can refer to :ref:`api_guide_model_save_reader_en`
    for more details.
1330

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

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

F
fengjiayi 已提交
1365
    Returns:
1366
        list, The fetch variables' name list.
F
fengjiayi 已提交
1367 1368 1369

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

1371
            import paddle
1372 1373
            import paddle.fluid as fluid

1374
            paddle.enable_static()
F
fengjiayi 已提交
1375 1376
            path = "./infer_model"

T
tianshuo78520a 已提交
1377
            # User defined network, here a softmax regession example
G
guofei 已提交
1378 1379
            image = fluid.data(name='img', shape=[None, 28, 28], dtype='float32')
            label = fluid.data(name='label', shape=[None, 1], dtype='int64')
1380 1381 1382 1383
            feeder = fluid.DataFeeder(feed_list=[image, label], place=fluid.CPUPlace())
            predict = fluid.layers.fc(input=image, size=10, act='softmax')

            loss = fluid.layers.cross_entropy(input=predict, label=label)
1384
            avg_loss = paddle.mean(loss)
1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396

            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 已提交
1397
            # In this example, the save_inference_mode inference will prune the default
1398
            # main program according to the network's input node (img) and output node(predict).
G
guofei 已提交
1399
            # The pruned inference program is going to be saved in the "./infer_model/__model__"
F
fengjiayi 已提交
1400
            # and parameters are going to be saved in separate files under folder
1401
            # "./infer_model".
1402 1403

    """
1404
    if isinstance(feeded_var_names, str):
F
fengjiayi 已提交
1405
        feeded_var_names = [feeded_var_names]
X
Xin Pan 已提交
1406
    elif export_for_deployment:
Q
Qiao Longfei 已提交
1407
        if len(feeded_var_names) > 0:
1408
            # TODO(paddle-dev): polish these code blocks
1409 1410 1411 1412
            if not (
                bool(feeded_var_names)
                and all(isinstance(name, str) for name in feeded_var_names)
            ):
M
minqiyang 已提交
1413
                raise ValueError("'feed_var_names' should be a list of str.")
F
fengjiayi 已提交
1414 1415

    if isinstance(target_vars, Variable):
F
fengjiayi 已提交
1416
        target_vars = [target_vars]
X
Xin Pan 已提交
1417
    elif export_for_deployment:
1418 1419 1420 1421
        if not (
            bool(target_vars)
            and all(isinstance(var, Variable) for var in target_vars)
        ):
F
fengjiayi 已提交
1422 1423
            raise ValueError("'target_vars' should be a list of Variable.")

C
chengduo 已提交
1424
    main_program = _get_valid_program(main_program)
T
tangwei12 已提交
1425

1426
    # remind user to set auc_states to zeros if the program contains auc op
1427 1428
    all_ops = main_program.global_block().ops
    for op in all_ops:
1429 1430 1431
        # clear device of Op
        device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName()
        op._set_attr(device_attr_name, "")
1432 1433 1434 1435 1436 1437
        if op.type == 'auc':
            warnings.warn(
                "please ensure that you have set the auc states to zeros before saving inference model"
            )
            break

1438 1439
    with program_guard(main_program):
        uniq_target_vars = []
F
flame 已提交
1440 1441
        for i, var in enumerate(target_vars):
            uniq_target_vars.append(var)
1442
        target_vars = uniq_target_vars
F
flame 已提交
1443
    target_var_name_list = [var.name for var in target_vars]
1444

1445
    # when a pserver and a trainer running on the same machine, mkdir may conflict
L
lujun 已提交
1446
    save_dirname = dirname
1447
    try:
L
lujun 已提交
1448 1449
        save_dirname = os.path.normpath(dirname)
        os.makedirs(save_dirname)
1450 1451 1452 1453
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise

X
Xin Pan 已提交
1454 1455 1456 1457
    if model_filename is not None:
        model_basename = os.path.basename(model_filename)
    else:
        model_basename = "__model__"
L
lujun 已提交
1458
    model_basename = os.path.join(save_dirname, model_basename)
1459

X
Xin Pan 已提交
1460 1461 1462 1463
    # 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.
1464 1465 1466

    origin_program = main_program.clone()

X
Xin Pan 已提交
1467
    if export_for_deployment:
X
Xin Pan 已提交
1468 1469
        main_program = main_program.clone()
        global_block = main_program.global_block()
1470
        need_to_remove_op_index = []
X
Xin Pan 已提交
1471 1472 1473
        for i, op in enumerate(global_block.ops):
            op.desc.set_is_target(False)
            if op.type == "feed" or op.type == "fetch":
1474 1475 1476 1477 1478
                need_to_remove_op_index.append(i)

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

X
Xin Pan 已提交
1479
        main_program.desc.flush()
X
Xin Pan 已提交
1480

1481
        main_program = main_program._prune_with_input(
1482 1483
            feeded_var_names=feeded_var_names, targets=target_vars
        )
X
Xin Pan 已提交
1484
        main_program = main_program._inference_optimize(prune_read_op=True)
X
Xin Pan 已提交
1485 1486
        fetch_var_names = [v.name for v in target_vars]

1487 1488 1489 1490 1491
        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,
1492
                    dtype=target_v.dtype,
1493 1494
                    persistable=target_v.persistable,
                )
1495

X
Xin Pan 已提交
1496 1497 1498
        prepend_feed_ops(main_program, feeded_var_names)
        append_fetch_ops(main_program, fetch_var_names)

1499
        main_program.desc._set_version()
1500
        paddle.fluid.core.save_op_version_info(main_program.desc)
X
Xin Pan 已提交
1501
        with open(model_basename, "wb") as f:
1502
            f.write(
1503
                main_program._remove_training_info(
1504 1505 1506
                    clip_extra=clip_extra
                ).desc.serialize_to_string()
            )
X
Xin Pan 已提交
1507 1508 1509
    else:
        # TODO(panyx0718): Save more information so that it can also be used
        # for training and more flexible post-processing.
X
Xin Pan 已提交
1510
        with open(model_basename + ".main_program", "wb") as f:
1511
            f.write(
1512
                main_program._remove_training_info(
1513 1514 1515
                    clip_extra=clip_extra
                ).desc.serialize_to_string()
            )
T
tangwei12 已提交
1516

T
tangwei12 已提交
1517 1518 1519 1520 1521 1522
    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

1523 1524
    main_program._copy_dist_param_info_from(origin_program)

X
fix  
Xin Pan 已提交
1525 1526
    if params_filename is not None:
        params_filename = os.path.basename(params_filename)
1527

L
lujun 已提交
1528
    save_persistables(executor, save_dirname, main_program, params_filename)
F
flame 已提交
1529
    return target_var_name_list
X
fix  
Xin Pan 已提交
1530

1531

1532 1533
@static_only
@deprecated(since="2.0.0", update_to="paddle.static.load_inference_model")
1534 1535 1536 1537 1538 1539 1540
def load_inference_model(
    dirname,
    executor,
    model_filename=None,
    params_filename=None,
    pserver_endpoints=None,
):
1541
    """
1542 1543 1544
    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.
1545
    You can refer to :ref:`api_guide_model_save_reader_en` for more details.
1546

F
fengjiayi 已提交
1547
    Args:
1548 1549 1550
        dirname(str): One of the following:
          - The given directory path.
          - Set to None when reading the model from memory.
F
fengjiayi 已提交
1551
        executor(Executor): The executor to run for loading inference model.
1552
                            See :ref:`api_guide_executor_en` for more details about it.
1553 1554 1555 1556 1557 1558 1559
        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:
1560
          - The name of file to load all parameters.
1561 1562 1563
          - 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``.
1564 1565 1566 1567

        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
1568
                                    a list of pserver endpoints.
F
fengjiayi 已提交
1569 1570

    Returns:
1571
        list: The return of this API is a list with three elements:
1572
        (program, feed_target_names, fetch_targets). The `program` is a
1573 1574 1575 1576 1577
        ``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 已提交
1578 1579 1580 1581 1582


    Examples:
        .. code-block:: python

1583
            import paddle
1584 1585
            import paddle.fluid as fluid
            import numpy as np
1586

1587
            paddle.enable_static()
1588
            # Build the model
1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599
            main_prog = fluid.Program()
            startup_prog = fluid.Program()
            with fluid.program_guard(main_prog, startup_prog):
                data = fluid.layers.data(name="img", shape=[64, 784], append_batch_size=False)
                w = fluid.layers.create_parameter(shape=[784, 200], dtype='float32')
                b = fluid.layers.create_parameter(shape=[200], dtype='float32')
                hidden_w = fluid.layers.matmul(x=data, y=w)
                hidden_b = fluid.layers.elementwise_add(hidden_w, b)
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            exe.run(startup_prog)
1600 1601

            # Save the inference model
F
fengjiayi 已提交
1602
            path = "./infer_model"
1603 1604
            fluid.io.save_inference_model(dirname=path, feeded_var_names=['img'],
                         target_vars=[hidden_b], executor=exe, main_program=main_prog)
1605 1606 1607

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

1615 1616 1617
            # 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.
1618
            endpoints = ["127.0.0.1:2023","127.0.0.1:2024"]
1619
            [dist_inference_program, dist_feed_target_names, dist_fetch_targets] = (
1620 1621
                fluid.io.load_inference_model(dirname=path,
                                              executor=exe,
1622
                                              pserver_endpoints=endpoints))
1623

1624
            # In this example, the inference program was saved in the file
1625
            # "./infer_model/__model__" and parameters were saved in
1626 1627 1628 1629
            # 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.
1630
    """
1631 1632 1633 1634
    load_from_memory = False
    if dirname is not None:
        load_dirname = os.path.normpath(dirname)
        if not os.path.isdir(load_dirname):
1635
            raise ValueError("There is no directory named '%s'" % dirname)
1636

1637 1638
        if model_filename is None:
            model_filename = '__model__'
1639

1640 1641 1642
        model_filename = os.path.join(
            load_dirname, os.path.basename(model_filename)
        )
1643

1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654 1655 1656 1657
        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
1658

1659
    program = Program.parse_from_string(program_desc_str)
X
Xin Pan 已提交
1660
    if not core._is_program_version_supported(program._version()):
1661 1662 1663
        raise ValueError(
            "Unsupported program version: %d\n" % program._version()
        )
X
version  
Xin Pan 已提交
1664
    # Binary data also need versioning.
L
lujun 已提交
1665
    load_persistables(executor, load_dirname, program, params_filename)
1666

T
tangwei12 已提交
1667
    if pserver_endpoints:
T
tangwei12 已提交
1668
        program = _endpoints_replacement(program, pserver_endpoints)
T
tangwei12 已提交
1669

1670 1671
    feed_target_names = program.desc.get_feed_target_names()
    fetch_target_names = program.desc.get_fetch_target_names()
1672 1673 1674 1675 1676
    fetch_targets = [
        program.global_block().var(name) for name in fetch_target_names
    ]

    return [program, feed_target_names, fetch_targets]
X
xuwei06 已提交
1677 1678


T
tangwei12 已提交
1679 1680 1681
def _endpoints_replacement(program, endpoints):
    ENDPOINT_MAP = "epmap"
    for op in program.global_block().ops:
T
tangwei12 已提交
1682 1683
        if op.has_attr(ENDPOINT_MAP):
            op.set_attr(ENDPOINT_MAP, endpoints)
T
fix  
tangwei12 已提交
1684
    program._sync_with_cpp()
T
tangwei12 已提交
1685
    return program
T
tangwei12 已提交
1686 1687


X
xuwei06 已提交
1688 1689
def get_parameter_value(para, executor):
    """
F
fengjiayi 已提交
1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700
    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 已提交
1701

F
fengjiayi 已提交
1702 1703
    Examples:
        .. code-block:: python
X
xuwei06 已提交
1704

1705
            import paddle
1706
            import paddle.fluid as fluid
1707 1708

            paddle.enable_static()
F
fengjiayi 已提交
1709 1710 1711
            exe = fluid.Executor(fluid.CPUPlace())
            param = fluid.default_main_program().global_block().var('fc.w')
            p = fluid.io.get_parameter_value(param, exe)
1712

X
xuwei06 已提交
1713
    """
1714
    assert is_parameter(para), "The input variable is not parameter."
X
xuwei06 已提交
1715

X
xuwei06 已提交
1716 1717 1718 1719 1720 1721 1722 1723
    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 已提交
1724
    Get the LoDTensor value of a certain parameter by its name.
X
xuwei06 已提交
1725

F
fengjiayi 已提交
1726 1727 1728 1729 1730 1731 1732
    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 已提交
1733

F
fengjiayi 已提交
1734 1735
    Returns:
        numpy.array: The parameter's values.
1736

F
fengjiayi 已提交
1737 1738 1739
    Examples:
        .. code-block:: python

1740
            import paddle
1741
            import paddle.fluid as fluid
1742 1743

            paddle.enable_static()
F
fengjiayi 已提交
1744 1745
            exe = fluid.Executor(fluid.CPUPlace())
            p = fluid.io.get_parameter_value('fc.w', exe)
X
xuwei06 已提交
1746 1747
    """
    if program is None:
Y
Yu Yang 已提交
1748
        program = default_main_program()
X
xuwei06 已提交
1749 1750
    var = program.global_block().var(name)
    return get_parameter_value(var, executor)
1751 1752 1753 1754 1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765


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:
1766
        name = node.name()
1767 1768 1769 1770 1771 1772 1773
        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()
1774 1775 1776 1777
        if (
            var_desc.type() == core.VarDesc.VarType.RAW
            or var_desc.type() == core.VarDesc.VarType.READER
        ):
1778 1779 1780 1781 1782 1783 1784
            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(),
1785 1786
            persistable=var_desc.persistable(),
        )
1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803
        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:
1804
        name = node.name()
1805 1806 1807 1808 1809 1810 1811 1812 1813 1814 1815
        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()
1816 1817 1818 1819
        if (
            var_desc.type() == core.VarDesc.VarType.RAW
            or var_desc.type() == core.VarDesc.VarType.READER
        ):
1820 1821 1822 1823 1824 1825 1826
            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(),
1827 1828
            persistable=var_desc.persistable(),
        )
1829 1830 1831 1832 1833
        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 已提交
1834 1835


W
WeiXin 已提交
1836
def _unpack_saved_dict(saved_obj, protocol):
1837 1838
    temp_saved_obj = {}
    unpack_infor = {}
W
WeiXin 已提交
1839 1840 1841 1842 1843 1844
    # 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(
1845 1846
                        (2**30 - 1) / value.dtype.itemsize
                    )
W
WeiXin 已提交
1847 1848 1849 1850 1851 1852 1853
                    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(
1854 1855 1856 1857 1858 1859
                            int(
                                math.ceil(
                                    num_element * 1.0 / MAX_NUMBER_OF_ELEMENT
                                )
                            )
                        ):
W
WeiXin 已提交
1860 1861 1862
                            part_name = key + "@@." + str(i)
                            unpack_infor[key]["slices"].append(part_name)
                            temp_saved_obj[part_name] = value[
1863 1864 1865 1866
                                i
                                * MAX_NUMBER_OF_ELEMENT : MAX_NUMBER_OF_ELEMENT
                                * (i + 1)
                            ]
1867 1868 1869 1870 1871 1872 1873 1874 1875 1876 1877 1878

    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 已提交
1879 1880 1881 1882 1883 1884
    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"]]
1885
                load_obj[key] = np.concatenate(slices).reshape(
1886 1887
                    value["OriginShape"]
                )
W
WeiXin 已提交
1888 1889 1890 1891 1892
                removes += value["slices"]
            for key in removes:
                load_obj.pop(key)
            load_obj.pop(unpack_info)

1893 1894 1895
    return load_obj


1896
@static_only
1897 1898
def _legacy_save(param_dict, model_path, protocol=2):
    def get_tensor(var):
J
Jiabin Yang 已提交
1899
        if isinstance(var, (core.VarBase, core.eager.Tensor)):
1900 1901 1902 1903 1904 1905 1906 1907
            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'
1908 1909 1910 1911 1912
    if (
        _is_file_path(model_path)
        and sys.platform == 'darwin'
        and sys.version_info.major == 3
    ):
1913 1914 1915 1916
        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):
1917
                f.write(pickle_bytes[i : i + max_bytes])
1918
    else:
1919
        with _open_file_buffer(model_path, 'wb') as f:
1920 1921 1922 1923
            pickle.dump(param_dict, f, protocol=protocol)


@static_only
1924
def save(program, model_path, protocol=4, **configs):
H
hong 已提交
1925
    """
1926

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

1929 1930
    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 已提交
1931
    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".
1932

H
hong 已提交
1933 1934 1935
    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
1936
        protocol(int, optional): The protocol version of pickle module must be greater than 1 and less than 5.
1937
                                 Default: 4
1938
        configs(dict, optional) : optional keyword arguments.
H
hong 已提交
1939 1940 1941 1942 1943 1944 1945

    Returns:
        None

    Examples:
        .. code-block:: python

1946
            import paddle
1947
            import paddle.static as static
H
hong 已提交
1948

1949
            paddle.enable_static()
H
hong 已提交
1950

1951 1952 1953 1954 1955 1956 1957 1958 1959 1960
            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 已提交
1961 1962 1963
    """

    base_name = os.path.basename(model_path)
1964 1965 1966
    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."
1967 1968 1969 1970 1971
    if 'pickle_protocol' in configs:
        protocol = configs['pickle_protocol']
        warnings.warn(
            "'pickle_protocol' is a deprecated argument. Please use 'protocol' instead."
        )
H
hong 已提交
1972

1973
    if not isinstance(protocol, int):
1974 1975 1976 1977 1978
        raise ValueError(
            "The 'protocol' MUST be `int`, but received {}".format(
                type(protocol)
            )
        )
W
WeiXin 已提交
1979

1980
    if protocol < 2 or protocol > 4:
1981
        raise ValueError(
1982 1983
            "Expected 1<'protocol'<5, but received protocol={}".format(protocol)
        )
W
WeiXin 已提交
1984

1985 1986 1987 1988
    dir_name = os.path.dirname(model_path)
    if dir_name and not os.path.exists(dir_name):
        os.makedirs(dir_name)

Y
Yang Zhang 已提交
1989 1990 1991 1992
    def get_tensor(var):
        t = global_scope().find_var(var.name).get_tensor()
        return np.array(t)

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

1996
    param_dict = _unpack_saved_dict(param_dict, protocol)
1997

1998 1999 2000
    # 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)
2001 2002 2003
        with open(model_path + ".pdparams", 'wb') as f:
            max_bytes = 2**30
            for i in range(0, len(pickle_bytes), max_bytes):
2004
                f.write(pickle_bytes[i : i + max_bytes])
2005 2006
    else:
        with open(model_path + ".pdparams", 'wb') as f:
2007
            pickle.dump(param_dict, f, protocol=protocol)
H
hong 已提交
2008 2009

    optimizer_var_list = list(
2010 2011
        filter(is_belong_to_optimizer, program.list_vars())
    )
H
hong 已提交
2012

Y
Yang Zhang 已提交
2013 2014
    opt_dict = {p.name: get_tensor(p) for p in optimizer_var_list}
    with open(model_path + ".pdopt", 'wb') as f:
2015
        pickle.dump(opt_dict, f, protocol=protocol)
H
hong 已提交
2016 2017 2018 2019

    main_program = program.clone()
    program.desc.flush()
    main_program.desc._set_version()
2020
    paddle.fluid.core.save_op_version_info(program.desc)
H
hong 已提交
2021 2022 2023 2024 2025

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


2026 2027 2028 2029 2030 2031
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 已提交
2032
    load_result = pickle.loads(pickle_bytes, encoding='latin1')
2033 2034 2035
    return load_result


2036
@static_only
H
hong 已提交
2037
def load(program, model_path, executor=None, var_list=None):
H
hong 已提交
2038
    """
2039 2040
    :api_attr: Static Graph

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

2044 2045
    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 已提交
2046 2047
    ( filename is not None When save_params, save_persistables or save_vars is called ).

2048
    Args:
2049 2050
        program(Program): The program will be loaded
        model_path(str): The file prefix store the program
2051
        executor(Executor, optional): The executor used for initialize the parameter
2052
                                      When startup program is not run.
2053
        var_list(list|tuple, optional): The Tensor list/tuple to load single model file saved with
2054
                                  [ save_params, save_persistables, save_vars ].
H
hong 已提交
2055
                                  Default: None
H
hong 已提交
2056 2057 2058

    Returns:
        None
2059

H
hong 已提交
2060 2061 2062
     Examples:
        .. code-block:: python

2063
            import paddle
2064
            import paddle.static as static
H
hong 已提交
2065

2066
            paddle.enable_static()
H
hong 已提交
2067

2068 2069 2070
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
H
hong 已提交
2071

2072 2073 2074 2075 2076 2077 2078
            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 已提交
2079 2080
    """

2081 2082
    assert executor is None or isinstance(executor, Executor)

H
hong 已提交
2083 2084 2085 2086 2087 2088 2089 2090 2091 2092 2093 2094 2095
    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]
2096
        _logger.debug(
2097 2098 2099 2100
            "{} not found, try to load model file saved with [ save_params, save_persistables, save_vars ]".format(
                parameter_file_name
            )
        )
H
hong 已提交
2101 2102 2103 2104
        if executor is None:
            raise ValueError(
                "executor is required when loading model file saved with [ save_params, save_persistables, save_vars ]"
            )
2105 2106 2107 2108 2109 2110

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

H
hong 已提交
2111 2112 2113 2114 2115
        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(
2116 2117
                        os.path.join(root, f).replace("\\", "/")
                    )
H
hong 已提交
2118 2119 2120 2121
            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("\\", "/")
2122 2123 2124
                load_condition = (
                    var_list_names is None or var.name in var_list_names
                )
2125
                if var_path in binary_file_set and load_condition:
H
hong 已提交
2126 2127 2128 2129
                    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))
2130 2131 2132 2133
                _logger.warning(
                    "variable file [ %s ] not used"
                    % (" ".join(list(binary_file_set)))
                )
H
hong 已提交
2134
            try:
2135 2136 2137
                load_vars(
                    executor=executor, dirname=model_path, vars=loaded_var_list
                )
H
hong 已提交
2138 2139 2140 2141 2142
            except RuntimeError as e:
                _logger.error(e)
                raise e
            except:
                raise RuntimeError(
2143
                    "Failed to load model file, please make sure model file is saved with the "
2144 2145
                    "following APIs: save_params, save_persistables, save_vars"
                )
H
hong 已提交
2146 2147 2148

            return
        elif os.path.isfile(model_path):
2149
            if var_list is None:
H
hong 已提交
2150 2151 2152 2153 2154 2155 2156 2157 2158 2159
                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(
2160 2161
                        "loaded var [{}] is not in program variable list"
                    )
H
hong 已提交
2162 2163 2164

            dir_name, file_name = os.path.split(model_path)
            try:
2165 2166 2167 2168 2169 2170
                load_vars(
                    executor=executor,
                    dirname=dir_name,
                    vars=var_list,
                    filename=file_name,
                )
H
hong 已提交
2171 2172 2173 2174
            except RuntimeError as e:
                _logger.error(e)
                raise e
            except:
2175 2176 2177 2178 2179
                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 已提交
2180 2181

            return
Y
Yang Zhang 已提交
2182 2183 2184 2185 2186 2187 2188 2189

    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()
2190 2191 2192 2193
        elif p.is_xpu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.XPUPlace(p.xpu_device_id())
2194 2195 2196 2197
        elif p.is_npu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.NPUPlace(p.npu_device_id())
2198 2199 2200 2201
        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 已提交
2202 2203 2204 2205 2206 2207
        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 已提交
2208 2209

    parameter_list = list(filter(is_parameter, program.list_vars()))
2210 2211

    if executor:
2212 2213 2214
        paddle.fluid.core._create_loaded_parameter(
            parameter_list, global_scope(), executor._default_executor
        )
Y
Yang Zhang 已提交
2215
    with open(parameter_file_name, 'rb') as f:
2216 2217 2218 2219 2220

        # 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 已提交
2221
            load_dict = pickle.load(f, encoding='latin1')
2222
        load_dict = _pack_loaded_dict(load_dict)
Y
Yang Zhang 已提交
2223
    for v in parameter_list:
2224 2225 2226 2227 2228
        assert (
            v.name in load_dict
        ), "Can not find [{}] in model file [{}]".format(
            v.name, parameter_file_name
        )
Y
Yang Zhang 已提交
2229
        set_var(v, load_dict[v.name])
H
hong 已提交
2230 2231

    optimizer_var_list = list(
2232 2233
        filter(is_belong_to_optimizer, program.list_vars())
    )
H
hong 已提交
2234 2235

    if len(optimizer_var_list) > 0:
H
hong 已提交
2236
        opt_file_name = model_prefix + ".pdopt"
2237 2238 2239
        assert os.path.exists(
            opt_file_name
        ), "Optimizer file [{}] not exits".format(opt_file_name)
2240 2241 2242

        if executor:
            paddle.fluid.core._create_loaded_parameter(
2243 2244
                optimizer_var_list, global_scope(), executor._default_executor
            )
Y
Yang Zhang 已提交
2245 2246

        with open(opt_file_name, 'rb') as f:
T
tianshuo78520a 已提交
2247
            load_dict = pickle.load(f, encoding='latin1')
Y
Yang Zhang 已提交
2248
        for v in optimizer_var_list:
2249 2250 2251 2252 2253
            assert (
                v.name in load_dict
            ), "Can not find [{}] in model file [{}]".format(
                v.name, opt_file_name
            )
Y
Yang Zhang 已提交
2254
            set_var(v, load_dict[v.name])
2255 2256


H
hong 已提交
2257
def load_program_state(model_path, var_list=None):
2258
    """
2259

2260
    Load program state from local file
2261

2262 2263
    Args:
        model_path(str): The file prefix store the program
2264
        var_list(list|tuple, optional): The Tensor list/tuple to load saved with
2265
                                  [ save_params, save_persistables, save_vars ].
H
hong 已提交
2266
                                  Default: None.
2267
                                  The var_list is only used to get name,
H
hong 已提交
2268
                                  will not be modified.
2269 2270 2271 2272
    Returns:
        state_dict(dict): the dict store Parameter and optimizer information

    Examples:
2273

2274 2275
        .. code-block:: python

2276
            import paddle
2277
            import paddle.static as static
2278 2279

            paddle.enable_static()
2280

2281 2282 2283
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
2284

2285 2286 2287 2288
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()
2289

2290 2291
            static.save(prog, "./temp")
            program_state = static.load_program_state("./temp")
2292
    """
H
hong 已提交
2293 2294 2295 2296 2297 2298 2299 2300 2301 2302 2303 2304
    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]
2305
        _logger.debug(
2306 2307 2308 2309
            "{} not found, try to load model file saved with [ save_params, save_persistables, save_vars ]".format(
                parameter_file_name
            )
        )
H
hong 已提交
2310 2311 2312 2313

        var_name_list = []
        if var_list is None and os.path.isfile(model_path):
            raise ValueError(
2314 2315
                "var_list can not be None when model_path is a file type"
            )
H
hong 已提交
2316 2317 2318 2319 2320 2321 2322 2323 2324 2325 2326 2327 2328 2329 2330 2331 2332 2333 2334 2335

        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,
2336 2337 2338 2339 2340 2341 2342 2343 2344
                    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
            ):
2345
                try:
2346 2347 2348 2349 2350 2351
                    load_vars(
                        executor=exe,
                        dirname=dirname,
                        vars=vars,
                        filename=filename,
                    )
2352 2353
                    return True
                except:
2354 2355 2356 2357 2358 2359 2360 2361 2362 2363
                    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
                    )
2364 2365 2366 2367 2368 2369 2370 2371 2372
                    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 已提交
2373 2374
            loaded_var_list = []

2375 2376 2377
            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 已提交
2378 2379
                for var in var_list:
                    loaded_var_list.append(clone_var_to_block(load_block, var))
2380 2381 2382
                _load_vars_with_try_catch(
                    exe, dir_name, loaded_var_list, file_name
                )
H
hong 已提交
2383
            else:
2384 2385 2386 2387
                # var_list can be None or not None
                if var_list is not None:
                    for var in var_list:
                        loaded_var_list.append(
2388 2389 2390 2391 2392
                            clone_var_to_block(load_block, var)
                        )
                    _load_vars_with_try_catch(
                        exe, model_path, loaded_var_list, None
                    )
H
hong 已提交
2393
                else:
2394
                    for var_name in var_name_list:
2395 2396 2397 2398
                        # 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
2399
                        # to ensure that other legal variables can be loaded.
2400 2401 2402 2403 2404 2405
                        temp_var = load_block.create_var(
                            name=var_name, persistable=True
                        )
                        if _load_vars_with_try_catch(
                            exe, model_path, [temp_var], None, False
                        ):
2406 2407
                            loaded_var_list.append(temp_var)

H
hong 已提交
2408 2409
            res_dict = {}
            for var in loaded_var_list:
2410
                res_dict[var.name] = np.asarray(
2411 2412
                    paddle.fluid.global_scope().find_var(var.name).get_tensor()
                )
H
hong 已提交
2413 2414 2415

            return res_dict

2416 2417 2418
    assert os.path.exists(
        parameter_file_name
    ), "Parameter file [{}] not exits".format(parameter_file_name)
2419 2420

    with open(parameter_file_name, 'rb') as f:
2421 2422 2423 2424
        # 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 已提交
2425
            para_dict = pickle.load(f, encoding='latin1')
2426
    para_dict = _pack_loaded_dict(para_dict)
2427

H
hong 已提交
2428
    opt_file_name = model_prefix + ".pdopt"
2429 2430
    if os.path.exists(opt_file_name):
        with open(opt_file_name, 'rb') as f:
T
tianshuo78520a 已提交
2431
            opti_dict = pickle.load(f, encoding='latin1')
2432 2433 2434 2435 2436 2437

        para_dict.update(opti_dict)

    return para_dict


2438
@static_only
2439 2440 2441 2442
def set_program_state(program, state_dict):
    """
    Set program parameter from state_dict

2443
    An exception will throw if shape or dtype of the parameters is not match.
2444 2445 2446 2447 2448 2449

    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
2450
    Returns:
2451
        None
2452

2453 2454
    Examples:
        .. code-block:: python
2455

2456
            import paddle
2457
            import paddle.static as static
2458 2459

            paddle.enable_static()
2460

2461 2462 2463
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
2464

2465 2466 2467 2468
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()
2469

2470 2471
            static.save(prog, "./temp")
            program_state = static.load_program_state("./temp")
H
hong 已提交
2472

2473
            static.set_program_state(prog, program_state)
2474
    """
2475
    state_dict = _pack_loaded_dict(state_dict)
2476 2477 2478 2479 2480
    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)
2481
        assert (
2482
            var_temp is not None
2483 2484 2485
        ), "Variable [ {} ] Not found, Please make sure run startup program".format(
            para.name
        )
2486 2487 2488 2489
        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]
2490 2491 2492 2493 2494 2495 2496 2497 2498 2499 2500 2501
            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
                )
            )
2502 2503 2504 2505

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

2506
            # assert ten_place.is_gpu_place() or ten_place.is_cpu_place(), \
Q
QingshuChen 已提交
2507
            #    "Place not support, only support CPUPlace and GPUPlace, now is {}".format(str(ten_place))
2508 2509 2510 2511 2512 2513 2514
            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 已提交
2515 2516 2517 2518
            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())
2519 2520 2521 2522
            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())
2523 2524 2525 2526
            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())
2527 2528 2529 2530 2531 2532 2533 2534 2535 2536 2537

            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(
2538 2539 2540 2541
            "This list is not set, Because of Paramerter not found in program. There are: {}".format(
                " ".join(unused_para_list)
            )
        )