jit.py 51.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# 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
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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 18
import os
import pickle
19
import warnings
20
import functools
21
from collections import OrderedDict
22 23

import six
24
import paddle
25
from paddle.fluid import core
26 27
from paddle.fluid.compiler import BuildStrategy, CompiledProgram, ExecutionStrategy
from paddle.fluid.data_feeder import check_type
28
from paddle.fluid.layers.utils import flatten, pack_sequence_as
29
from paddle.fluid.dygraph.base import program_desc_tracing_guard, switch_to_static_graph
30
from paddle.fluid.dygraph.dygraph_to_static import logging_utils
31
from paddle.fluid.dygraph.dygraph_to_static.convert_call_func import ConversionOptions, CONVERSION_OPTIONS
32
from paddle.fluid.dygraph.dygraph_to_static.logging_utils import set_code_level, set_verbosity
33
from paddle.fluid.dygraph.dygraph_to_static.program_translator import ProgramTranslator, StaticFunction, unwrap_decorators
34
from paddle.fluid.dygraph.io import TranslatedLayer, INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX, INFER_PARAMS_INFO_SUFFIX
35 36
from paddle.fluid.dygraph.layers import Layer
from paddle.fluid.executor import Executor, scope_guard
37 38 39
from paddle.fluid.framework import Block, ParamBase, Program, Variable
from paddle.fluid.framework import _current_expected_place, _dygraph_guard, _dygraph_tracer
from paddle.fluid.framework import dygraph_only, in_dygraph_mode
40
from paddle.fluid.wrapped_decorator import wrap_decorator
41

42 43
__all__ = [
    'TracedLayer', 'declarative', 'dygraph_to_static_func', 'set_code_level',
44
    'set_verbosity', 'save', 'load', 'not_to_static'
45
]
46 47 48 49 50 51 52 53 54 55


def create_program_from_desc(program_desc):
    program = Program()
    program.desc = program_desc
    program.blocks = [Block(program, 0)]
    program._sync_with_cpp()
    return program


56
def _extract_vars(inputs, result_list, err_tag='inputs'):
57
    if isinstance(inputs, Variable):
58
        result_list.append(inputs)
59
    elif isinstance(inputs, (list, tuple)):
60
        for var in inputs:
61
            _extract_vars(var, result_list, err_tag)
62 63
    else:
        raise TypeError(
64 65
            "The type of 'each element of {}' in fluid.dygraph.jit.TracedLayer.trace must be fluid.Variable, but received {}.".
            format(err_tag, type(inputs)))
66 67


68
def extract_vars(inputs, err_tag='inputs'):
69
    result_list = []
70
    _extract_vars(inputs, result_list, err_tag)
71 72 73
    return result_list


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 121 122
def _dygraph_to_static_func_(dygraph_func):
    """
    Converts imperative dygraph APIs into declarative function APIs. Decorator
    @dygraph_to_static_func only converts imperative dygraph APIs into
    declarative net-building APIs, which means it doesn't return immediate
    digital result as imperative mode. Users should handle Program and Executor
    by themselves.

    Note:
    This decorator is NOT our recommended way to transform imperative function
    to declarative function. We will remove this decorator after we finalize
    cleaning up code.

    Args:
        dygraph_func (callable): callable imperative function.

    Returns:
        Callable: converting imperative dygraph APIs into declarative
        net-building APIs.

    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          import numpy as np
          from paddle.fluid.dygraph.jit import dygraph_to_static_func

          @dygraph_to_static_func
          def func(x):
              if fluid.layers.mean(x) < 0:
                  x_v = x - 1
              else:
                  x_v = x + 1

               return x_v

          x = fluid.layers.fill_constant(shape=[3, 3], value=0, dtype='float64')

          x_v = func(x)
          exe = fluid.Executor(fluid.CPUPlace())
          out = exe.run(fetch_list=[x_v])
          print(out[0])
          # [[1. 1. 1.]
          #  [1. 1. 1.]
          #  [1. 1. 1.]]

    """

    # TODO: remove this decorator after we finalize training API
123 124
    def __impl__(*args, **kwargs):
        program_translator = ProgramTranslator()
125
        if in_dygraph_mode() or not program_translator.enable_to_static:
126
            logging_utils.warn(
127
                "The decorator 'dygraph_to_static_func' doesn't work in "
128
                "dygraph mode or set ProgramTranslator.enable to False. "
129 130 131 132
                "We will just return dygraph output.")
            return dygraph_func(*args, **kwargs)
        static_func = program_translator.get_func(dygraph_func)
        return static_func(*args, **kwargs)
133 134 135 136

    return __impl__


137
dygraph_to_static_func = wrap_decorator(_dygraph_to_static_func_)
138

139

140 141 142 143 144 145
def copy_decorator_attrs(original_func, decorated_obj):
    """
    Copies some necessary attributes from original function into decorated function.

    Args:
        original_func(callable): the original decorated function.
146
        decorated_obj(StaticFunction): the target decorated StaticFunction object.
147 148 149 150 151 152 153 154 155 156 157 158 159 160
    """
    decorator_name = "declarative"

    decorated_obj.__name__ = original_func.__name__
    decorated_obj._decorator_name = decorator_name
    decorated_obj.__wrapped__ = original_func
    decorated_obj.__doc__ = original_func.__doc__
    if hasattr(original_func, "__module__"):
        decorated_obj.__module__ = original_func.__module__

    return decorated_obj


def declarative(function=None, input_spec=None):
161 162 163
    """
    Converts imperative dygraph APIs into declarative function APIs. Decorator
    @declarative handles the Program and Executor of static mode and returns
164 165 166 167
    the result as dygraph Tensor(s). Users could use the returned dygraph
    Tensor(s) to do imperative training, inference, or other operations. If the
    decorated function calls other imperative function, the called one will be
    converted into declarative function as well.
168

169
    Args:
170
        function (callable): callable imperative function.
171
        input_spec(list[InputSpec]|tuple[InputSpec]): list/tuple of InputSpec to specific the shape/dtype/name
172
            information of each input Tensor.
173

174
    Returns:
175
        Tensor(s): containing the numerical result.
176

177 178
    Examples:
        .. code-block:: python
179

180 181 182 183 184 185 186 187 188 189 190 191 192 193
            import paddle
            from paddle.jit import to_static

            @to_static
            def func(x):
                if paddle.mean(x) < 0:
                    x_v = x - 1
                else:
                    x_v = x + 1
                return x_v

            x = paddle.ones([1, 2], dtype='float32')
            x_v = func(x)
            print(x_v) # [[2. 2.]]
194

195
    """
196

197 198
    def decorated(python_func):
        """
199
        Decorates a python function into a StaticFunction object.
200 201 202
        """
        # Step 1. unwrap the function if it is already decorated.
        _, python_func = unwrap_decorators(python_func)
203

204 205 206
        # Step 2. copy some attributes from original python function.
        static_layer = copy_decorator_attrs(
            original_func=python_func,
207
            decorated_obj=StaticFunction(
208 209 210
                function=python_func, input_spec=input_spec))

        return static_layer
211

212 213
    # for usage: `declarative(foo, ...)`
    if function is not None:
214
        if isinstance(function, Layer):
215
            if isinstance(function.forward, StaticFunction):
216
                class_name = function.__class__.__name__
217
                logging_utils.warn(
218 219 220 221 222 223
                    "`{}.forward` has already been decorated somewhere. It will be redecorated to replace previous one.".
                    format(class_name))
            function.forward = decorated(function.forward)
            return function
        else:
            return decorated(function)
224

225 226
    # for usage: `@declarative`
    return decorated
227 228


229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268
def not_to_static(func=None):
    """
    A Decorator to suppresses the convertion of a function.

    Args:
        func(callable): The function to decorate.

    Returns:
        callable: A function which won't be converted in Dynamic-to-Static.

    Examples:
        .. code-block:: python

            import paddle

            @paddle.jit.not_to_static
            def func_not_to_static(x):
                res = x - 1
                return res

            @paddle.jit.to_static
            def func(x):
                if paddle.mean(x) < 0:
                    out = func_not_to_static(x)
                else:
                    out = x + 1
                return out

            x = paddle.ones([1, 2], dtype='float32')
            out = func(x)
            print(out) # [[2. 2.]]
    """
    if func is None:
        return not_to_static

    options = ConversionOptions(not_convert=True)
    setattr(func, CONVERSION_OPTIONS, options)
    return func


269
class _SaveLoadConfig(object):
270 271 272 273 274
    def __init__(self):
        self._output_spec = None
        self._model_filename = None
        self._params_filename = None
        self._separate_params = False
275 276
        # used for `paddle.load`
        self._keep_name_table = False
277 278 279 280

        # NOTE: Users rarely use following configs, so these configs are not open to users,
        # reducing user learning costs, but we retain the configuration capabilities

281 282
        # If True, programs are modified to only support direct inference deployment.
        # Otherwise,more information will be stored for flexible optimization and re-training.
283 284 285 286 287 288 289 290 291 292 293 294
        # Currently, only True is supported
        self._export_for_deployment = True

        # If True, It will save inference program only, and do not save params of Program
        self._program_only = False

    @property
    def output_spec(self):
        return self._output_spec

    @output_spec.setter
    def output_spec(self, spec):
295 296
        if spec is None:
            return
297 298
        if not isinstance(spec, list):
            raise TypeError(
299
                "The config `output_spec` should be 'list', but received input type is %s."
300 301 302 303
                % type(input))
            for var in spec:
                if not isinstance(var, core.VarBase):
                    raise TypeError(
304
                        "The element in config `output_spec` list should be 'Variable', but received element's type is %s."
305 306 307 308 309 310 311 312 313
                        % type(var))
        self._output_spec = spec

    @property
    def model_filename(self):
        return self._model_filename

    @model_filename.setter
    def model_filename(self, filename):
314 315
        if filename is None:
            return
316 317
        if not isinstance(filename, six.string_types):
            raise TypeError(
318
                "The config `model_filename` should be str, but received input's type is %s."
319 320
                % type(filename))
        if len(filename) == 0:
321
            raise ValueError("The config `model_filename` is empty string.")
322 323 324 325 326 327 328 329
        self._model_filename = filename

    @property
    def params_filename(self):
        return self._params_filename

    @params_filename.setter
    def params_filename(self, filename):
330 331
        if filename is None:
            return
332 333
        if not isinstance(filename, six.string_types):
            raise TypeError(
334
                "The config `params_filename` should be str, but received input's type is %s."
335 336
                % type(filename))
        if len(filename) == 0:
337
            raise ValueError("The config `params_filename` is empty string.")
338 339
        self._params_filename = filename

340 341 342 343 344 345
    @property
    def keep_name_table(self):
        return self._keep_name_table

    @keep_name_table.setter
    def keep_name_table(self, value):
346 347
        if value is None:
            return
348 349
        if not isinstance(value, bool):
            raise TypeError(
350
                "The config `keep_name_table` should be bool value, but received input's type is %s."
351 352 353
                % type(value))
        self._keep_name_table = value

354

355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389
def _parse_save_configs(configs):
    supported_configs = ['output_spec']

    # input check
    for key in configs:
        if key not in supported_configs:
            raise ValueError(
                "The additional config (%s) of `paddle.jit.save` is not supported."
                % (key))

    # construct inner config
    inner_config = _SaveLoadConfig()
    inner_config.output_spec = configs.get('output_spec', None)

    return inner_config


def _parse_load_config(configs):
    supported_configs = ['model_filename', 'params_filename']

    # input check
    for key in configs:
        if key not in supported_configs:
            raise ValueError(
                "The additional config (%s) of `paddle.jit.load` is not supported."
                % (key))

    # construct inner config
    inner_config = _SaveLoadConfig()
    inner_config.model_filename = configs.get('model_filename', None)
    inner_config.params_filename = configs.get('params_filename', None)

    return inner_config


390 391 392 393 394 395 396 397 398 399
def _get_input_var_names(inputs, input_spec):
    name_none_error = "The %s's name is None. " \
        "When using jit.save, please set InputSepc's name in " \
        "to_static(input_spec=[]) and jit.save(input_spec=[]) " \
        "and make sure they are consistent."
    name_no_exists_error = "The tensor `%s` does not exists. " \
        "Please make sure the name of InputSpec or example Tensor " \
        "in input_spec is the same as the name of InputSpec in " \
        "`to_static` decorated on the Layer.forward method."
    result_list = []
400 401 402
    input_var_names = [
        var.name for var in flatten(inputs) if isinstance(var, Variable)
    ]
403 404 405 406 407 408
    if input_spec is None:
        # no prune
        result_list = input_var_names
    elif input_spec is not None and len(input_spec) == len(input_var_names):
        # no prune
        result_list = input_var_names
409
        # if input spec name not in input_var_names, only raise warning
410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439
        for spec in input_spec:
            if spec.name is None:
                warnings.warn(name_none_error % spec)
            elif spec.name not in input_var_names:
                warnings.warn(name_no_exists_error % spec.name)
            else:
                # do nothing
                pass
    else:
        # prune
        for spec in input_spec:
            if spec.name is None:
                # name is None, the input_spec only can be InputSpec
                raise ValueError(name_none_error % spec)
            elif spec.name not in input_var_names:
                # the input_spec can be `InputSpec` or `VarBase`
                raise ValueError(name_no_exists_error % spec.name)
            else:
                result_list.append(spec.name)

    return result_list


def _get_output_vars(outputs, output_spec):
    name_no_exists_error = "The tensor `%s` does not exists. " \
        "Please make sure the name of example Tensor " \
        "in configs.output_spec is the output tensor of " \
        "Layer.forward method."
    result_list = []
    output_vars_dict = OrderedDict()
440
    for var in flatten(outputs):
441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458
        if isinstance(var, Variable):
            output_vars_dict[var.name] = var
    if output_spec is None:
        result_list = output_vars_dict.values()
    elif output_spec is not None and len(output_spec) == len(output_vars_dict):
        result_list = output_vars_dict.values()
        for var in output_spec:
            if var.name not in output_vars_dict:
                warnings.warn(name_no_exists_error % var.name)
    else:
        for var in output_spec:
            if var.name not in output_vars_dict:
                raise ValueError(name_no_exists_error % var.name)
            else:
                result_list.append(output_vars_dict[var.name])
    return result_list


459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501
# NOTE(chenweihang): [ Handling of use cases of API paddle.jit.load ]
# `paddle.jit.load` may be used to load saved results of:
# 1. Expected cases:
#   - paddle.jit.save
#   - paddle.static.save_inference_model
#   - paddle.fluid.io.save_inference_model
# 2. Error cases:
#   - paddle.save: no .pdmodel for prefix
#   - paddle.static.save: no .pdiparams but .pdparams exists
#   - paddle.fluid.io.save_params/save_persistables: no __model__
# TODO(chenweihang): polish error message in above error cases
def _build_load_path_and_config(path, config):
    # NOTE(chenweihang): If both [prefix save format] and [directory save format] exist,
    # raise error, avoid confusing behavior
    prefix_format_path = path + INFER_MODEL_SUFFIX
    prefix_format_exist = os.path.exists(prefix_format_path)
    directory_format_exist = os.path.isdir(path)
    if prefix_format_exist and directory_format_exist:
        raise ValueError(
            "The %s.pdmodel and %s directory exist at the same time, "
            "don't know which one to load, please make sure that the specified target "
            "of ``path`` is unique." % (path, path))
    elif not prefix_format_exist and not directory_format_exist:
        raise ValueError("The ``path`` (%s) to load model not exists." % path)
    else:
        if prefix_format_exist:
            file_prefix = os.path.basename(path)
            model_path = os.path.dirname(path)
            if config.model_filename is not None:
                warnings.warn(
                    "When loading the result saved with the "
                    "specified file prefix, the ``model_filename`` config does "
                    "not take effect.")
            config.model_filename = file_prefix + INFER_MODEL_SUFFIX
            if config.params_filename is not None:
                warnings.warn(
                    "When loading the result saved with the "
                    "specified file prefix, the ``params_filename`` config does "
                    "not take effect.")
            config.params_filename = file_prefix + INFER_PARAMS_SUFFIX
        else:
            # Compatible with the old save_inference_model format
            model_path = path
502

503
    return model_path, config
504 505


506
@switch_to_static_graph
507
def save(layer, path, input_spec=None, **configs):
508
    """
509
    Saves input Layer as ``paddle.jit.TranslatedLayer``
510 511
    format model, which can be used for inference or fine-tuning after loading.

512
    It will save the translated program and all related persistable
513
    variables of input Layer to given ``path`` .
514 515

    ``path`` is the prefix of saved objects, and the saved translated program file
516
    suffix is ``.pdmodel`` , the saved persistable variables file suffix is ``.pdiparams`` ,
517
    and here also saved some additional variable description information to a file,
518
    its suffix is ``.pdiparams.info``, these additional information is used in fine-tuning.
519 520

    The saved model can be loaded by follow APIs:
521 522
      - ``paddle.jit.load``
      - ``paddle.static.load_inference_model``
523 524 525
      - Other C++ inference APIs

    Args:
526
        layer (Layer): The Layer to be saved.
527
        path (str): The path prefix to save model. The format is ``dirname/file_prefix`` or ``file_prefix``.
528
        input_spec (list[InputSpec|Tensor]|tuple[InputSpec|Tensor], optional): Describes the input of the saved model's forward
529
            method, which can be described by InputSpec or example Tensor. If None, all input variables of
530
            the original Layer's forward method would be the inputs of the saved model. Default None.
531 532
        **configs (dict, optional): Other save configuration options for compatibility. We do not
            recommend using these configurations, they may be removed in the future. If not necessary,
533 534 535
            DO NOT use them. Default None.
            The following options are currently supported:
            (1) output_spec (list[Tensor]): Selects the output targets of the saved model.
536 537 538
            By default, all return variables of original Layer's forward method are kept as the
            output of the saved model. If the provided ``output_spec`` list is not all output variables,
            the saved model will be pruned according to the given ``output_spec`` list.
539

540 541 542 543 544 545 546
    Returns:
        None

    Examples:
        .. code-block:: python

            import numpy as np
547 548 549
            import paddle
            import paddle.nn as nn
            import paddle.optimizer as opt
550

551 552 553
            BATCH_SIZE = 16
            BATCH_NUM = 4
            EPOCH_NUM = 4
554

555 556 557 558 559 560 561
            IMAGE_SIZE = 784
            CLASS_NUM = 10

            # define a random dataset
            class RandomDataset(paddle.io.Dataset):
                def __init__(self, num_samples):
                    self.num_samples = num_samples
562

563 564 565 566
                def __getitem__(self, idx):
                    image = np.random.random([IMAGE_SIZE]).astype('float32')
                    label = np.random.randint(0, CLASS_NUM - 1, (1, )).astype('int64')
                    return image, label
567

568 569
                def __len__(self):
                    return self.num_samples
570

571 572
            class LinearNet(nn.Layer):
                def __init__(self):
573
                    super(LinearNet, self).__init__()
574
                    self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
575

576
                @paddle.jit.to_static
577 578 579
                def forward(self, x):
                    return self._linear(x)

580 581 582 583 584 585 586 587 588 589 590 591
            def train(layer, loader, loss_fn, opt):
                for epoch_id in range(EPOCH_NUM):
                    for batch_id, (image, label) in enumerate(loader()):
                        out = layer(image)
                        loss = loss_fn(out, label)
                        loss.backward()
                        opt.step()
                        opt.clear_grad()
                        print("Epoch {} batch {}: loss = {}".format(
                            epoch_id, batch_id, np.mean(loss.numpy())))

            # 1. train & save model.
592

593 594 595 596
            # create network
            layer = LinearNet()
            loss_fn = nn.CrossEntropyLoss()
            adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())
597

598 599 600 601 602 603 604
            # create data loader
            dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
            loader = paddle.io.DataLoader(dataset,
                batch_size=BATCH_SIZE,
                shuffle=True,
                drop_last=True,
                num_workers=2)
605

606 607
            # train
            train(layer, loader, loss_fn, adam)
608

609
            # save
610 611
            path = "example_model/linear"
            paddle.jit.save(layer, path)
612 613
    """

614
    # 1. input build & check
615
    prog_translator = ProgramTranslator()
616
    if not prog_translator.enable_to_static:
617
        raise RuntimeError(
618
            "The paddle.jit.save doesn't work when setting ProgramTranslator.enable to False."
619 620 621
        )
    if not isinstance(layer, Layer):
        raise TypeError(
622
            "The input layer of paddle.jit.save should be 'Layer', but received layer type is %s."
623 624
            % type(layer))

625 626
    # NOTE(chenweihang): If the input layer be wrapped by DataParallel,
    # the args and kwargs of forward method will can't be parsed by
627
    # function_spec, so here we save DataParallel._layers instead
628 629 630 631 632 633 634
    # DataParallel it self
    # NOTE(chenweihang): using inner_layer, do not change input layer
    if isinstance(layer, paddle.DataParallel):
        inner_layer = layer._layers
    else:
        inner_layer = layer

635 636 637 638 639 640 641 642 643 644 645
    # path check
    file_prefix = os.path.basename(path)
    if file_prefix == "":
        raise ValueError(
            "The input path MUST be format of dirname/file_prefix "
            "[dirname\\file_prefix in Windows system], but received "
            "file_prefix is empty string.")

    dirname = os.path.dirname(path)
    if dirname and not os.path.exists(dirname):
        os.makedirs(dirname)
646

647 648
    # avoid change user given input_spec
    inner_input_spec = None
649
    if input_spec is not None:
650 651
        for attr_func in dir(inner_layer):
            static_func = getattr(inner_layer, attr_func, None)
652 653 654 655 656
            if isinstance(static_func,
                          StaticFunction) and 'forward' != attr_func:
                raise ValueError(
                    "If there are static functions other than 'forward' that need to be saved, the input 'input_spec' should be None, but received the type of 'input_spec' is %s."
                    % type(input_spec))
657
        if not isinstance(input_spec, (list, tuple)):
658 659 660
            raise TypeError(
                "The input input_spec should be 'list', but received input_spec's type is %s."
                % type(input_spec))
661
        inner_input_spec = []
662
        for var in flatten(input_spec):
663 664 665 666 667 668
            if isinstance(var, paddle.static.InputSpec):
                inner_input_spec.append(var)
            elif isinstance(var, (core.VarBase, Variable)):
                inner_input_spec.append(
                    paddle.static.InputSpec.from_tensor(var))
            else:
669
                raise TypeError(
670
                    "The element in input_spec list should be 'Variable' or `paddle.static.InputSpec`, but received element's type is %s."
671 672
                    % type(var))

673 674
    # parse configs
    configs = _parse_save_configs(configs)
675 676
    scope = core.Scope()
    extra_var_info = dict()
677 678
    for attr_func in dir(inner_layer):
        static_func = getattr(inner_layer, attr_func, None)
679
        if isinstance(static_func, StaticFunction):
680 681
            concrete_program = static_func.concrete_program_specify_input_spec(
                inner_input_spec)
682 683
        elif 'forward' == attr_func:
            # transform in jit.save, if input_spec is incomplete, declarative will throw error
684 685 686 687 688
            # inner_input_spec is list[InputSpec], it should be packed with same sturcture
            # as original input_spec here.
            if inner_input_spec:
                inner_input_spec = pack_sequence_as(input_spec,
                                                    inner_input_spec)
689
            static_forward = declarative(
690
                inner_layer.forward, input_spec=inner_input_spec)
691
            concrete_program = static_forward.concrete_program
692
            # the input_spec has been used in declarative, which is equal to
693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711
            # @declarative with input_spec and jit.save without input_spec,
            # avoid needless warning
            inner_input_spec = None
        else:
            continue

        # 3. build input & output of save_infernece_model
        # NOTE(chenweihang): [ Get input variables name ]
        # There are two cases, whether to prune the inputs or not
        # - not prune inputs (recommend):
        #   - the len(input_spec) == len((concrete_program.inputs) - 1
        #   - here can use concrete_program.inputs directly
        # - prune inputs:
        #   - the input_spec length < len((concrete_program.inputs) - 1
        #   - the input_spec's name should be in concrete_program.inputs
        input_var_names = _get_input_var_names(concrete_program.inputs,
                                               inner_input_spec)

        # NOTE(chenweihang): [ Get output variables ]
712 713
        # the rule is like [ Get input variables name ]. For output var,
        # we only support VarBase spec, and actually, we only need the
714 715 716 717 718 719
        # var name of output, and we don't recommended to use output_spec
        output_vars = _get_output_vars(concrete_program.outputs,
                                       configs.output_spec)

        # NOTE(chenweihang): we maintain the mapping of variable name to
        # structured name, the buffer variable (non-persistable)
720
        # saved to inference program may not need by dygraph Layer,
721 722
        # we only record the state_dict variable's structured name
        state_names_dict = dict()
723
        for structured_name, var in six.iteritems(inner_layer.state_dict()):
724 725
            state_names_dict[var.name] = structured_name

726
        # 4. share parameters from Layer to scope & record var info
727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749
        for param_or_buffer in concrete_program.parameters:
            # share to scope
            param_or_buffer_tensor = scope.var(param_or_buffer.name).get_tensor(
            )
            src_tensor = param_or_buffer.value().get_tensor()
            param_or_buffer_tensor._share_data_with(src_tensor)
            # record var info
            if param_or_buffer.name not in extra_var_info:
                extra_info_dict = dict()
                if param_or_buffer.name in state_names_dict:
                    extra_info_dict['structured_name'] = state_names_dict[
                        param_or_buffer.name]
                extra_info_dict['stop_gradient'] = param_or_buffer.stop_gradient
                if isinstance(param_or_buffer, ParamBase):
                    extra_info_dict['trainable'] = param_or_buffer.trainable
                extra_var_info[param_or_buffer.name] = extra_info_dict

        # 5. save inference model
        from paddle.fluid.io import save_inference_model

        # construct new save_inference_model arguments
        model_path = dirname
        # NOTE(chenweihang): because prefix contains model and params filename,
750
        # so we don't support set model_filename & params_filename
751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776
        if 'forward' == attr_func:
            model_filename = file_prefix + INFER_MODEL_SUFFIX
            params_filename = file_prefix + INFER_PARAMS_SUFFIX
        else:
            model_filename = file_prefix + '.' + attr_func + INFER_MODEL_SUFFIX
            params_filename = file_prefix + '.' + attr_func + INFER_PARAMS_SUFFIX

        with scope_guard(scope):
            save_inference_model(
                dirname=model_path,
                feeded_var_names=input_var_names,
                target_vars=output_vars,
                executor=Executor(_current_expected_place()),
                main_program=concrete_program.main_program.clone(),
                model_filename=model_filename,
                params_filename=params_filename,
                export_for_deployment=configs._export_for_deployment,
                program_only=configs._program_only)

    # NOTE(chenweihang): [ Save extra variable info ]
    # save_inference_model will lose some important variable information, including:
    #   - Variable name and correspondence (when saved variables as one file)
    #   - Variable.stop_gradient information
    #   - Which persistent variable are parameter and which are not
    #   - Parameter.trainable information
    #
777 778
    # The lost information cannot be recovered when it is loaded again,
    # so if we want to perform fine-tune after loading, we may need to
779 780
    # configure redundant information to proceed.
    #
781 782
    # Due to compatibility issues, we cannot change the original storage structure,
    # but we can save these information in `jit.save` without changing the original
783 784
    # storage to improve user experience. So we save extra information into
    # file `***.pdiparams.info`
785
    with scope_guard(scope):
786
        extra_var_info_path = path + INFER_PARAMS_INFO_SUFFIX
787 788 789 790 791
        with open(extra_var_info_path, 'wb') as f:
            pickle.dump(extra_var_info, f, protocol=2)


@dygraph_only
792
def load(path, **configs):
793 794 795
    """
    :api_attr: imperative

796 797
    Load model saved by ``paddle.jit.save`` or ``paddle.static.save_inference_model`` or
    paddle 1.x API ``paddle.fluid.io.save_inference_model`` as ``paddle.jit.TranslatedLayer``,
798
    then performing inference or fine-tune training.
799 800

    .. note::
801
        If you load model saved by ``paddle.static.save_inference_model`` ,
802 803
        there will be the following limitations when using it in fine-tuning:
        1. Imperative mode do not support LoDTensor. All original model's feed targets or parametars that depend on LoD are temporarily unavailable.
804
        2. All saved model's feed targets need to be passed into TranslatedLayer's forward function.
805 806 807 808
        3. The variable's ``stop_gradient`` information is lost and can not be recovered.
        4. The parameter's ``trainable`` information is lost and can not be recovered.

    Args:
809
        path (str): The path prefix to load model. The format is ``dirname/file_prefix`` or ``file_prefix`` .
810 811
        **configs (dict, optional): Other load configuration options for compatibility. We do not
            recommend using these configurations, they may be removed in the future. If not necessary,
812 813
            DO NOT use them. Default None.
            The following options are currently supported:
814 815 816 817
            (1) model_filename (str): The inference model file name of the paddle 1.x
            ``save_inference_model`` save format. Default file name is :code:`__model__` .
            (2) params_filename (str): The persistable variables file name of the paddle 1.x
            ``save_inference_model`` save format. No default file name, save variables separately
818 819
            by default.

820 821 822 823 824

    Returns:
        TranslatedLayer: A Layer object can run saved translated model.

    Examples:
825
        1. Load model saved by ``paddle.jit.save`` then performing inference and fine-tune training.
826 827 828 829

        .. code-block:: python

            import numpy as np
830 831 832
            import paddle
            import paddle.nn as nn
            import paddle.optimizer as opt
833

834 835 836
            BATCH_SIZE = 16
            BATCH_NUM = 4
            EPOCH_NUM = 4
837

838 839
            IMAGE_SIZE = 784
            CLASS_NUM = 10
840

841 842 843 844
            # define a random dataset
            class RandomDataset(paddle.io.Dataset):
                def __init__(self, num_samples):
                    self.num_samples = num_samples
845

846 847 848 849
                def __getitem__(self, idx):
                    image = np.random.random([IMAGE_SIZE]).astype('float32')
                    label = np.random.randint(0, CLASS_NUM - 1, (1, )).astype('int64')
                    return image, label
850

851 852 853 854 855
                def __len__(self):
                    return self.num_samples

            class LinearNet(nn.Layer):
                def __init__(self):
856
                    super(LinearNet, self).__init__()
857
                    self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
858

859
                @paddle.jit.to_static
860 861 862
                def forward(self, x):
                    return self._linear(x)

863 864 865 866 867 868 869 870 871 872 873
            def train(layer, loader, loss_fn, opt):
                for epoch_id in range(EPOCH_NUM):
                    for batch_id, (image, label) in enumerate(loader()):
                        out = layer(image)
                        loss = loss_fn(out, label)
                        loss.backward()
                        opt.step()
                        opt.clear_grad()
                        print("Epoch {} batch {}: loss = {}".format(
                            epoch_id, batch_id, np.mean(loss.numpy())))

874
            # 1. train & save model.
875

876
            # create network
877 878 879 880
            layer = LinearNet()
            loss_fn = nn.CrossEntropyLoss()
            adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())

881
            # create data loader
882 883 884 885 886 887
            dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
            loader = paddle.io.DataLoader(dataset,
                batch_size=BATCH_SIZE,
                shuffle=True,
                drop_last=True,
                num_workers=2)
888

889 890
            # train
            train(layer, loader, loss_fn, adam)
891

892
            # save
893 894
            path = "example_model/linear"
            paddle.jit.save(layer, path)
895

896
            # 2. load model
897

898
            # load
899
            loaded_layer = paddle.jit.load(path)
900 901

            # inference
902 903 904
            loaded_layer.eval()
            x = paddle.randn([1, IMAGE_SIZE], 'float32')
            pred = loaded_layer(x)
905 906

            # fine-tune
907 908 909
            loaded_layer.train()
            adam = opt.Adam(learning_rate=0.001, parameters=loaded_layer.parameters())
            train(loaded_layer, loader, loss_fn, adam)
910 911


912
        2. Load model saved by ``paddle.fluid.io.save_inference_model`` then performing and fine-tune training.
913 914 915 916

        .. code-block:: python

            import numpy as np
917
            import paddle
918
            import paddle.static as static
919 920
            import paddle.nn as nn
            import paddle.optimizer as opt
921
            import paddle.nn.functional as F
922

923 924 925
            BATCH_SIZE = 16
            BATCH_NUM = 4
            EPOCH_NUM = 4
926

927 928 929 930 931 932 933
            IMAGE_SIZE = 784
            CLASS_NUM = 10

            # define a random dataset
            class RandomDataset(paddle.io.Dataset):
                def __init__(self, num_samples):
                    self.num_samples = num_samples
934

935 936 937 938
                def __getitem__(self, idx):
                    image = np.random.random([IMAGE_SIZE]).astype('float32')
                    label = np.random.randint(0, CLASS_NUM - 1, (1, )).astype('int64')
                    return image, label
939

940 941
                def __len__(self):
                    return self.num_samples
942

943 944
            paddle.enable_static()

945 946
            image = static.data(name='image', shape=[None, 784], dtype='float32')
            label = static.data(name='label', shape=[None, 1], dtype='int64')
947
            pred = static.nn.fc(x=image, size=10, activation='softmax')
948 949
            loss = F.cross_entropy(input=pred, label=label)
            avg_loss = paddle.mean(loss)
950

951
            optimizer = paddle.optimizer.SGD(learning_rate=0.001)
952 953
            optimizer.minimize(avg_loss)

954 955 956
            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())
957

958 959 960 961 962
            # create data loader
            dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
            loader = paddle.io.DataLoader(dataset,
                feed_list=[image, label],
                places=place,
963
                batch_size=BATCH_SIZE,
964 965 966
                shuffle=True,
                drop_last=True,
                num_workers=2)
967 968 969 970

            # 1. train and save inference model
            for data in loader():
                exe.run(
971
                    static.default_main_program(),
972
                    feed=data,
973 974 975
                    fetch_list=[avg_loss])

            model_path = "fc.example.model"
976
            paddle.fluid.io.save_inference_model(
977 978 979
                model_path, ["image"], [pred], exe)

            # 2. load model
980 981

            # enable dygraph mode
982 983 984 985
            paddle.disable_static(place)

            # load
            fc = paddle.jit.load(model_path)
986

987 988 989
            # inference
            fc.eval()
            x = paddle.randn([1, IMAGE_SIZE], 'float32')
990 991
            pred = fc(x)

992
            # fine-tune
993
            fc.train()
994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010
            loss_fn = nn.CrossEntropyLoss()
            adam = opt.Adam(learning_rate=0.001, parameters=fc.parameters())
            loader = paddle.io.DataLoader(dataset,
                places=place,
                batch_size=BATCH_SIZE,
                shuffle=True,
                drop_last=True,
                num_workers=2)
            for epoch_id in range(EPOCH_NUM):
                for batch_id, (image, label) in enumerate(loader()):
                    out = fc(image)
                    loss = loss_fn(out, label)
                    loss.backward()
                    adam.step()
                    adam.clear_grad()
                    print("Epoch {} batch {}: loss = {}".format(
                        epoch_id, batch_id, np.mean(loss.numpy())))
1011
    """
1012 1013 1014 1015
    # 1. construct correct config
    config = _parse_load_config(configs)
    model_path, config = _build_load_path_and_config(path, config)

1016
    return TranslatedLayer._construct(model_path, config)
1017 1018


1019
@dygraph_only
Z
Zeng Jinle 已提交
1020 1021 1022 1023 1024
def _trace(layer,
           inputs,
           feed_prefix='feed_',
           fetch_prefix='fetch_',
           tmp_prefix='t_'):
1025
    assert isinstance(layer, Layer)
1026 1027 1028 1029 1030 1031 1032 1033 1034

    if not isinstance(inputs, (list, tuple)):
        inputs = [inputs]

    tracer = _dygraph_tracer()._get_program_desc_tracer()

    var_list = extract_vars(inputs)

    with program_desc_tracing_guard(True):
1035
        original_outputs = layer(*inputs)
1036 1037 1038 1039
        if not isinstance(original_outputs, (list, tuple)):
            outputs = [original_outputs]
        else:
            outputs = original_outputs
1040
        out_vars = extract_vars(outputs, err_tag='outputs')
1041

1042
        program_desc, feed_names, fetch_names, parameters = tracer.create_program_desc(
Z
Zeng Jinle 已提交
1043
            var_list, feed_prefix, out_vars, fetch_prefix, tmp_prefix)
1044 1045 1046 1047 1048
        tracer.reset()

    with _dygraph_guard(None):
        program = create_program_from_desc(program_desc)

1049
    return original_outputs, program, feed_names, fetch_names, parameters
1050 1051 1052 1053


class TracedLayer(object):
    """
1054
    :api_attr: imperative
1055

1056 1057 1058 1059 1060
    TracedLayer is used to convert a forward dygraph model to a static
    graph model. This is mainly used to save the dygraph model for online
    inference using C++. Besides, users can also do inference in Python
    using the converted static graph model, which usually has better
    performance than the original dygraph model.
1061 1062 1063 1064

    TracedLayer would run the static graph model using :code:`Executor`
    and :code:`CompiledProgram` . The static graph model would share
    parameters with the dygraph model.
1065 1066

    All TracedLayer objects should not be created by constructor and should
1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077
    be created by static method :code:`TracedLayer.trace(layer, inputs)` .

    The TracedLayer can only be used to convert the data-independent dygraph
    model into the static graph model, which means the dygraph model should
    be independent with the tensor data and shape.
    """

    def __init__(self, program, parameters, feed_names, fetch_names):
        self._program = program
        self._feed_names = feed_names
        self._fetch_names = fetch_names
1078
        self._params = parameters
1079 1080 1081 1082 1083

        self._place = _current_expected_place()

        self._scope = core.Scope()
        for p in parameters:
1084
            src_tensor = p.value().get_tensor()
1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107
            dst_tensor = self._scope.var(p.name).get_tensor()
            dst_tensor._share_data_with(src_tensor)

        self._exe = Executor(self._place)
        self._compiled_program = None
        self._build_strategy = None
        self._exec_strategy = None

    @property
    def program(self):
        return self._program

    def _switch(self, is_test=True):
        for block_id in range(self._program.num_blocks):
            block = self._program.block(block_id)
            for op in block.ops:
                if op.has_attr("is_test"):
                    op._set_attr("is_test", is_test)

    @staticmethod
    @dygraph_only
    def trace(layer, inputs):
        """
1108
        This method is the only allowed method to create TracedLayer object.
1109 1110 1111 1112
        It would call the :code:`layer(*inputs)` method to run the dygraph
        model and convert it into a static graph model.

        Args:
1113
            layer (paddle.nn.Layer): the layer object to be traced.
1114 1115
            inputs (list(Tensor)|tuple(Tensor)|Tensor): the input tensors of
                the layer object.
1116 1117

        Returns:
1118
            tuple: A tuple of 2 items, whose the first item is the output of
1119 1120
                :code:`layer(*inputs)` , and the second item is the created
                TracedLayer object.
1121

1122
        Examples:
1123 1124
            .. code-block:: python:

1125
                import paddle
1126

1127
                class ExampleLayer(paddle.nn.Layer):
1128 1129
                    def __init__(self):
                        super(ExampleLayer, self).__init__()
1130
                        self._fc = paddle.nn.Linear(3, 10)
1131 1132 1133 1134

                    def forward(self, input):
                        return self._fc(input)

1135

1136 1137 1138 1139 1140 1141
                layer = ExampleLayer()
                in_var = paddle.uniform(shape=[2, 3], dtype='float32')
                out_dygraph, static_layer = paddle.jit.TracedLayer.trace(layer, inputs=[in_var])

                # run the static graph model using Executor inside
                out_static_graph = static_layer([in_var])
1142

1143 1144
                print(len(out_static_graph)) # 1
                print(out_static_graph[0].shape) # (2, 10)
1145

1146 1147
                # save the static graph model for inference
                static_layer.save_inference_model(dirname='./saved_infer_model')
1148

1149
        """
1150 1151 1152 1153
        assert isinstance(
            layer, Layer
        ), "The type of 'layer' in fluid.dygraph.jit.TracedLayer.trace must be fluid.dygraph.Layer, but received {}.".format(
            type(layer))
1154 1155
        outs, prog, feed, fetch, parameters = _trace(layer, inputs)
        traced = TracedLayer(prog, parameters, feed, fetch)
1156 1157 1158 1159 1160 1161 1162
        return outs, traced

    def set_strategy(self, build_strategy=None, exec_strategy=None):
        """
        Set the strategies when running static graph model.

        Args:
1163
            build_strategy (BuildStrategy, optional): build strategy of
1164 1165 1166 1167 1168 1169 1170 1171 1172 1173
                :code:`CompiledProgram` inside TracedLayer. Default None.
            exec_strategy (ExecutionStrategy, optional): execution strategy of
                :code:`CompiledProgram` inside TracedLayer. Default None.

        Returns:
            None

        Examples:
            .. code-block:: python:

1174
                import paddle
1175

1176
                class ExampleLayer(paddle.nn.Layer):
1177 1178
                    def __init__(self):
                        super(ExampleLayer, self).__init__()
1179
                        self._fc = paddle.nn.Linear(3, 10)
1180 1181 1182 1183

                    def forward(self, input):
                        return self._fc(input)

1184 1185 1186 1187
                layer = ExampleLayer()
                in_var = paddle.uniform(shape=[2, 3], dtype='float32')

                out_dygraph, static_layer = paddle.jit.TracedLayer.trace(layer, inputs=[in_var])
1188

1189 1190
                build_strategy = paddle.static.BuildStrategy()
                build_strategy.enable_inplace = True
1191

1192 1193
                exec_strategy = paddle.static.ExecutionStrategy()
                exec_strategy.num_threads = 2
1194

1195 1196
                static_layer.set_strategy(build_strategy=build_strategy, exec_strategy=exec_strategy)
                out_static_graph = static_layer([in_var])
1197 1198 1199

        """
        assert self._compiled_program is None, "Cannot set strategy after run"
1200 1201 1202 1203 1204 1205 1206 1207
        assert isinstance(
            build_strategy, (type(None), BuildStrategy)
        ), "The type of 'build_strategy' in fluid.dygraph.jit.TracedLayer.set_strategy must be fluid.BuildStrategy, but received {}.".format(
            type(build_strategy))
        assert isinstance(
            exec_strategy, (type(None), ExecutionStrategy)
        ), "The type of 'exec_strategy' in fluid.dygraph.jit.TracedLayer.set_strategy must be fluid.ExecutionStrategy, but received {}.".format(
            type(exec_strategy))
1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225
        self._build_strategy = build_strategy
        self._exec_strategy = exec_strategy

    @switch_to_static_graph
    def _compile(self):
        self._compiled_program = CompiledProgram(
            self._program).with_data_parallel(
                build_strategy=self._build_strategy,
                exec_strategy=self._exec_strategy,
                places=self._place)

    def _build_feed(self, inputs):
        assert isinstance(inputs, (list, tuple)), \
            "Inputs should be a list or tuple of variables"
        assert len(inputs) == len(self._feed_names)
        feed_dict = {}
        if in_dygraph_mode():
            for x, name in zip(inputs, self._feed_names):
1226
                feed_dict[name] = x.value().get_tensor()
1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246
        else:
            for x, name in zip(inputs, self._feed_names):
                feed_dict[name] = x

        return feed_dict

    @switch_to_static_graph
    def _run(self, feed):
        return self._exe.run(self._compiled_program,
                             feed=feed,
                             fetch_list=self._fetch_names)

    def __call__(self, inputs):
        with scope_guard(self._scope):
            if self._compiled_program is None:
                self._compile()

            return self._run(self._build_feed(inputs))

    @switch_to_static_graph
1247
    def save_inference_model(self, path, feed=None, fetch=None):
1248
        """
1249 1250
        Save the TracedLayer to a model for inference. The saved
        inference model can be loaded by C++ inference APIs.
1251

1252 1253 1254
        ``path`` is the prefix of saved objects, and the saved translated program file
        suffix is ``.pdmodel`` , the saved persistable variables file suffix is ``.pdiparams`` .

1255
        Args:
1256
            path(str): The path prefix to save model. The format is ``dirname/file_prefix`` or ``file_prefix``.
1257
            feed (list[int], optional): the input variable indices of the saved
1258
                inference model. If None, all input variables of the
1259 1260 1261 1262 1263 1264 1265 1266
                TracedLayer object would be the inputs of the saved inference
                model. Default None.
            fetch (list[int], optional): the output variable indices of the
                saved inference model. If None, all output variables of the
                TracedLayer object would be the outputs of the saved inference
                model. Default None.

        Returns:
1267
            None
1268 1269 1270 1271 1272

        Examples:
            .. code-block:: python:

                import numpy as np
1273
                import paddle
1274

1275
                class ExampleLayer(paddle.nn.Layer):
1276 1277
                    def __init__(self):
                        super(ExampleLayer, self).__init__()
1278
                        self._fc = paddle.nn.Linear(3, 10)
1279 1280 1281 1282

                    def forward(self, input):
                        return self._fc(input)

1283 1284
                save_dirname = './saved_infer_model'
                in_np = np.random.random([2, 3]).astype('float32')
1285 1286
                in_var = paddle.to_tensor(in_np)
                layer = ExampleLayer()
1287

1288 1289
                out_dygraph, static_layer = paddle.jit.TracedLayer.trace(layer, inputs=[in_var])
                static_layer.save_inference_model(save_dirname, feed=[0], fetch=[0])
1290

1291 1292 1293 1294
                paddle.enable_static()
                place = paddle.CPUPlace()
                exe = paddle.static.Executor(place)
                program, feed_vars, fetch_vars = paddle.static.load_inference_model(save_dirname,
1295
                                                    exe)
1296 1297 1298

                fetch, = exe.run(program, feed={feed_vars[0]: in_np}, fetch_list=fetch_vars)
                print(fetch.shape) # (2, 10)
1299
        """
1300
        check_type(path, "path", str,
1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314
                   "fluid.dygraph.jit.TracedLayer.save_inference_model")
        check_type(feed, "feed", (type(None), list),
                   "fluid.dygraph.jit.TracedLayer.save_inference_model")
        if isinstance(feed, list):
            for f in feed:
                check_type(f, "each element of feed", int,
                           "fluid.dygraph.jit.TracedLayer.save_inference_model")
        check_type(fetch, "fetch", (type(None), list),
                   "fluid.dygraph.jit.TracedLayer.save_inference_model")
        if isinstance(fetch, list):
            for f in fetch:
                check_type(f, "each element of fetch", int,
                           "fluid.dygraph.jit.TracedLayer.save_inference_model")

1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326
        # path check
        file_prefix = os.path.basename(path)
        if file_prefix == "":
            raise ValueError(
                "The input path MUST be format of dirname/file_prefix "
                "[dirname\\file_prefix in Windows system], but received "
                "file_prefix is empty string.")

        dirname = os.path.dirname(path)
        if dirname and not os.path.exists(dirname):
            os.makedirs(dirname)

1327
        from paddle.fluid.io import save_inference_model
1328 1329 1330 1331 1332

        def get_feed_fetch(all_vars, partial_vars):
            if partial_vars is None:
                return all_vars

1333
            return [all_vars[idx] for idx in partial_vars]
1334 1335 1336 1337 1338 1339 1340 1341 1342 1343

        with scope_guard(self._scope):
            feeded_var_names = get_feed_fetch(self._feed_names, feed)
            target_var_names = get_feed_fetch(self._fetch_names, fetch)
            target_vars = []
            for name in target_var_names:
                target_var = self._program.global_block().vars.get(name, None)
                assert target_var is not None, "{} cannot be found".format(name)
                target_vars.append(target_var)

1344 1345 1346
            model_filename = file_prefix + INFER_MODEL_SUFFIX
            params_filename = file_prefix + INFER_PARAMS_SUFFIX

1347
            save_inference_model(
1348 1349 1350 1351
                dirname=dirname,
                feeded_var_names=feeded_var_names,
                target_vars=target_vars,
                executor=self._exe,
1352 1353 1354
                main_program=self._program.clone(),
                model_filename=model_filename,
                params_filename=params_filename)