io.py 93.8 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13 14
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

15 16
from __future__ import print_function

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

H
hong 已提交
28
import numpy as np
29
import math
30
import paddle
31
from paddle.fluid import layers
H
hong 已提交
32
from paddle.fluid.executor import Executor, global_scope
33
from paddle.fluid.evaluator import Evaluator
T
tangwei12 已提交
34
from paddle.fluid.framework import Program, Parameter, default_main_program, default_startup_program, Variable, \
35
    program_guard, dygraph_not_support, static_only
36 37
from paddle.reader import cache, map_readers, buffered, compose, chain, shuffle, \
    ComposeNotAligned, firstn, xmap_readers, multiprocess_reader
38
from .wrapped_decorator import signature_safe_contextmanager
T
tangwei12 已提交
39
from paddle.fluid.compiler import CompiledProgram
40
from paddle.fluid.log_helper import get_logger
S
sneaxiy 已提交
41
from . import reader
42
from . import unique_name
S
sneaxiy 已提交
43
from .reader import *
44 45
from . import dataloader
from .dataloader import *
K
fix bug  
Kexin Zhao 已提交
46
from . import core
47
from .. import compat as cpt
48 49
from paddle.utils import deprecated
from paddle.fluid.framework import static_only
50

51 52
batch = paddle.batch

53
__all__ = [
54 55 56 57 58 59 60 61 62 63 64 65 66
    '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 已提交
67 68
    'get_program_parameter',
    'get_program_persistable_vars',
69
] + reader.__all__
70

71 72
_logger = get_logger(
    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s')
73

74

75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120
class _open_buffer(object):
    def __init__(self, buffer):
        self.buffer = buffer

    def __enter__(self):
        return self.buffer


class _buffer_reader(_open_buffer):
    def __init__(self, buffer):
        super(_buffer_reader, self).__init__(buffer)
        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:
            raise ValueError("Expected 'r' or 'w' in mode but got {}".format(
                mode))


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


121
def is_parameter(var):
F
fengjiayi 已提交
122 123
    """
    Check whether the given variable is an instance of Parameter.
124 125

    Args:
F
fengjiayi 已提交
126
        var(Variable): The variable to be checked.
127 128

    Returns:
F
fengjiayi 已提交
129 130 131 132 133 134
        bool: True if the given `var` is an instance of Parameter,
        False if not.

    Examples:
        .. code-block:: python

135
            import paddle
136
            import paddle.fluid as fluid
137 138

            paddle.enable_static()
F
fengjiayi 已提交
139 140
            param = fluid.default_main_program().global_block().var('fc.w')
            res = fluid.io.is_parameter(param)
141
    """
142 143 144 145
    return isinstance(var, Parameter)


def is_persistable(var):
F
fengjiayi 已提交
146 147 148 149 150 151 152 153 154 155 156 157 158
    """
    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

159
            import paddle
160
            import paddle.fluid as fluid
161 162

            paddle.enable_static()
163
            param = fluid.default_main_program().global_block().var('fc.b')
F
fengjiayi 已提交
164 165
            res = fluid.io.is_persistable(param)
    """
166
    if var.desc.type() == core.VarDesc.VarType.FEED_MINIBATCH or \
167 168
                    var.desc.type() == core.VarDesc.VarType.FETCH_LIST or \
                    var.desc.type() == core.VarDesc.VarType.READER:
169
        return False
170 171 172
    return var.persistable


H
hong 已提交
173
def is_belong_to_optimizer(var):
174
    if not (isinstance(var, Parameter) or var.desc.need_check_feed()):
175 176 177
        return is_persistable(var)

    return False
H
hong 已提交
178 179


180
@dygraph_not_support
H
hong 已提交
181 182
def get_program_parameter(program):
    """
183 184
    :api_attr: Static Graph

H
hong 已提交
185 186 187 188 189 190 191 192 193 194 195
    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

196
            import paddle
H
hong 已提交
197
            import paddle.fluid as fluid
198 199

            paddle.enable_static()
H
hong 已提交
200 201 202 203 204 205 206 207
            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()))


208
@dygraph_not_support
H
hong 已提交
209 210
def get_program_persistable_vars(program):
    """
211 212
    :api_attr: Static Graph

H
hong 已提交
213 214 215 216 217 218 219 220 221 222 223
    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

224
            import paddle
H
hong 已提交
225
            import paddle.fluid as fluid
226 227

            paddle.enable_static()
H
hong 已提交
228 229 230 231 232 233 234 235
            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()))


236 237
def _clone_var_in_block_(block, var):
    assert isinstance(var, Variable)
238 239 240 241 242 243 244 245 246 247 248 249 250 251 252
    if var.desc.type() == core.VarDesc.VarType.LOD_TENSOR:
        return block.create_var(
            name=var.name,
            shape=var.shape,
            dtype=var.dtype,
            type=var.type,
            lod_level=var.lod_level,
            persistable=True)
    else:
        return block.create_var(
            name=var.name,
            shape=var.shape,
            dtype=var.dtype,
            type=var.type,
            persistable=True)
253 254


255
@signature_safe_contextmanager
H
hong 已提交
256 257 258 259 260 261 262
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():
263 264
                with paddle.fluid.framework._dygraph_guard(None):
                    yield
H
hong 已提交
265 266


267
def _get_valid_program(main_program=None):
C
chengduo 已提交
268 269 270 271 272
    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:
273 274 275
            raise TypeError(
                "The type of input main_program is invalid, expected tyep is Program, but received None"
            )
C
chengduo 已提交
276 277 278
        warnings.warn(
            "The input is a CompiledProgram, this is not recommended.")
    if not isinstance(main_program, Program):
279 280 281
        raise TypeError(
            "The type of input main_program is invalid, expected type is fluid.Program, but received %s"
            % type(main_program))
C
chengduo 已提交
282 283 284
    return main_program


285
@dygraph_not_support
286 287 288 289 290
def save_vars(executor,
              dirname,
              main_program=None,
              vars=None,
              predicate=None,
291
              filename=None):
292
    """
293 294
    :api_attr: Static Graph

295
    This API saves specific variables in the `Program` to files.
F
fengjiayi 已提交
296

297
    There are two ways to specify the variables to be saved: set variables in
298 299
    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.
300

301
    The `dirname` is used to specify the folder where to save variables.
T
tianshuo78520a 已提交
302
    If you prefer to save variables in separate files in the `dirname` folder,
303
    do not set `filename`. If you prefer to save all variables in a single file,
F
fengjiayi 已提交
304
    use `filename` to specify it.
305

F
fengjiayi 已提交
306 307
    Args:
        executor(Executor): The executor to run for saving variables.
308 309
        dirname(str, optional): The folder where to save variables.
                            When you need to save the parameter to the memory, set it to None.
310
        main_program(Program, optional): The program whose variables will be saved.
311
                                    If it is None, the default main program will
F
fengjiayi 已提交
312 313
                                    be used automatically.
                                    Default: None
314 315 316
        vars(list[Variable], optional): The list contains all variables to be saved.
                                        Default: None
        predicate(function, optional): The function selects the variables that make
317
                                       `predicate(variable) == True`.
318 319
                                       Default: None
        filename(str, optional): If you prefer to save all variables in a single file,
320
                                 use `filename` to specify it. Otherwise, let `filename` be None.
321
                                 Default: None
F
fengjiayi 已提交
322 323

    Returns:
324 325
        str: When saving parameters to a file, returns None.
             When saving parameters to memory, returns a binary string containing parameters.
F
fengjiayi 已提交
326 327 328 329 330 331 332

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

    Examples:
        .. code-block:: python

333
            import paddle
334
            import paddle.fluid as fluid
335

336
            paddle.enable_static()
337 338 339 340 341 342 343 344 345 346 347
            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 已提交
348

349
            # The first usage: use `vars` to set the saved variables.
350 351
            var_list = [w, b]
            path = "./my_paddle_vars"
352
            fluid.io.save_vars(executor=exe, dirname=path, vars=var_list,
353 354 355 356 357 358 359 360 361 362
                            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.
363
    """
364 365 366 367
    save_to_memory = False
    if dirname is None and filename is None:
        save_to_memory = True

C
chengduo 已提交
368
    main_program = _get_valid_program(main_program)
T
tangwei12 已提交
369

370
    if vars is None:
371
        return save_vars(
372
            executor,
373
            main_program=main_program,
374
            dirname=dirname,
375
            vars=list(filter(predicate, main_program.list_vars())),
376
            filename=filename)
377
    else:
378
        params_var_name = unique_name.generate("saved_params")
379 380 381 382 383 384 385
        # 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

386 387
        save_program = Program()
        save_block = save_program.global_block()
388 389

        save_var_map = {}
390
        for each_var in vars:
391 392 393
            # NOTE: don't save the variable which type is RAW
            if each_var.type == core.VarDesc.VarType.RAW:
                continue
394
            new_var = _clone_var_in_block_(save_block, each_var)
395 396 397
            if filename is None and save_to_memory is False:
                save_file_path = os.path.join(
                    os.path.normpath(dirname), new_var.name)
398 399 400 401
                save_block.append_op(
                    type='save',
                    inputs={'X': [new_var]},
                    outputs={},
402
                    attrs={'file_path': os.path.normpath(save_file_path)})
403 404 405
            else:
                save_var_map[new_var.name] = new_var

406
        if filename is not None or save_to_memory:
407 408 409 410
            save_var_list = []
            for name in sorted(save_var_map.keys()):
                save_var_list.append(save_var_map[name])

411 412 413 414 415 416 417
            save_path = str()
            if save_to_memory is False:
                save_path = os.path.join(os.path.normpath(dirname), filename)

            saved_params = save_block.create_var(
                type=core.VarDesc.VarType.RAW, name=params_var_name)
            saved_params.desc.set_persistable(True)
418
            save_block.append_op(
419 420
                type='save_combine',
                inputs={'X': save_var_list},
421 422 423 424 425
                outputs={'Y': saved_params},
                attrs={
                    'file_path': save_path,
                    'save_to_memory': save_to_memory
                })
426

427
        # NOTE(zhiqiu): save op will add variable kLookupTablePath in save_program.desc,
428 429 430
        # which leads to diff on save_program and its desc. Call _sync_with_cpp
        # to keep consistency.
        save_program._sync_with_cpp()
431
        executor.run(save_program)
432 433
        if save_to_memory:
            return global_scope().find_var(params_var_name).get_bytes()
434 435


436
@dygraph_not_support
437
def save_params(executor, dirname, main_program=None, filename=None):
438
    """
439 440
    :api_attr: Static Graph

G
guofei 已提交
441
    This operator saves all parameters from the :code:`main_program` to
442
    the folder :code:`dirname` or file :code:`filename`. You can refer to
G
guofei 已提交
443
    :ref:`api_guide_model_save_reader_en` for more details.
F
fengjiayi 已提交
444

G
guofei 已提交
445 446 447
    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 已提交
448 449
    the file name.

450
    Note:
G
guofei 已提交
451
        Some variables are not Parameter while they are necessary for
452
        training, such as learning rate, global step, etc. So you can NOT save
G
guofei 已提交
453 454
        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`
455 456 457
        and :ref:`api_fluid_io_load_persistables` instead.

        If you want to save your model for the inference, please use the
G
guofei 已提交
458 459
        :ref:`api_fluid_io_save_inference_model`. You can refer to
        :ref:`api_guide_model_save_reader_en` for more details.
F
fengjiayi 已提交
460 461

    Args:
462
        executor(Executor): The executor to run for saving parameters, You can
G
guofei 已提交
463
                            refer to :ref:`api_guide_executor_en`.
464 465
        dirname(str, optional): The saving directory path.
                            When you need to save the parameter to the memory, set it to None.
G
guofei 已提交
466
        main_program(Program, optional): The program whose parameters will be
467 468
                                         saved. You can refer to
                                         :ref:`api_guide_Program_en` for more
G
guofei 已提交
469 470 471 472 473 474 475
                                         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 已提交
476 477

    Returns:
478 479
        str: When saving parameters to a file, returns None.
             When saving parameters to memory, returns a binary string containing parameters.
F
fengjiayi 已提交
480 481 482 483

    Examples:
        .. code-block:: python

484
            import paddle
H
Huihuang Zheng 已提交
485
            import paddle.fluid as fluid
486

487 488

            paddle.enable_static()
G
guofei 已提交
489 490 491 492 493
            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')
494

G
guofei 已提交
495 496
            loss = fluid.layers.cross_entropy(input=predict, label=label)
            avg_loss = fluid.layers.mean(loss)
497

F
fengjiayi 已提交
498
            exe = fluid.Executor(fluid.CPUPlace())
G
guofei 已提交
499 500
            exe.run(fluid.default_startup_program())
            fluid.io.save_params(executor=exe, dirname=params_path)
501 502
            # 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"
503
    """
504
    return save_vars(
505 506
        executor,
        dirname=dirname,
507
        main_program=main_program,
508
        vars=None,
509
        predicate=is_parameter,
510
        filename=filename)
511 512


513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534
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

535
            import paddle
536
            import paddle.fluid as fluid
537 538

            paddle.enable_static()
539 540 541 542 543 544 545 546 547 548
            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 已提交
549
        receive params on pserver through rpc.
550 551 552 553 554 555 556 557 558 559
        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 已提交
560 561 562 563 564 565 566
            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)
567 568 569

            for idx, optimizer in enumerate(remote_params):
                block_id = optimizer.block_id
T
tangwei12 已提交
570
                slice = optimizer.slice
571 572 573
                endpoint = optimizer.endpoint

                index = block_id if is_slice else idx
T
tangwei12 已提交
574 575 576
                slices[index] = slice
                slice_varnames[index] = "{}.slice.{}".format(slice.name, idx)
                remote_varnames[index] = slice.name
577 578
                endpoints[index] = endpoint

T
tangwei12 已提交
579 580 581 582 583
            slice_shapes = []
            for slice in slices:
                tmp = [str(dim) for dim in slice.shape]
                slice_shapes.append(",".join(tmp))

584
            block.append_op(
T
tangwei12 已提交
585 586 587 588 589 590 591 592 593 594 595
                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)
                })

596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624
        executor.run(prog)

    def __save_distributed_lookup_tables(executor, dirname,
                                         distributed_lookup_table, endpoints):
        """
        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
        block.append_op(
            type='checkpoint_notify', inputs={}, outputs={}, attrs=attrs)
        executor.run(prog)

    def __exclude_vars(exclude_var_names=[]):
        def is_valid(var):
            if var.name in exclude_var_names:
                return False
            if var.desc.type() == core.VarDesc.VarType.FEED_MINIBATCH or \
625 626
                            var.desc.type() == core.VarDesc.VarType.FETCH_LIST or \
                            var.desc.type() == core.VarDesc.VarType.READER:
627 628 629 630 631 632
                return False
            return var.persistable

        return is_valid

    if not isinstance(main_program, Program):
T
tangwei12 已提交
633
        raise TypeError("'main_program' should be an instance of Program.")
634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666

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

    remote_params_map = main_program._parameters_on_pservers.get_distributed_vars_by_vtypes(
        ["Optimizer", "RemotePrefetch"], groupby=True)

    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(
        filter(__exclude_vars(exclude_var_names), main_program.list_vars()))
    save_vars(
        executor, main_program=main_program, dirname=dirname, vars=local_vars)

    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(
                executor, dirname, main_program._distributed_lookup_table,
                main_program._endpoints)


667
@dygraph_not_support
668
def save_persistables(executor, dirname, main_program=None, filename=None):
669
    """
670 671
    :api_attr: Static Graph

G
guofei 已提交
672 673 674
    This operator saves all persistable variables from :code:`main_program` to 
    the folder :code:`dirname` or file :code:`filename`. You can refer to 
    :ref:`api_guide_model_save_reader_en` for more details. And then
675 676
    saves these persistables variables to the folder :code:`dirname` or file
    :code:`filename`.
F
fengjiayi 已提交
677

G
guofei 已提交
678
    The :code:`dirname` is used to specify the folder where persistable variables
679
    are going to be saved. If you would like to save variables in separate
G
guofei 已提交
680 681
    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 已提交
682 683 684

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

688 689 690
        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
G
guofei 已提交
691 692
                                         be saved. You can refer to 
                                         :ref:`api_guide_Program_en` for more details.
693
                                         If it is None, the default main program will
G
guofei 已提交
694 695 696 697 698
                                         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 已提交
699 700

    Returns:
701 702
        str: When saving parameters to a file, returns None.
             When saving parameters to memory, returns a binary string containing parameters.
F
fengjiayi 已提交
703 704 705 706

    Examples:
        .. code-block:: python

707
            import paddle
H
Huihuang Zheng 已提交
708
            import paddle.fluid as fluid
709

710
            paddle.enable_static()
G
guofei 已提交
711 712 713 714 715
            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())
716

G
guofei 已提交
717 718 719
            predict = fluid.layers.fc(input=image, size=10, act='softmax')
            loss = fluid.layers.cross_entropy(input=predict, label=label)
            avg_loss = fluid.layers.mean(loss)
F
fengjiayi 已提交
720
            exe = fluid.Executor(fluid.CPUPlace())
G
guofei 已提交
721 722
            exe.run(fluid.default_startup_program())
            fluid.io.save_persistables(executor=exe, dirname=dir_path, filename=file_name)
723
            # The persistables variables weights and bias in the fc layer of the network
G
guofei 已提交
724 725
            # are going to be saved in the same file named "persistables" in the path
            # "./my_paddle_model"
726
    """
727
    if main_program and main_program._is_distributed:
728
        return _save_distributed_persistables(
729 730
            executor, dirname=dirname, main_program=main_program)
    else:
731
        return save_vars(
732 733 734 735 736 737
            executor,
            dirname=dirname,
            main_program=main_program,
            vars=None,
            predicate=is_persistable,
            filename=filename)
738 739


740 741 742 743 744
def load_vars(executor,
              dirname,
              main_program=None,
              vars=None,
              predicate=None,
745
              filename=None):
746
    """
747 748
    :api_attr: Static Graph

749
    This API loads variables from files by executor.
F
fengjiayi 已提交
750

751
    There are two ways to specify the variables to be loaded: the first way, set
752 753
    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`.
754
    The first way has a higher priority.
F
fengjiayi 已提交
755

756
    The `dirname` is used to specify the folder where to load variables.
757
    If variables were saved in separate files in the folder `dirname`,
758
    set `filename` None. If all variables were saved in a single file,
F
fengjiayi 已提交
759
    use `filename` to specify it.
760

F
fengjiayi 已提交
761 762
    Args:
        executor(Executor): The executor to run for loading variables.
763 764
        dirname(str): The folder where to load the variables.
        main_program(Program, optional): The program whose variables will be loaded.
765
                                    If it is None, the default main program will
F
fengjiayi 已提交
766 767
                                    be used automatically.
                                    Default: None
768
        vars(list[Variable], optional): The list that contains all variables to be loaded.
F
fengjiayi 已提交
769
                                   Default: None
770
        predicate(function, optional): The function selects variables that make
771 772 773 774 775
                                        `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 已提交
776 777 778 779 780 781 782 783 784 785

    Returns:
        None

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

    Examples:
        .. code-block:: python

786
            import paddle
787
            import paddle.fluid as fluid
788

789
            paddle.enable_static()
790 791 792 793 794 795 796 797 798 799 800
            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 已提交
801

802 803 804 805 806 807 808 809 810 811 812
            # 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
813
            param_path = "./my_paddle_model"
F
fengjiayi 已提交
814 815 816
            def name_has_fc(var):
                res = "fc" in var.name
                return res
817 818 819
            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 已提交
820
                               vars=None, predicate=name_has_fc)
821 822
            # 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 已提交
823

824
    """
825 826 827 828 829
    vars_from_memory = False
    if dirname is not None:
        dirname = os.path.normpath(dirname)
    else:
        vars_from_memory = True
T
tangwei12 已提交
830

831
    if vars is None:
832
        if main_program is None:
Y
Yu Yang 已提交
833
            main_program = default_main_program()
834
        if not isinstance(main_program, Program):
835 836 837
            raise TypeError(
                "The type of input main_program is invalid, expected type is fluid.Program, but received %s"
                % type(main_program))
838 839 840

        load_vars(
            executor,
841
            dirname=dirname,
T
tangwei12 已提交
842
            main_program=main_program,
843
            vars=list(filter(predicate, main_program.list_vars())),
844
            filename=filename)
845 846 847
    else:
        load_prog = Program()
        load_block = load_prog.global_block()
848

849 850
        if main_program is None:
            main_program = default_main_program()
T
tangwei12 已提交
851

852
        if not isinstance(main_program, Program):
853 854 855
            raise TypeError(
                "The type of input main_program is invalid, expected type is fluid.Program, but received %s"
                % type(main_program))
856

T
tangwei12 已提交
857
        # save origin param shape
H
hong 已提交
858
        orig_para_shape = {}
859
        load_var_map = {}
860 861 862 863

        check_vars = []
        sparse_vars = []

864 865
        for each_var in vars:
            assert isinstance(each_var, Variable)
866

T
tangwei12 已提交
867 868
            if each_var.type == core.VarDesc.VarType.RAW:
                continue
H
hong 已提交
869 870

            if isinstance(each_var, Parameter):
871 872
                orig_para_shape[each_var.name] = tuple(each_var.desc.get_shape(
                ))
873 874 875 876 877

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

878
            new_var = _clone_var_in_block_(load_block, each_var)
879 880
            check_vars.append(each_var)

881
            if filename is None:
882 883 884 885
                if dirname is None:
                    raise ValueError(
                        "The directory path and params cannot be None at the same time."
                    )
886 887 888 889
                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [new_var]},
890
                    attrs={'file_path': os.path.join(dirname, new_var.name)})
891 892 893
            else:
                load_var_map[new_var.name] = new_var

894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944
        for each_var in sparse_vars:
            assert isinstance(each_var, Variable)

            if filename is not None:
                raise ValueError(
                    "SelectedRows can not be load with load_combine")

            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):
                raise ValueError("SelectedRows var {} can not find at {}".
                                 format(new_var.name, var_path))

            if os.path.isfile(var_path):
                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [new_var]},
                    attrs={'file_path': os.path.join(dirname, new_var.name)})
            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:
                    slice = load_block.create_var(
                        name=block,
                        type=new_var.type,
                        shape=new_var.shape,
                        dtype=new_var.dtype,
                        persistable=False)
                    slices.append(slice)

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

                load_block.append_op(
                    type='lookup_sparse_table_merge',
                    inputs={'X': slices},
                    outputs={'Out': new_var},
                    attrs={})

945
        if filename is not None:
946 947 948 949
            load_var_list = []
            for name in sorted(load_var_map.keys()):
                load_var_list.append(load_var_map[name])

950 951 952
            if vars_from_memory is False:
                filename = os.path.join(dirname, filename)

953
            load_block.append_op(
954
                type='load_combine',
955
                inputs={},
956
                outputs={"Out": load_var_list},
957 958 959 960
                attrs={
                    'file_path': filename,
                    'model_from_memory': vars_from_memory
                })
961 962
        executor.run(load_prog)

T
tangwei12 已提交
963
        # check var shape
964
        for each_var in check_vars:
H
hong 已提交
965 966 967 968 969
            if not isinstance(each_var, Parameter):
                continue
            var_temp = paddle.fluid.global_scope().find_var(each_var.name)
            assert var_temp != None, "can't not find var: " + each_var.name
            new_shape = (np.array(var_temp.get_tensor())).shape
970
            assert each_var.name in orig_para_shape, each_var.name + "MUST in var list"
H
hong 已提交
971 972 973
            orig_shape = orig_para_shape.get(each_var.name)
            if new_shape != orig_shape:
                raise RuntimeError(
974
                    "Variable's shape does not match, the Program requires a parameter with the shape of ({}), "
H
hong 已提交
975 976 977
                    "while the loaded parameter (namely [ {} ]) has a shape of  ({}).".
                    format(orig_shape, each_var.name, new_shape))

978

979
@dygraph_not_support
980
def load_params(executor, dirname, main_program=None, filename=None):
981
    """
982 983
    :api_attr: Static Graph

984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002
    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 已提交
1003 1004

    Args:
1005 1006
        executor(Executor): The executor used for loading parameters.
                            See :ref:`api_guide_executor_en` for more details about it.
F
fengjiayi 已提交
1007
        dirname(str): The directory path.
1008 1009 1010 1011 1012 1013 1014 1015
        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 已提交
1016 1017 1018 1019 1020 1021 1022

    Returns:
        None

    Examples:
        .. code-block:: python

1023
            import paddle
1024
            import paddle.fluid as fluid
1025

1026
            paddle.enable_static()
F
fengjiayi 已提交
1027 1028 1029
            exe = fluid.Executor(fluid.CPUPlace())
            param_path = "./my_paddle_model"
            prog = fluid.default_main_program()
1030
            fluid.io.load_params(executor=exe, dirname=param_path,
F
fengjiayi 已提交
1031
                                main_program=None)
1032 1033
    """
    load_vars(
1034 1035 1036
        executor,
        dirname=dirname,
        main_program=main_program,
1037
        predicate=is_parameter,
1038
        filename=filename)
1039 1040


1041
@dygraph_not_support
1042
def load_persistables(executor, dirname, main_program=None, filename=None):
1043
    """
1044
    :api_attr: Static Graph
1045

1046 1047
    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 已提交
1048
    directory ``dirname`` or the file ``filename``.
F
fengjiayi 已提交
1049

1050 1051 1052 1053
    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 已提交
1054 1055

    Args:
1056 1057
        executor(Executor): The executor used for loading persistable variables.
                            See :ref:`api_guide_executor_en` for more details about it.
F
fengjiayi 已提交
1058
        dirname(str): The directory path.
T
tianshuo78520a 已提交
1059
        main_program(Program, optional): The program whose persistable variables will
1060 1061 1062 1063 1064 1065 1066
                                    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 已提交
1067 1068 1069 1070 1071 1072 1073

    Returns:
        None

    Examples:
        .. code-block:: python

1074
            import paddle
1075
            import paddle.fluid as fluid
1076

1077
            paddle.enable_static()
F
fengjiayi 已提交
1078 1079 1080
            exe = fluid.Executor(fluid.CPUPlace())
            param_path = "./my_paddle_model"
            prog = fluid.default_main_program()
1081
            fluid.io.load_persistables(executor=exe, dirname=param_path,
F
fengjiayi 已提交
1082
                                       main_program=None)
1083
    """
1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114

    if main_program and main_program._is_distributed:
        _load_distributed_persistables(
            executor, dirname=dirname, main_program=main_program)
    else:
        load_vars(
            executor,
            dirname=dirname,
            main_program=main_program,
            predicate=is_persistable,
            filename=filename)


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

1115
            import paddle
1116
            import paddle.fluid as fluid
1117 1118

            paddle.enable_static()
1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151
            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:
                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(
T
tangwei12 已提交
1152 1153 1154 1155 1156 1157 1158 1159
                    type='load',
                    inputs={},
                    outputs={'Out': [slice]},
                    attrs={
                        'file_path': os.path.join(dirname, origin_var.name),
                        'seek': offset,
                        'shape': slice.shape
                    })
1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181
            else:
                origin = load_block.create_var(
                    name="{}".format(origin_var.name),
                    type=origin_var.type,
                    shape=origin_var.shape,
                    dtype=origin_var.dtype,
                    persistable=True)
                load_block.append_op(
                    type='load',
                    inputs={},
                    outputs={'Out': [origin]},
                    attrs={
                        'file_path': os.path.join(dirname, origin_var.name)
                    })

        load_block.append_op(
            type='delete_var',
            inputs={'X': need_delete_vars}, )

        executor.run(load_prog)

    if not isinstance(main_program, Program):
T
tangwei12 已提交
1182
        raise TypeError("'main_program' should be an instance of Program.")
1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196

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

    need_load_vars = main_program._parameters_on_pservers.get_distributed_vars_by_ep(
        main_program._ps_endpoint)
    __load_persistable_vars(executor, dirname, need_load_vars)
1197 1198


1199 1200 1201
def prepend_feed_ops(inference_program,
                     feed_target_names,
                     feed_holder_name='feed'):
Q
Qiao Longfei 已提交
1202 1203 1204
    if len(feed_target_names) == 0:
        return

K
Kexin Zhao 已提交
1205 1206
    global_block = inference_program.global_block()
    feed_var = global_block.create_var(
1207 1208 1209
        name=feed_holder_name,
        type=core.VarDesc.VarType.FEED_MINIBATCH,
        persistable=True)
K
Kexin Zhao 已提交
1210

1211
    for i, name in enumerate(feed_target_names):
1212 1213 1214 1215 1216 1217 1218
        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 "
                "if '{name}' is not involved in the target_vars calculation.".
                format(
                    i=i, name=name))
K
fix bug  
Kexin Zhao 已提交
1219
        out = global_block.var(name)
W
Wu Yi 已提交
1220
        global_block._prepend_op(
K
Kexin Zhao 已提交
1221 1222
            type='feed',
            inputs={'X': [feed_var]},
K
fix bug  
Kexin Zhao 已提交
1223
            outputs={'Out': [out]},
K
Kexin Zhao 已提交
1224 1225 1226
            attrs={'col': i})


1227 1228 1229
def append_fetch_ops(inference_program,
                     fetch_target_names,
                     fetch_holder_name='fetch'):
K
Kexin Zhao 已提交
1230 1231
    global_block = inference_program.global_block()
    fetch_var = global_block.create_var(
1232 1233 1234
        name=fetch_holder_name,
        type=core.VarDesc.VarType.FETCH_LIST,
        persistable=True)
K
Kexin Zhao 已提交
1235

1236
    for i, name in enumerate(fetch_target_names):
K
Kexin Zhao 已提交
1237 1238 1239 1240 1241 1242 1243
        global_block.append_op(
            type='fetch',
            inputs={'X': [name]},
            outputs={'Out': [fetch_var]},
            attrs={'col': i})


1244 1245
@static_only
@deprecated(since="2.0.0", update_to="paddle.static.save_inference_model")
1246 1247 1248 1249
def save_inference_model(dirname,
                         feeded_var_names,
                         target_vars,
                         executor,
1250
                         main_program=None,
1251
                         model_filename=None,
1252
                         params_filename=None,
T
tangwei12 已提交
1253
                         export_for_deployment=True,
1254 1255
                         program_only=False,
                         clip_extra=False):
1256
    """
1257 1258
    :api_attr: Static Graph

F
fengjiayi 已提交
1259
    Prune the given `main_program` to build a new program especially for inference,
G
guofei 已提交
1260
    and then save it and all related parameters to given `dirname` .
1261
    If you just want to save parameters of your trained model, please use the
G
guofei 已提交
1262 1263
    :ref:`api_fluid_io_save_params` . You can refer to :ref:`api_guide_model_save_reader_en`
    for more details.
1264

G
guofei 已提交
1265
    Note:
1266
        The :code:`dirname` is used to specify the folder where inference model
G
guofei 已提交
1267
        structure and parameters are going to be saved. If you would like to save params of
1268
        Program in separate files, set `params_filename` None; if you would like to save all
G
guofei 已提交
1269
        params of Program in a single file, use `params_filename` to specify the file name.
F
fengjiayi 已提交
1270 1271 1272

    Args:
        dirname(str): The directory path to save the inference model.
T
tianshuo78520a 已提交
1273
        feeded_var_names(list[str]): list of string. Names of variables that need to be fed
G
guofei 已提交
1274
                                     data during inference.
1275
        target_vars(list[Variable]): list of Variable. Variables from which we can get
G
guofei 已提交
1276
                                     inference results.
1277
        executor(Executor): The executor that saves the inference model. You can refer
G
guofei 已提交
1278 1279
                            to :ref:`api_guide_executor_en` for more details.
        main_program(Program, optional): The original program, which will be pruned to
T
tianshuo78520a 已提交
1280
                                         build the inference model. If is set None,
G
guofei 已提交
1281 1282 1283
                                         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 已提交
1284
                                       itself. If is set None, a default filename
G
guofei 已提交
1285 1286
                                       :code:`__model__` will be used.
        params_filename(str, optional): The name of file to save all related parameters.
T
tianshuo78520a 已提交
1287
                                        If it is set None, parameters will be saved
G
guofei 已提交
1288
                                        in separate files .
X
Xin Pan 已提交
1289 1290 1291 1292 1293
        export_for_deployment(bool): If True, programs are modified to only support
                                     direct inference deployment. Otherwise,
                                     more information will be stored for flexible
                                     optimization and re-training. Currently, only
                                     True is supported.
G
guofei 已提交
1294
                                     Default: True.
1295
        program_only(bool, optional): If True, It will save inference program only, and do not
G
guofei 已提交
1296 1297
                                      save params of Program.
                                      Default: False.
1298

F
fengjiayi 已提交
1299
    Returns:
G
guofei 已提交
1300 1301 1302 1303
        The fetch variables' name list

     Return Type:
        list
F
fengjiayi 已提交
1304 1305

    Raises:
G
guofei 已提交
1306 1307
        ValueError: If `feed_var_names` is not a list of basestring, an exception is thrown.
        ValueError: If `target_vars` is not a list of Variable, an exception is thrown.
F
fengjiayi 已提交
1308 1309 1310

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

1312
            import paddle
1313 1314
            import paddle.fluid as fluid

1315
            paddle.enable_static()
F
fengjiayi 已提交
1316 1317
            path = "./infer_model"

T
tianshuo78520a 已提交
1318
            # User defined network, here a softmax regession example
G
guofei 已提交
1319 1320
            image = fluid.data(name='img', shape=[None, 28, 28], dtype='float32')
            label = fluid.data(name='label', shape=[None, 1], dtype='int64')
1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337
            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)
            avg_loss = fluid.layers.mean(loss)

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

    """
M
minqiyang 已提交
1345
    if isinstance(feeded_var_names, six.string_types):
F
fengjiayi 已提交
1346
        feeded_var_names = [feeded_var_names]
X
Xin Pan 已提交
1347
    elif export_for_deployment:
Q
Qiao Longfei 已提交
1348
        if len(feeded_var_names) > 0:
1349
            # TODO(paddle-dev): polish these code blocks
Q
Qiao Longfei 已提交
1350
            if not (bool(feeded_var_names) and all(
M
minqiyang 已提交
1351
                    isinstance(name, six.string_types)
1352
                    for name in feeded_var_names)):
M
minqiyang 已提交
1353
                raise ValueError("'feed_var_names' should be a list of str.")
F
fengjiayi 已提交
1354 1355

    if isinstance(target_vars, Variable):
F
fengjiayi 已提交
1356
        target_vars = [target_vars]
X
Xin Pan 已提交
1357
    elif export_for_deployment:
1358 1359
        if not (bool(target_vars) and
                all(isinstance(var, Variable) for var in target_vars)):
F
fengjiayi 已提交
1360 1361
            raise ValueError("'target_vars' should be a list of Variable.")

C
chengduo 已提交
1362
    main_program = _get_valid_program(main_program)
T
tangwei12 已提交
1363

1364
    # remind user to set auc_states to zeros if the program contains auc op
1365 1366
    all_ops = main_program.global_block().ops
    for op in all_ops:
1367 1368 1369
        # clear device of Op
        device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName()
        op._set_attr(device_attr_name, "")
1370 1371 1372 1373 1374 1375
        if op.type == 'auc':
            warnings.warn(
                "please ensure that you have set the auc states to zeros before saving inference model"
            )
            break

1376 1377
    with program_guard(main_program):
        uniq_target_vars = []
F
flame 已提交
1378 1379
        for i, var in enumerate(target_vars):
            uniq_target_vars.append(var)
1380
        target_vars = uniq_target_vars
F
flame 已提交
1381
    target_var_name_list = [var.name for var in target_vars]
1382

1383
    # when a pserver and a trainer running on the same machine, mkdir may conflict
L
lujun 已提交
1384
    save_dirname = dirname
1385
    try:
L
lujun 已提交
1386 1387
        save_dirname = os.path.normpath(dirname)
        os.makedirs(save_dirname)
1388 1389 1390 1391
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise

X
Xin Pan 已提交
1392 1393 1394 1395
    if model_filename is not None:
        model_basename = os.path.basename(model_filename)
    else:
        model_basename = "__model__"
L
lujun 已提交
1396
    model_basename = os.path.join(save_dirname, model_basename)
1397

X
Xin Pan 已提交
1398 1399 1400 1401
    # 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.
1402 1403 1404

    origin_program = main_program.clone()

X
Xin Pan 已提交
1405
    if export_for_deployment:
X
Xin Pan 已提交
1406 1407
        main_program = main_program.clone()
        global_block = main_program.global_block()
1408
        need_to_remove_op_index = []
X
Xin Pan 已提交
1409 1410 1411
        for i, op in enumerate(global_block.ops):
            op.desc.set_is_target(False)
            if op.type == "feed" or op.type == "fetch":
1412 1413 1414 1415 1416
                need_to_remove_op_index.append(i)

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

X
Xin Pan 已提交
1417
        main_program.desc.flush()
X
Xin Pan 已提交
1418

1419 1420
        main_program = main_program._prune_with_input(
            feeded_var_names=feeded_var_names, targets=target_vars)
X
Xin Pan 已提交
1421
        main_program = main_program._inference_optimize(prune_read_op=True)
X
Xin Pan 已提交
1422 1423
        fetch_var_names = [v.name for v in target_vars]

1424 1425 1426 1427 1428
        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,
1429 1430
                    dtype=target_v.dtype,
                    persistable=target_v.persistable)
1431

X
Xin Pan 已提交
1432 1433 1434
        prepend_feed_ops(main_program, feeded_var_names)
        append_fetch_ops(main_program, fetch_var_names)

1435
        main_program.desc._set_version()
1436
        paddle.fluid.core.save_op_version_info(main_program.desc)
X
Xin Pan 已提交
1437
        with open(model_basename, "wb") as f:
1438 1439 1440
            f.write(
                main_program._remove_training_info(clip_extra=clip_extra)
                .desc.serialize_to_string())
X
Xin Pan 已提交
1441 1442 1443
    else:
        # TODO(panyx0718): Save more information so that it can also be used
        # for training and more flexible post-processing.
X
Xin Pan 已提交
1444
        with open(model_basename + ".main_program", "wb") as f:
1445 1446 1447
            f.write(
                main_program._remove_training_info(clip_extra=clip_extra)
                .desc.serialize_to_string())
T
tangwei12 已提交
1448

T
tangwei12 已提交
1449 1450 1451 1452 1453 1454
    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

1455 1456
    main_program._copy_dist_param_info_from(origin_program)

X
fix  
Xin Pan 已提交
1457 1458
    if params_filename is not None:
        params_filename = os.path.basename(params_filename)
1459

L
lujun 已提交
1460
    save_persistables(executor, save_dirname, main_program, params_filename)
F
flame 已提交
1461
    return target_var_name_list
X
fix  
Xin Pan 已提交
1462

1463

1464 1465
@static_only
@deprecated(since="2.0.0", update_to="paddle.static.load_inference_model")
1466 1467 1468
def load_inference_model(dirname,
                         executor,
                         model_filename=None,
T
tangwei12 已提交
1469 1470
                         params_filename=None,
                         pserver_endpoints=None):
1471
    """
1472 1473
    :api_attr: Static Graph

1474 1475 1476
    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.
1477
    You can refer to :ref:`api_guide_model_save_reader_en` for more details.
1478

F
fengjiayi 已提交
1479
    Args:
1480 1481 1482
        dirname(str): One of the following:
          - The given directory path.
          - Set to None when reading the model from memory.
F
fengjiayi 已提交
1483
        executor(Executor): The executor to run for loading inference model.
1484
                            See :ref:`api_guide_executor_en` for more details about it.
1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495
        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:
          - The name of file to load all parameters.  
          - 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``.
1496 1497 1498 1499

        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
1500
                                    a list of pserver endpoints.
F
fengjiayi 已提交
1501 1502

    Returns:
1503
        list: The return of this API is a list with three elements:
1504
        (program, feed_target_names, fetch_targets). The `program` is a
1505 1506 1507 1508 1509
        ``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 已提交
1510 1511 1512 1513 1514 1515 1516

    Raises:
        ValueError: If `dirname` is not a existing directory.

    Examples:
        .. code-block:: python

1517
            import paddle
1518 1519
            import paddle.fluid as fluid
            import numpy as np
1520

1521
            paddle.enable_static()
1522
            # Build the model
1523 1524 1525 1526 1527 1528 1529 1530 1531 1532 1533
            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)
1534 1535

            # Save the inference model
F
fengjiayi 已提交
1536
            path = "./infer_model"
1537 1538
            fluid.io.save_inference_model(dirname=path, feeded_var_names=['img'],
                         target_vars=[hidden_b], executor=exe, main_program=main_prog)
1539 1540 1541

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

1549 1550 1551
            # 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.
1552
            endpoints = ["127.0.0.1:2023","127.0.0.1:2024"]
1553
            [dist_inference_program, dist_feed_target_names, dist_fetch_targets] = (
1554 1555
                fluid.io.load_inference_model(dirname=path,
                                              executor=exe,
1556
                                              pserver_endpoints=endpoints))
1557

1558
            # In this example, the inference program was saved in the file
1559
            # "./infer_model/__model__" and parameters were saved in
1560 1561 1562 1563
            # 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.
1564
    """
1565 1566 1567 1568
    load_from_memory = False
    if dirname is not None:
        load_dirname = os.path.normpath(dirname)
        if not os.path.isdir(load_dirname):
1569
            raise ValueError("There is no directory named '%s'" % dirname)
1570

1571 1572
        if model_filename is None:
            model_filename = '__model__'
1573

1574 1575
        model_filename = os.path.join(load_dirname,
                                      os.path.basename(model_filename))
1576

1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590
        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
1591

1592
    program = Program.parse_from_string(program_desc_str)
X
Xin Pan 已提交
1593
    if not core._is_program_version_supported(program._version()):
X
version  
Xin Pan 已提交
1594 1595 1596
        raise ValueError("Unsupported program version: %d\n" %
                         program._version())
    # Binary data also need versioning.
L
lujun 已提交
1597
    load_persistables(executor, load_dirname, program, params_filename)
1598

T
tangwei12 已提交
1599
    if pserver_endpoints:
T
tangwei12 已提交
1600
        program = _endpoints_replacement(program, pserver_endpoints)
T
tangwei12 已提交
1601

1602 1603
    feed_target_names = program.desc.get_feed_target_names()
    fetch_target_names = program.desc.get_fetch_target_names()
1604 1605 1606 1607 1608
    fetch_targets = [
        program.global_block().var(name) for name in fetch_target_names
    ]

    return [program, feed_target_names, fetch_targets]
X
xuwei06 已提交
1609 1610


T
tangwei12 已提交
1611 1612 1613
def _endpoints_replacement(program, endpoints):
    ENDPOINT_MAP = "epmap"
    for op in program.global_block().ops:
T
tangwei12 已提交
1614 1615
        if op.has_attr(ENDPOINT_MAP):
            op.set_attr(ENDPOINT_MAP, endpoints)
T
fix  
tangwei12 已提交
1616
    program._sync_with_cpp()
T
tangwei12 已提交
1617
    return program
T
tangwei12 已提交
1618 1619


X
xuwei06 已提交
1620 1621
def get_parameter_value(para, executor):
    """
F
fengjiayi 已提交
1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632
    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 已提交
1633

F
fengjiayi 已提交
1634 1635
    Examples:
        .. code-block:: python
X
xuwei06 已提交
1636

1637
            import paddle
1638
            import paddle.fluid as fluid
1639 1640

            paddle.enable_static()
F
fengjiayi 已提交
1641 1642 1643
            exe = fluid.Executor(fluid.CPUPlace())
            param = fluid.default_main_program().global_block().var('fc.w')
            p = fluid.io.get_parameter_value(param, exe)
1644

X
xuwei06 已提交
1645
    """
1646
    assert is_parameter(para), "The input variable is not parameter."
X
xuwei06 已提交
1647

X
xuwei06 已提交
1648 1649 1650 1651 1652 1653 1654 1655
    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 已提交
1656
    Get the LoDTensor value of a certain parameter by its name.
X
xuwei06 已提交
1657

F
fengjiayi 已提交
1658 1659 1660 1661 1662 1663 1664
    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 已提交
1665

F
fengjiayi 已提交
1666 1667
    Returns:
        numpy.array: The parameter's values.
1668

F
fengjiayi 已提交
1669 1670 1671
    Raises:
        TypeError: If given `name` is not an instance of basestring.
        TypeError: If the parameter with the given name doesn't exist.
T
tianshuo78520a 已提交
1672
        AssertionError: If there is a variable named `name` in the
F
fengjiayi 已提交
1673
                        given program but it is not a Parameter.
1674

F
fengjiayi 已提交
1675 1676 1677
    Examples:
        .. code-block:: python

1678
            import paddle
1679
            import paddle.fluid as fluid
1680 1681

            paddle.enable_static()
F
fengjiayi 已提交
1682 1683
            exe = fluid.Executor(fluid.CPUPlace())
            p = fluid.io.get_parameter_value('fc.w', exe)
X
xuwei06 已提交
1684 1685
    """
    if program is None:
Y
Yu Yang 已提交
1686
        program = default_main_program()
X
xuwei06 已提交
1687 1688
    var = program.global_block().var(name)
    return get_parameter_value(var, executor)
1689 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700 1701 1702 1703 1704 1705 1706 1707 1708 1709 1710 1711 1712


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:
        name = cpt.to_text(node.name())
        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()
        if var_desc.type() == core.VarDesc.VarType.RAW or \
1713
                        var_desc.type() == core.VarDesc.VarType.READER:
1714 1715 1716 1717 1718 1719 1720 1721 1722 1723 1724 1725 1726 1727 1728 1729 1730 1731 1732 1733 1734 1735 1736 1737 1738 1739 1740 1741 1742 1743 1744 1745 1746 1747 1748 1749 1750 1751
            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(),
            persistable=var_desc.persistable())
        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:
        name = cpt.to_text(node.name())
        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()
        if var_desc.type() == core.VarDesc.VarType.RAW or \
1752
                        var_desc.type() == core.VarDesc.VarType.READER:
1753 1754 1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765
            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(),
            persistable=var_desc.persistable())
        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 已提交
1766 1767


W
WeiXin 已提交
1768
def _unpack_saved_dict(saved_obj, protocol):
1769 1770
    temp_saved_obj = {}
    unpack_infor = {}
W
WeiXin 已提交
1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786 1787 1788 1789 1790 1791 1792
    # 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(
                        (2**30 - 1) / value.dtype.itemsize)
                    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(
                                int(
                                    math.ceil(num_element * 1.0 /
                                              MAX_NUMBER_OF_ELEMENT))):
                            part_name = key + "@@." + str(i)
                            unpack_infor[key]["slices"].append(part_name)
                            temp_saved_obj[part_name] = value[
                                i * MAX_NUMBER_OF_ELEMENT:MAX_NUMBER_OF_ELEMENT
                                * (i + 1)]
1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804

    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 已提交
1805 1806 1807 1808 1809 1810 1811 1812 1813 1814 1815 1816 1817
    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"]]
                load_obj[key] = np.concatenate(slices).reshape(value[
                    "OriginShape"])
                removes += value["slices"]
            for key in removes:
                load_obj.pop(key)
            load_obj.pop(unpack_info)

1818 1819 1820
    return load_obj


1821
@static_only
1822 1823
def _legacy_save(param_dict, model_path, protocol=2):
    def get_tensor(var):
J
Jiabin Yang 已提交
1824
        if isinstance(var, (core.VarBase, core.eager.Tensor)):
1825 1826 1827 1828 1829 1830 1831 1832
            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'
1833 1834 1835
    if _is_file_path(
            model_path
    ) and sys.platform == 'darwin' and sys.version_info.major == 3:
1836 1837 1838 1839 1840 1841
        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):
                f.write(pickle_bytes[i:i + max_bytes])
    else:
1842
        with _open_file_buffer(model_path, 'wb') as f:
1843 1844 1845 1846
            pickle.dump(param_dict, f, protocol=protocol)


@static_only
1847
def save(program, model_path, protocol=4, **configs):
H
hong 已提交
1848
    """
1849 1850
    :api_attr: Static Graph

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

1853 1854
    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 已提交
1855
    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".
1856

H
hong 已提交
1857 1858 1859
    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
1860
        protocol(int, optional): The protocol version of pickle module must be greater than 1 and less than 5.
1861
                                 Default: 4
1862
        configs(dict, optional) : optional keyword arguments.                        
H
hong 已提交
1863 1864 1865 1866 1867 1868 1869

    Returns:
        None

    Examples:
        .. code-block:: python

1870
            import paddle
1871
            import paddle.static as static
H
hong 已提交
1872

1873
            paddle.enable_static()
H
hong 已提交
1874

1875 1876 1877 1878 1879 1880 1881 1882 1883 1884
            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 已提交
1885 1886 1887 1888
    """

    base_name = os.path.basename(model_path)
    assert base_name != "", \
1889
        "The input model_path MUST be format of dirname/filename [dirname\\filename in Windows system], but received model_path is empty string."
1890 1891 1892 1893 1894
    if 'pickle_protocol' in configs:
        protocol = configs['pickle_protocol']
        warnings.warn(
            "'pickle_protocol' is a deprecated argument. Please use 'protocol' instead."
        )
H
hong 已提交
1895

1896
    if not isinstance(protocol, int):
W
WeiXin 已提交
1897
        raise ValueError("The 'protocol' MUST be `int`, but received {}".format(
1898
            type(protocol)))
W
WeiXin 已提交
1899

1900
    if protocol < 2 or protocol > 4:
W
WeiXin 已提交
1901
        raise ValueError("Expected 1<'protocol'<5, but received protocol={}".
1902
                         format(protocol))
W
WeiXin 已提交
1903

1904 1905 1906 1907
    dir_name = os.path.dirname(model_path)
    if dir_name and not os.path.exists(dir_name):
        os.makedirs(dir_name)

Y
Yang Zhang 已提交
1908 1909 1910 1911
    def get_tensor(var):
        t = global_scope().find_var(var.name).get_tensor()
        return np.array(t)

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

1915
    param_dict = _unpack_saved_dict(param_dict, protocol)
1916

1917 1918 1919
    # 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)
1920 1921 1922 1923 1924 1925
        with open(model_path + ".pdparams", 'wb') as f:
            max_bytes = 2**30
            for i in range(0, len(pickle_bytes), max_bytes):
                f.write(pickle_bytes[i:i + max_bytes])
    else:
        with open(model_path + ".pdparams", 'wb') as f:
1926
            pickle.dump(param_dict, f, protocol=protocol)
H
hong 已提交
1927 1928 1929 1930

    optimizer_var_list = list(
        filter(is_belong_to_optimizer, program.list_vars()))

Y
Yang Zhang 已提交
1931 1932
    opt_dict = {p.name: get_tensor(p) for p in optimizer_var_list}
    with open(model_path + ".pdopt", 'wb') as f:
1933
        pickle.dump(opt_dict, f, protocol=protocol)
H
hong 已提交
1934 1935 1936 1937

    main_program = program.clone()
    program.desc.flush()
    main_program.desc._set_version()
1938
    paddle.fluid.core.save_op_version_info(program.desc)
H
hong 已提交
1939 1940 1941 1942 1943

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


1944 1945 1946 1947 1948 1949
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 已提交
1950
    load_result = pickle.loads(pickle_bytes, encoding='latin1')
1951 1952 1953
    return load_result


1954
@static_only
H
hong 已提交
1955
def load(program, model_path, executor=None, var_list=None):
H
hong 已提交
1956
    """
1957 1958
    :api_attr: Static Graph

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

1962 1963
    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 已提交
1964 1965
    ( filename is not None When save_params, save_persistables or save_vars is called ).

1966
    Args:
1967 1968
        program(Program): The program will be loaded
        model_path(str): The file prefix store the program
1969
        executor(Executor, optional): The executor used for initialize the parameter
1970
                                      When startup program is not run.
1971
        var_list(list|tuple, optional): The Tensor list/tuple to load single model file saved with
1972
                                  [ save_params, save_persistables, save_vars ].
H
hong 已提交
1973
                                  Default: None
H
hong 已提交
1974 1975 1976

    Returns:
        None
1977

H
hong 已提交
1978 1979 1980
     Examples:
        .. code-block:: python

1981
            import paddle
1982
            import paddle.static as static
H
hong 已提交
1983

1984
            paddle.enable_static()
H
hong 已提交
1985

1986 1987 1988
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
H
hong 已提交
1989

1990 1991 1992 1993 1994 1995 1996
            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 已提交
1997 1998
    """

1999 2000
    assert executor is None or isinstance(executor, Executor)

H
hong 已提交
2001 2002 2003 2004 2005 2006 2007 2008 2009 2010 2011 2012 2013
    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]
2014
        _logger.debug(
H
hong 已提交
2015 2016 2017 2018 2019 2020
            "{} not found, try to load model file saved with [ save_params, save_persistables, save_vars ]".
            format(parameter_file_name))
        if executor is None:
            raise ValueError(
                "executor is required when loading model file saved with [ save_params, save_persistables, save_vars ]"
            )
2021 2022 2023 2024 2025 2026

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

H
hong 已提交
2027 2028 2029 2030 2031 2032 2033 2034 2035 2036
        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(
                        os.path.join(root, f).replace("\\", "/"))
            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("\\", "/")
2037 2038
                load_condition = var_list_names is None or var.name in var_list_names
                if var_path in binary_file_set and load_condition:
H
hong 已提交
2039 2040 2041 2042 2043 2044 2045 2046 2047 2048 2049 2050 2051 2052
                    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))
                _logger.warning("variable file [ %s ] not used" %
                                (" ".join(list(binary_file_set))))
            try:
                load_vars(
                    executor=executor, dirname=model_path, vars=loaded_var_list)
            except RuntimeError as e:
                _logger.error(e)
                raise e
            except:
                raise RuntimeError(
2053
                    "Failed to load model file, please make sure model file is saved with the "
H
hong 已提交
2054 2055 2056 2057 2058 2059 2060 2061 2062 2063 2064 2065 2066 2067 2068
                    "following APIs: save_params, save_persistables, save_vars")

            return
        elif os.path.isfile(model_path):
            if var_list == None:
                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(
2069
                        "loaded var [{}] is not in program variable list")
H
hong 已提交
2070 2071 2072 2073 2074 2075 2076 2077 2078 2079 2080 2081

            dir_name, file_name = os.path.split(model_path)
            try:
                load_vars(
                    executor=executor,
                    dirname=dir_name,
                    vars=var_list,
                    filename=file_name)
            except RuntimeError as e:
                _logger.error(e)
                raise e
            except:
2082 2083 2084
                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 已提交
2085 2086

            return
Y
Yang Zhang 已提交
2087 2088 2089 2090 2091 2092 2093 2094

    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()
2095 2096 2097 2098
        elif p.is_xpu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.XPUPlace(p.xpu_device_id())
2099 2100 2101 2102
        elif p.is_npu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.NPUPlace(p.npu_device_id())
2103 2104 2105 2106
        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 已提交
2107 2108 2109 2110 2111 2112
        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 已提交
2113 2114

    parameter_list = list(filter(is_parameter, program.list_vars()))
2115 2116 2117 2118 2119

    if executor:
        paddle.fluid.core._create_loaded_parameter(parameter_list,
                                                   global_scope(),
                                                   executor._default_executor)
Y
Yang Zhang 已提交
2120
    with open(parameter_file_name, 'rb') as f:
2121 2122 2123 2124 2125

        # 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 已提交
2126
            load_dict = pickle.load(f, encoding='latin1')
2127
        load_dict = _pack_loaded_dict(load_dict)
Y
Yang Zhang 已提交
2128 2129 2130 2131 2132
    for v in parameter_list:
        assert v.name in load_dict, \
            "Can not find [{}] in model file [{}]".format(
                v.name, parameter_file_name)
        set_var(v, load_dict[v.name])
H
hong 已提交
2133 2134 2135 2136 2137

    optimizer_var_list = list(
        filter(is_belong_to_optimizer, program.list_vars()))

    if len(optimizer_var_list) > 0:
H
hong 已提交
2138
        opt_file_name = model_prefix + ".pdopt"
H
hong 已提交
2139
        assert os.path.exists(opt_file_name), \
T
tangwei12 已提交
2140
            "Optimizer file [{}] not exits".format(opt_file_name)
2141 2142 2143 2144

        if executor:
            paddle.fluid.core._create_loaded_parameter(
                optimizer_var_list, global_scope(), executor._default_executor)
Y
Yang Zhang 已提交
2145 2146

        with open(opt_file_name, 'rb') as f:
T
tianshuo78520a 已提交
2147
            load_dict = pickle.load(f, encoding='latin1')
Y
Yang Zhang 已提交
2148 2149 2150 2151 2152
        for v in optimizer_var_list:
            assert v.name in load_dict, \
                "Can not find [{}] in model file [{}]".format(
                    v.name, opt_file_name)
            set_var(v, load_dict[v.name])
2153 2154


H
hong 已提交
2155
def load_program_state(model_path, var_list=None):
2156
    """
2157 2158
    :api_attr: Static Graph

2159
    Load program state from local file
2160

2161 2162
    Args:
        model_path(str): The file prefix store the program
2163
        var_list(list|tuple, optional): The Tensor list/tuple to load saved with
2164
                                  [ save_params, save_persistables, save_vars ].
H
hong 已提交
2165
                                  Default: None.
2166
                                  The var_list is only used to get name,
H
hong 已提交
2167
                                  will not be modified.
2168 2169 2170 2171 2172 2173
    Returns:
        state_dict(dict): the dict store Parameter and optimizer information

    Examples:
        .. code-block:: python

2174
            import paddle
2175
            import paddle.static as static
2176 2177

            paddle.enable_static()
2178

2179 2180 2181
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
2182

2183 2184 2185 2186
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()
2187

2188 2189
            static.save(prog, "./temp")
            program_state = static.load_program_state("./temp")
2190
    """
H
hong 已提交
2191 2192 2193 2194 2195 2196 2197 2198 2199 2200 2201 2202
    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]
2203
        _logger.debug(
H
hong 已提交
2204 2205 2206 2207 2208 2209 2210 2211 2212 2213 2214 2215 2216 2217 2218 2219 2220 2221 2222 2223 2224 2225 2226 2227 2228 2229 2230 2231 2232 2233 2234 2235
            "{} not found, try to load model file saved with [ save_params, save_persistables, save_vars ]".
            format(parameter_file_name))

        var_name_list = []
        if var_list is None and os.path.isfile(model_path):
            raise ValueError(
                "var_list can not be None when model_path is a file type")

        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,
                    lod_level=var.lod_level
                    if var.desc.type() == core.VarDesc.VarType.LOD_TENSOR else
                    None,
                    persistable=True)

2236 2237 2238 2239 2240 2241 2242 2243 2244 2245 2246 2247 2248 2249 2250 2251 2252 2253 2254 2255 2256 2257 2258 2259 2260 2261 2262
            def _load_vars_with_try_catch(exe,
                                          dirname,
                                          vars,
                                          filename,
                                          raise_error=True):
                try:
                    load_vars(
                        executor=exe,
                        dirname=dirname,
                        vars=vars,
                        filename=filename)
                    return True
                except:
                    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
                    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 已提交
2263 2264
            loaded_var_list = []

2265 2266 2267
            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 已提交
2268 2269
                for var in var_list:
                    loaded_var_list.append(clone_var_to_block(load_block, var))
2270 2271
                _load_vars_with_try_catch(exe, dir_name, loaded_var_list,
                                          file_name)
H
hong 已提交
2272
            else:
2273 2274 2275 2276 2277 2278 2279
                # var_list can be None or not None
                if var_list is not None:
                    for var in var_list:
                        loaded_var_list.append(
                            clone_var_to_block(load_block, var))
                    _load_vars_with_try_catch(exe, model_path, loaded_var_list,
                                              None)
H
hong 已提交
2280
                else:
2281 2282 2283 2284 2285 2286 2287 2288 2289 2290 2291 2292
                    for var_name in var_name_list:
                        # 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 
                        # to ensure that other legal variables can be loaded.
                        temp_var = load_block.create_var(
                            name=var_name, persistable=True)
                        if _load_vars_with_try_catch(exe, model_path,
                                                     [temp_var], None, False):
                            loaded_var_list.append(temp_var)

H
hong 已提交
2293 2294 2295 2296 2297 2298 2299
            res_dict = {}
            for var in loaded_var_list:
                res_dict[var.name] = np.asarray(paddle.fluid.global_scope(
                ).find_var(var.name).get_tensor())

            return res_dict

2300
    assert os.path.exists(parameter_file_name), \
T
tangwei12 已提交
2301
        "Parameter file [{}] not exits".format(parameter_file_name)
2302 2303

    with open(parameter_file_name, 'rb') as f:
2304 2305 2306 2307
        # 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 已提交
2308
            para_dict = pickle.load(f, encoding='latin1')
2309
    para_dict = _pack_loaded_dict(para_dict)
2310

H
hong 已提交
2311
    opt_file_name = model_prefix + ".pdopt"
2312 2313
    if os.path.exists(opt_file_name):
        with open(opt_file_name, 'rb') as f:
T
tianshuo78520a 已提交
2314
            opti_dict = pickle.load(f, encoding='latin1')
2315 2316 2317 2318 2319 2320

        para_dict.update(opti_dict)

    return para_dict


2321
@static_only
2322 2323
def set_program_state(program, state_dict):
    """
2324 2325
    :api_attr: Static Graph

2326 2327
    Set program parameter from state_dict

2328
    An exception will throw if shape or dtype of the parameters is not match.
2329 2330 2331 2332 2333 2334

    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
2335
    Returns:
2336
        None
2337

2338 2339
    Examples:
        .. code-block:: python
2340

2341
            import paddle
2342
            import paddle.static as static
2343 2344

            paddle.enable_static()
2345

2346 2347 2348
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
2349

2350 2351 2352 2353
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()
2354

2355 2356
            static.save(prog, "./temp")
            program_state = static.load_program_state("./temp")
H
hong 已提交
2357

2358
            static.set_program_state(prog, program_state)
2359
    """
2360
    state_dict = _pack_loaded_dict(state_dict)
2361 2362 2363 2364 2365 2366
    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)
        assert var_temp != None, \
T
tangwei12 已提交
2367
            "Variable [ {} ] Not found, Please make sure run startup program".format(para.name)
2368 2369 2370 2371
        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]
T
tangwei12 已提交
2372
            assert orig_para_np.shape == new_para_np.shape, \
2373
                "Parameter's shape does not match, the Program requires a parameter with the shape of ({}), " \
T
tangwei12 已提交
2374
                "while the loaded parameter (namely [ {} ]) has a shape of  ({})." \
2375
                    .format(orig_para_np.shape, para.name, new_para_np.shape)
T
tangwei12 已提交
2376
            assert orig_para_np.dtype == new_para_np.dtype, \
2377
                "Parameter's data type does not match, the Program requires a parameter with a dtype of ({}), " \
T
tangwei12 已提交
2378
                "while the loaded parameter (namely [ {} ]) has a dtype of  ({})." \
2379 2380 2381 2382 2383
                    .format(orig_para_np.dtype, para.name, new_para_np.dtype)

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

Q
QingshuChen 已提交
2384 2385
            #assert ten_place.is_gpu_place() or ten_place.is_cpu_place(), \
            #    "Place not support, only support CPUPlace and GPUPlace, now is {}".format(str(ten_place))
2386 2387 2388 2389 2390 2391 2392
            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 已提交
2393 2394 2395 2396
            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())
2397 2398 2399 2400
            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())
2401 2402 2403 2404
            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())
2405 2406 2407 2408 2409 2410 2411 2412 2413 2414 2415 2416 2417

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