io.py 93.4 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 1045
    :api_attr: Static Graph
    
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 1254
                         export_for_deployment=True,
                         program_only=False):
1255
    """
1256 1257
    :api_attr: Static Graph

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

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

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

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

     Return Type:
        list
F
fengjiayi 已提交
1303 1304

    Raises:
G
guofei 已提交
1305 1306
        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 已提交
1307 1308 1309

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

1311
            import paddle
1312 1313
            import paddle.fluid as fluid

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

T
tianshuo78520a 已提交
1317
            # User defined network, here a softmax regession example
G
guofei 已提交
1318 1319
            image = fluid.data(name='img', shape=[None, 28, 28], dtype='float32')
            label = fluid.data(name='label', shape=[None, 1], dtype='int64')
1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336
            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 已提交
1337
            # In this example, the save_inference_mode inference will prune the default
1338
            # main program according to the network's input node (img) and output node(predict).
G
guofei 已提交
1339
            # The pruned inference program is going to be saved in the "./infer_model/__model__"
F
fengjiayi 已提交
1340
            # and parameters are going to be saved in separate files under folder
1341
            # "./infer_model".
1342 1343

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

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

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

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

1375 1376 1377 1378 1379
    # fix the bug that the activation op's output as target will be pruned.
    # will affect the inference performance.
    # TODO(Superjomn) add an IR pass to remove 1-scale op.
    with program_guard(main_program):
        uniq_target_vars = []
F
flame 已提交
1380
        for i, var in enumerate(target_vars):
1381
            if isinstance(var, Variable):
F
flame 已提交
1382 1383 1384
                var = layers.scale(
                    var, 1., name="save_infer_model/scale_{}".format(i))
            uniq_target_vars.append(var)
1385
        target_vars = uniq_target_vars
F
flame 已提交
1386
    target_var_name_list = [var.name for var in target_vars]
1387

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

X
Xin Pan 已提交
1397 1398 1399 1400
    if model_filename is not None:
        model_basename = os.path.basename(model_filename)
    else:
        model_basename = "__model__"
L
lujun 已提交
1401
    model_basename = os.path.join(save_dirname, model_basename)
1402

X
Xin Pan 已提交
1403 1404 1405 1406
    # 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.
1407 1408 1409

    origin_program = main_program.clone()

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

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

X
Xin Pan 已提交
1422
        main_program.desc.flush()
X
Xin Pan 已提交
1423

1424 1425
        main_program = main_program._prune_with_input(
            feeded_var_names=feeded_var_names, targets=target_vars)
X
Xin Pan 已提交
1426
        main_program = main_program._inference_optimize(prune_read_op=True)
X
Xin Pan 已提交
1427 1428
        fetch_var_names = [v.name for v in target_vars]

X
Xin Pan 已提交
1429 1430 1431
        prepend_feed_ops(main_program, feeded_var_names)
        append_fetch_ops(main_program, fetch_var_names)

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

T
tangwei12 已提交
1444 1445 1446 1447 1448 1449
    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

1450 1451
    main_program._copy_dist_param_info_from(origin_program)

X
fix  
Xin Pan 已提交
1452 1453
    if params_filename is not None:
        params_filename = os.path.basename(params_filename)
1454

L
lujun 已提交
1455
    save_persistables(executor, save_dirname, main_program, params_filename)
F
flame 已提交
1456
    return target_var_name_list
X
fix  
Xin Pan 已提交
1457

1458

1459 1460
@static_only
@deprecated(since="2.0.0", update_to="paddle.static.load_inference_model")
1461 1462 1463
def load_inference_model(dirname,
                         executor,
                         model_filename=None,
T
tangwei12 已提交
1464 1465
                         params_filename=None,
                         pserver_endpoints=None):
1466
    """
1467 1468
    :api_attr: Static Graph

1469 1470 1471
    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.
1472
    You can refer to :ref:`api_guide_model_save_reader_en` for more details.
1473

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

        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
1495
                                    a list of pserver endpoints.
F
fengjiayi 已提交
1496 1497

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

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

    Examples:
        .. code-block:: python

1512
            import paddle
1513 1514
            import paddle.fluid as fluid
            import numpy as np
1515

1516
            paddle.enable_static()
1517
            # Build the model
1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528
            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)
1529 1530

            # Save the inference model
F
fengjiayi 已提交
1531
            path = "./infer_model"
1532 1533
            fluid.io.save_inference_model(dirname=path, feeded_var_names=['img'],
                         target_vars=[hidden_b], executor=exe, main_program=main_prog)
1534 1535 1536

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

1544 1545 1546
            # 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.
1547
            endpoints = ["127.0.0.1:2023","127.0.0.1:2024"]
1548
            [dist_inference_program, dist_feed_target_names, dist_fetch_targets] = (
1549 1550
                fluid.io.load_inference_model(dirname=path,
                                              executor=exe,
1551
                                              pserver_endpoints=endpoints))
1552

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

1566 1567
        if model_filename is None:
            model_filename = '__model__'
1568

1569 1570
        model_filename = os.path.join(load_dirname,
                                      os.path.basename(model_filename))
1571

1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585
        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
1586

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

T
tangwei12 已提交
1594
    if pserver_endpoints:
T
tangwei12 已提交
1595
        program = _endpoints_replacement(program, pserver_endpoints)
T
tangwei12 已提交
1596

1597 1598
    feed_target_names = program.desc.get_feed_target_names()
    fetch_target_names = program.desc.get_fetch_target_names()
1599 1600 1601 1602 1603
    fetch_targets = [
        program.global_block().var(name) for name in fetch_target_names
    ]

    return [program, feed_target_names, fetch_targets]
X
xuwei06 已提交
1604 1605


T
tangwei12 已提交
1606 1607 1608
def _endpoints_replacement(program, endpoints):
    ENDPOINT_MAP = "epmap"
    for op in program.global_block().ops:
T
tangwei12 已提交
1609 1610
        if op.has_attr(ENDPOINT_MAP):
            op.set_attr(ENDPOINT_MAP, endpoints)
T
fix  
tangwei12 已提交
1611
    program._sync_with_cpp()
T
tangwei12 已提交
1612
    return program
T
tangwei12 已提交
1613 1614


X
xuwei06 已提交
1615 1616
def get_parameter_value(para, executor):
    """
F
fengjiayi 已提交
1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627
    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 已提交
1628

F
fengjiayi 已提交
1629 1630
    Examples:
        .. code-block:: python
X
xuwei06 已提交
1631

1632
            import paddle
1633
            import paddle.fluid as fluid
1634 1635

            paddle.enable_static()
F
fengjiayi 已提交
1636 1637 1638
            exe = fluid.Executor(fluid.CPUPlace())
            param = fluid.default_main_program().global_block().var('fc.w')
            p = fluid.io.get_parameter_value(param, exe)
1639

X
xuwei06 已提交
1640
    """
1641
    assert is_parameter(para), "The input variable is not parameter."
X
xuwei06 已提交
1642

X
xuwei06 已提交
1643 1644 1645 1646 1647 1648 1649 1650
    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 已提交
1651
    Get the LoDTensor value of a certain parameter by its name.
X
xuwei06 已提交
1652

F
fengjiayi 已提交
1653 1654 1655 1656 1657 1658 1659
    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 已提交
1660

F
fengjiayi 已提交
1661 1662
    Returns:
        numpy.array: The parameter's values.
1663

F
fengjiayi 已提交
1664 1665 1666
    Raises:
        TypeError: If given `name` is not an instance of basestring.
        TypeError: If the parameter with the given name doesn't exist.
T
tianshuo78520a 已提交
1667
        AssertionError: If there is a variable named `name` in the
F
fengjiayi 已提交
1668
                        given program but it is not a Parameter.
1669

F
fengjiayi 已提交
1670 1671 1672
    Examples:
        .. code-block:: python

1673
            import paddle
1674
            import paddle.fluid as fluid
1675 1676

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


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 \
1708
                        var_desc.type() == core.VarDesc.VarType.READER:
1709 1710 1711 1712 1713 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
            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 \
1747
                        var_desc.type() == core.VarDesc.VarType.READER:
1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758 1759 1760
            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 已提交
1761 1762


W
WeiXin 已提交
1763
def _unpack_saved_dict(saved_obj, protocol):
1764 1765
    temp_saved_obj = {}
    unpack_infor = {}
W
WeiXin 已提交
1766 1767 1768 1769 1770 1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786 1787
    # 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)]
1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799

    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 已提交
1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811 1812
    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)

1813 1814 1815
    return load_obj


1816
@static_only
1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827
def _legacy_save(param_dict, model_path, protocol=2):
    def get_tensor(var):
        if isinstance(var, core.VarBase):
            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'
1828 1829 1830
    if _is_file_path(
            model_path
    ) and sys.platform == 'darwin' and sys.version_info.major == 3:
1831 1832 1833 1834 1835 1836
        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:
1837
        with _open_file_buffer(model_path, 'wb') as f:
1838 1839 1840 1841
            pickle.dump(param_dict, f, protocol=protocol)


@static_only
1842
def save(program, model_path, protocol=4, **configs):
H
hong 已提交
1843
    """
1844 1845
    :api_attr: Static Graph

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

1848 1849
    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 已提交
1850
    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".
1851

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

    Returns:
        None

    Examples:
        .. code-block:: python

1865
            import paddle
1866
            import paddle.static as static
H
hong 已提交
1867

1868
            paddle.enable_static()
H
hong 已提交
1869

1870 1871 1872 1873 1874 1875 1876 1877 1878 1879
            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 已提交
1880 1881 1882 1883
    """

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

1891
    if not isinstance(protocol, int):
W
WeiXin 已提交
1892
        raise ValueError("The 'protocol' MUST be `int`, but received {}".format(
1893
            type(protocol)))
W
WeiXin 已提交
1894

1895
    if protocol < 2 or protocol > 4:
W
WeiXin 已提交
1896
        raise ValueError("Expected 1<'protocol'<5, but received protocol={}".
1897
                         format(protocol))
W
WeiXin 已提交
1898

1899 1900 1901 1902
    dir_name = os.path.dirname(model_path)
    if dir_name and not os.path.exists(dir_name):
        os.makedirs(dir_name)

Y
Yang Zhang 已提交
1903 1904 1905 1906
    def get_tensor(var):
        t = global_scope().find_var(var.name).get_tensor()
        return np.array(t)

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

1910
    param_dict = _unpack_saved_dict(param_dict, protocol)
1911

1912 1913 1914
    # 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)
1915 1916 1917 1918 1919 1920
        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:
1921
            pickle.dump(param_dict, f, protocol=protocol)
H
hong 已提交
1922 1923 1924 1925

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

Y
Yang Zhang 已提交
1926 1927
    opt_dict = {p.name: get_tensor(p) for p in optimizer_var_list}
    with open(model_path + ".pdopt", 'wb') as f:
1928
        pickle.dump(opt_dict, f, protocol=protocol)
H
hong 已提交
1929 1930 1931 1932

    main_program = program.clone()
    program.desc.flush()
    main_program.desc._set_version()
1933
    paddle.fluid.core.save_op_version_info(program.desc)
H
hong 已提交
1934 1935 1936 1937 1938

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


1939 1940 1941 1942 1943 1944
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 已提交
1945
    load_result = pickle.loads(pickle_bytes, encoding='latin1')
1946 1947 1948
    return load_result


1949
@static_only
H
hong 已提交
1950
def load(program, model_path, executor=None, var_list=None):
H
hong 已提交
1951
    """
1952 1953
    :api_attr: Static Graph

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

1957 1958
    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 已提交
1959 1960
    ( filename is not None When save_params, save_persistables or save_vars is called ).

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

    Returns:
        None
1972

H
hong 已提交
1973 1974 1975
     Examples:
        .. code-block:: python

1976
            import paddle
1977
            import paddle.static as static
H
hong 已提交
1978

1979
            paddle.enable_static()
H
hong 已提交
1980

1981 1982 1983
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
H
hong 已提交
1984

1985 1986 1987 1988 1989 1990 1991
            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 已提交
1992 1993
    """

1994 1995
    assert executor is None or isinstance(executor, Executor)

H
hong 已提交
1996 1997 1998 1999 2000 2001 2002 2003 2004 2005 2006 2007 2008
    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]
2009
        _logger.debug(
H
hong 已提交
2010 2011 2012 2013 2014 2015
            "{} 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 ]"
            )
2016 2017 2018 2019 2020 2021

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

H
hong 已提交
2022 2023 2024 2025 2026 2027 2028 2029 2030 2031
        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("\\", "/")
2032 2033
                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 已提交
2034 2035 2036 2037 2038 2039 2040 2041 2042 2043 2044 2045 2046 2047
                    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(
2048
                    "Failed to load model file, please make sure model file is saved with the "
H
hong 已提交
2049 2050 2051 2052 2053 2054 2055 2056 2057 2058 2059 2060 2061 2062 2063
                    "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(
2064
                        "loaded var [{}] is not in program variable list")
H
hong 已提交
2065 2066 2067 2068 2069 2070 2071 2072 2073 2074 2075 2076

            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:
2077 2078 2079
                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 已提交
2080 2081

            return
Y
Yang Zhang 已提交
2082 2083 2084 2085 2086 2087 2088 2089

    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()
2090 2091 2092 2093
        elif p.is_xpu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.XPUPlace(p.xpu_device_id())
2094 2095 2096 2097
        elif p.is_npu_place():
            p = paddle.fluid.core.Place()
            p.set_place(t._place())
            place = paddle.fluid.NPUPlace(p.npu_device_id())
Y
Yang Zhang 已提交
2098 2099 2100 2101 2102 2103
        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 已提交
2104 2105

    parameter_list = list(filter(is_parameter, program.list_vars()))
2106 2107 2108 2109 2110

    if executor:
        paddle.fluid.core._create_loaded_parameter(parameter_list,
                                                   global_scope(),
                                                   executor._default_executor)
Y
Yang Zhang 已提交
2111
    with open(parameter_file_name, 'rb') as f:
2112 2113 2114 2115 2116

        # 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 已提交
2117
            load_dict = pickle.load(f, encoding='latin1')
2118
        load_dict = _pack_loaded_dict(load_dict)
Y
Yang Zhang 已提交
2119 2120 2121 2122 2123
    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 已提交
2124 2125 2126 2127 2128

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

    if len(optimizer_var_list) > 0:
H
hong 已提交
2129
        opt_file_name = model_prefix + ".pdopt"
H
hong 已提交
2130
        assert os.path.exists(opt_file_name), \
T
tangwei12 已提交
2131
            "Optimizer file [{}] not exits".format(opt_file_name)
2132 2133 2134 2135

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

        with open(opt_file_name, 'rb') as f:
T
tianshuo78520a 已提交
2138
            load_dict = pickle.load(f, encoding='latin1')
Y
Yang Zhang 已提交
2139 2140 2141 2142 2143
        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])
2144 2145


H
hong 已提交
2146
def load_program_state(model_path, var_list=None):
2147
    """
2148 2149
    :api_attr: Static Graph

2150
    Load program state from local file
2151

2152 2153
    Args:
        model_path(str): The file prefix store the program
2154
        var_list(list|tuple, optional): The Tensor list/tuple to load saved with
2155
                                  [ save_params, save_persistables, save_vars ].
H
hong 已提交
2156
                                  Default: None.
2157
                                  The var_list is only used to get name,
H
hong 已提交
2158
                                  will not be modified.
2159 2160 2161 2162 2163 2164
    Returns:
        state_dict(dict): the dict store Parameter and optimizer information

    Examples:
        .. code-block:: python

2165
            import paddle
2166
            import paddle.static as static
2167 2168

            paddle.enable_static()
2169

2170 2171 2172
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
2173

2174 2175 2176 2177
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()
2178

2179 2180
            static.save(prog, "./temp")
            program_state = static.load_program_state("./temp")
2181
    """
H
hong 已提交
2182 2183 2184 2185 2186 2187 2188 2189 2190 2191 2192 2193
    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]
2194
        _logger.debug(
H
hong 已提交
2195 2196 2197 2198 2199 2200 2201 2202 2203 2204 2205 2206 2207 2208 2209 2210 2211 2212 2213 2214 2215 2216 2217 2218 2219 2220 2221 2222 2223 2224 2225 2226
            "{} 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)

2227 2228 2229 2230 2231 2232 2233 2234 2235 2236 2237 2238 2239 2240 2241 2242 2243 2244 2245 2246 2247 2248 2249 2250 2251 2252 2253
            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 已提交
2254 2255
            loaded_var_list = []

2256 2257 2258
            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 已提交
2259 2260
                for var in var_list:
                    loaded_var_list.append(clone_var_to_block(load_block, var))
2261 2262
                _load_vars_with_try_catch(exe, dir_name, loaded_var_list,
                                          file_name)
H
hong 已提交
2263
            else:
2264 2265 2266 2267 2268 2269 2270
                # 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 已提交
2271
                else:
2272 2273 2274 2275 2276 2277 2278 2279 2280 2281 2282 2283
                    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 已提交
2284 2285 2286 2287 2288 2289 2290
            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

2291
    assert os.path.exists(parameter_file_name), \
T
tangwei12 已提交
2292
        "Parameter file [{}] not exits".format(parameter_file_name)
2293 2294

    with open(parameter_file_name, 'rb') as f:
2295 2296 2297 2298
        # 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 已提交
2299
            para_dict = pickle.load(f, encoding='latin1')
2300
    para_dict = _pack_loaded_dict(para_dict)
2301

H
hong 已提交
2302
    opt_file_name = model_prefix + ".pdopt"
2303 2304
    if os.path.exists(opt_file_name):
        with open(opt_file_name, 'rb') as f:
T
tianshuo78520a 已提交
2305
            opti_dict = pickle.load(f, encoding='latin1')
2306 2307 2308 2309 2310 2311

        para_dict.update(opti_dict)

    return para_dict


2312
@static_only
2313 2314
def set_program_state(program, state_dict):
    """
2315 2316
    :api_attr: Static Graph

2317 2318
    Set program parameter from state_dict

2319
    An exception will throw if shape or dtype of the parameters is not match.
2320 2321 2322 2323 2324 2325

    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
2326
    Returns:
2327
        None
2328

2329 2330
    Examples:
        .. code-block:: python
2331

2332
            import paddle
2333
            import paddle.static as static
2334 2335

            paddle.enable_static()
2336

2337 2338 2339
            x = static.data(name="x", shape=[10, 10], dtype='float32')
            y = static.nn.fc(x, 10)
            z = static.nn.fc(y, 10)
2340

2341 2342 2343 2344
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
            prog = static.default_main_program()
2345

2346 2347
            static.save(prog, "./temp")
            program_state = static.load_program_state("./temp")
H
hong 已提交
2348

2349
            static.set_program_state(prog, program_state)
2350
    """
2351
    state_dict = _pack_loaded_dict(state_dict)
2352 2353 2354 2355 2356 2357
    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 已提交
2358
            "Variable [ {} ] Not found, Please make sure run startup program".format(para.name)
2359 2360 2361 2362
        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 已提交
2363
            assert orig_para_np.shape == new_para_np.shape, \
2364
                "Parameter's shape does not match, the Program requires a parameter with the shape of ({}), " \
T
tangwei12 已提交
2365
                "while the loaded parameter (namely [ {} ]) has a shape of  ({})." \
2366
                    .format(orig_para_np.shape, para.name, new_para_np.shape)
T
tangwei12 已提交
2367
            assert orig_para_np.dtype == new_para_np.dtype, \
2368
                "Parameter's data type does not match, the Program requires a parameter with a dtype of ({}), " \
T
tangwei12 已提交
2369
                "while the loaded parameter (namely [ {} ]) has a dtype of  ({})." \
2370 2371 2372 2373 2374
                    .format(orig_para_np.dtype, para.name, new_para_np.dtype)

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

Q
QingshuChen 已提交
2375 2376
            #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))
2377 2378 2379 2380 2381 2382 2383
            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 已提交
2384 2385 2386 2387
            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())
2388 2389 2390 2391
            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())
2392 2393 2394 2395 2396 2397 2398 2399 2400 2401 2402 2403 2404

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