program_translator.py 52.0 KB
Newer Older
1
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14
#
# 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
import collections
16
from paddle.utils import gast
17
import inspect
18 19
import textwrap
import threading
20
import weakref
21 22

from paddle.fluid import framework
J
Jiabin Yang 已提交
23
from paddle.fluid import _non_static_mode
24
from paddle.fluid.dygraph import layers
25 26
from paddle.fluid.data_feeder import check_type
from paddle.fluid.layers.utils import flatten
27
from paddle.fluid.dygraph.base import param_guard
28
from paddle.fluid.dygraph.base import switch_to_static_graph
29
from paddle.fluid.dygraph.dygraph_to_static import error
30
from paddle.fluid.dygraph.dygraph_to_static import logging_utils
31 32 33 34 35 36 37 38 39 40 41 42
from paddle.fluid.dygraph.dygraph_to_static.origin_info import (
    attach_origin_info,
)
from paddle.fluid.dygraph.dygraph_to_static.origin_info import (
    create_and_update_origin_info_map,
)
from paddle.fluid.dygraph.dygraph_to_static.origin_info import (
    update_op_callstack_with_origin_info,
)
from paddle.fluid.dygraph.dygraph_to_static.partial_program import (
    partial_program_from,
)
43
from paddle.fluid.dygraph.dygraph_to_static.utils import ast_to_func
44
from paddle.fluid.dygraph.dygraph_to_static.utils import ast_to_source_code
45
from paddle.fluid.dygraph.dygraph_to_static.utils import func_to_source_code
46
from paddle.fluid.dygraph.dygraph_to_static.utils import input_specs_compatible
47
from paddle.fluid.dygraph.dygraph_to_static.utils import type_name
48
from paddle.fluid.dygraph.dygraph_to_static.utils import unwrap
49 50 51 52 53 54 55 56 57 58 59 60
from paddle.fluid.dygraph.dygraph_to_static.utils import (
    make_hashable,
    ALREADY_D2S,
)
from paddle.fluid.dygraph.dygraph_to_static.function_spec import (
    FunctionSpec,
    _hash_spec_names,
)
from paddle.fluid.dygraph.dygraph_to_static.function_spec import (
    get_buffers,
    get_parameters,
)
61 62

from .ast_transformer import DygraphToStaticAst
63

64
__all__ = ['ProgramTranslator', 'convert_to_static']
65

66 67 68 69
# For each traced function, we set `max_traced_program_count` = 10 to consider caching performance.
# Once exceeding the threshold, we will raise warning to users to make sure the conversion is as expected.
MAX_TRACED_PROGRAM_COUNT = 10

70

71
class FunctionCache:
72 73 74 75 76
    """
    Caches the transformed functions to avoid redundant conversions of the same function.
    """

    def __init__(self):
77
        # Caches the converted static functions. {dygraph_func: static_func}
X
xiongkun 已提交
78
        self._converted_static_func_caches = weakref.WeakKeyDictionary()
79 80 81
        # Caches the converted ast node for same source code. {source_code: ast_root}
        self._code_to_ast_caches = dict()
        self._dygraph_to_static = DygraphToStaticAst()
82

83 84 85 86 87 88
    def convert_with_cache(self, func):
        """
        Returns the cached static function or converts it when first encounters the function.
        """
        # If hit cache, return it directly.
        static_func = self._converted_static_func_caches.get(func, None)
89 90

        if static_func is None:
91 92
            static_func = self._convert(func)
            self._converted_static_func_caches[func] = static_func
93 94 95

        return static_func

96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116
    def _convert(self, func):
        """
        Converts dygraph function into static function. For two functions with same dedent code,
        the second function will reuse the transformed ast node of previous one.

        For example:
            # A.py
            def foo(x, y):
                z = x + y
                return z

            # B.py
            def foo(x, y):
                z = x + y
                return z

        If the conversion of A.foo happens after B.foo, it will reuse the transformed ast node of B.foo
        to speed up the conversion.
        """
        # Note: In Python2, it will raise OSError when inspect function
        # with decorator directly and function.__wrapped__ holds the actual function.
117
        func = unwrap(func)
118
        source_code = func_to_source_code(func)
119 120 121 122 123

        # TODO(liym27):
        #  Consider this case: source_code in self._code_to_ast_caches,
        #  but actually they are methods in different classes.
        #  Maybe use (__class__, source_code) as key
124 125 126 127
        if source_code in self._code_to_ast_caches:
            root_wrapper = self._code_to_ast_caches[source_code]
        else:
            root = gast.parse(source_code)
128
            root = attach_origin_info(root, func)
129 130
            root_wrapper = self._dygraph_to_static.get_static_ast(root)
            self._code_to_ast_caches[source_code] = root_wrapper
131

132 133
        # Get static function from AST
        static_func, file_name = ast_to_func(root_wrapper.node, func)
134 135

        create_and_update_origin_info_map(root_wrapper.node, static_func)
136
        return static_func
137 138

    def exist(self, func):
139
        return func in self._converted_static_func_caches
140 141


142 143 144 145
_CACHE_LOCK = threading.Lock()
_FUNCTION_CACHE = FunctionCache()


146
def convert_to_static(function):
147
    """
148
    Transforms function of dygraph into static function using the cache mechanism.
149 150 151

    Args:
        function(callable): The function with dygraph layers that will be converted into static layers.
152
    """
153 154
    if getattr(function, ALREADY_D2S, None):
        return function
155
    with _CACHE_LOCK:
156
        static_func = _FUNCTION_CACHE.convert_with_cache(function)
157
        setattr(static_func, ALREADY_D2S, True)
158 159 160
        return static_func


161
class CacheKey:
162 163 164
    """
    Cached key for ProgramCache.
    """
165

166
    __slots__ = [
167 168 169 170 171 172
        'function_spec',
        'input_args_with_spec',
        'input_kwargs_with_spec',
        'class_instance',
        'kwargs',
        '_spec_names_id',
173
    ]
174

175 176 177 178 179 180 181 182
    def __init__(
        self,
        function_spec,
        input_args_with_spec,
        input_kwargs_with_spec,
        class_instance,
        **kwargs
    ):
183 184
        """
        Initializes a cache key.
185

186 187
        Args:
            functions_spec(FunctionSpec): a FunctionSpec instance of decorated function.
188 189
            input_args_with_spec(list[InputSpec]): actual input args with some arguments replaced by InputSpec.
            input_kwargs_with_spec(list[{string:InputSpec}]): actual input kwargs with some arguments replaced by InputSpec.
190
            class_instance(object): a instance of class `Layer`.
191
            **kwargs(dict): manage other arguments used for better scalability
192
        """
193
        self.function_spec = function_spec
194 195
        self.input_args_with_spec = input_args_with_spec
        self.input_kwargs_with_spec = input_kwargs_with_spec
196
        self.class_instance = class_instance
197 198
        # NOTE: `kwargs` is usually not considered as basic member for `__hash__`
        self.kwargs = kwargs
199 200 201
        self._spec_names_id = _hash_spec_names(
            input_args_with_spec, input_kwargs_with_spec
        )
202 203 204

    @classmethod
    def from_func_and_args(cls, function_spec, args, kwargs, class_instance):
205
        """
206 207 208 209 210 211 212 213 214 215 216
        Generated a CacheKey instance by given inputs.

        Args:
            functions_spec(FunctionSpec): a FunctionSpec instance of decorated function.
            args(tuple): tuple of actual inputs arguments.
            kwargs(dict): dict of actual inputs keyword arguments.
            class_instance(object): a instance of class `Layer`.
        """
        # 1. filter `self` in args
        if args and isinstance(args[0], layers.Layer):
            args = args[1:]
217
        # 2. convert tensor and numpy array into InputSpec
218
        _args, _kwargs = function_spec.unified_args_and_kwargs(args, kwargs)
219 220 221 222
        (
            input_args_with_spec,
            input_kwargs_with_spec,
        ) = function_spec.args_to_input_spec(_args, _kwargs)
223 224

        # 3. check whether hit the cache or build a new program for the input arguments
225 226 227 228 229 230
        return CacheKey(
            function_spec,
            input_args_with_spec,
            input_kwargs_with_spec,
            class_instance,
        )
231 232 233

    def __hash__(self):
        error_msg = "Arguments to a `@paddle.jit.to_static` must be a hashable Python objects (or nested structures of these types)."
234
        with_hook = self.kwargs.get("with_hook", False)
235
        is_train = self.kwargs.get("is_train", False)
236 237 238 239 240 241 242 243 244 245 246
        return hash(
            (
                id(self.function_spec),
                make_hashable(self.input_args_with_spec, error_msg),
                make_hashable(self.input_kwargs_with_spec, error_msg),
                self._spec_names_id,
                self.class_instance,
                with_hook,
                is_train,
            )
        )
247 248 249 250 251 252 253 254

    def __eq__(self, other):
        return (type(self) is type(other)) and hash(self) == hash(other)

    def __neq__(self, other):
        return not self == other

    def __repr__(self):
255
        return "id(function_spec): {}, input_args_with_spec: {}, input_kwargs_with_spec: {}, class_instance: {}".format(
256 257 258 259 260
            id(self.function_spec),
            self.input_args_with_spec,
            self.input_kwargs_with_spec,
            self.class_instance,
        )
261 262 263 264 265 266 267 268 269


def unwrap_decorators(func):
    """
    Unwraps a decorated function and returns the decorator list and inner target.
    """
    decorators = []
    cur = func
    while True:
270
        if isinstance(cur, StaticFunction):
271 272 273 274 275
            decorators.append(cur)
            # Note: if `cur` is a method, keep it as bound method of class.
            instance = cur._class_instance
            if instance is not None:
                cur = cur.dygraph_function.__get__(instance)
276
            else:
277 278 279 280
                cur = cur.dygraph_function
        else:
            break
    return decorators, cur
281

282

283
class StaticFunction:
284 285 286 287 288
    """
    Wrapper class to Manage program conversion of decorated function.

    """

289
    def __init__(self, function, input_spec=None, **kwargs):
290
        """
291
        Initializes a `StaticFunction`.
292 293 294 295

        Args:
            function(callable): A function or method that will be converted into static program.
            input_spec(list[InputSpec]): list of InputSpec to specify the `shape/dtype/name` information for each input argument, default None.
296
            **kwargs(dict): other arguments like `build_strategy` et.al.
297 298
        """
        # save the instance `self` while decorating a method of class.
299

300 301 302
        if inspect.ismethod(function):
            self._dygraph_function = getattr(function, '__func__')
            self._class_instance = getattr(function, '__self__')
303

304 305 306
            if not hasattr(self._class_instance, '_original_funcs'):
                raise TypeError(
                    "When using 'to_static' to convert method of a class, "
307 308
                    "please ensure the class inherits from nn.Layer"
                )
309
            self._class_instance._original_funcs[
310 311
                function.__name__
            ] = self._dygraph_function
312 313 314 315 316 317 318
        else:
            self._dygraph_function = function
            self._class_instance = None

        self._input_spec = input_spec
        self._function_spec = FunctionSpec(function, input_spec)
        self._program_cache = ProgramCache()
319
        self._descriptor_cache = weakref.WeakKeyDictionary()
320
        # Note: Hold a reference to ProgramTranslator for switching `enable_to_static`.
321
        self._program_trans = ProgramTranslator()
322
        self._kwargs = kwargs
323
        self._training = True
324 325
        self._cuda_graph_capture_mode = ""
        self._cuda_graph_pool_id = 0
326

327 328 329 330 331 332 333
        self._property = kwargs.get("property", False)

    @property
    def is_property(self):
        # whether is class proproty to be exported.
        return self._property

334
    def train(self):
335 336
        if (
            isinstance(self._class_instance, layers.Layer)
337
            and self._class_instance.training is False
338
        ):
339 340 341
            raise RuntimeError(
                "Failed to switch train mode. {} is a Layer's method, "
                "please use Layer.train() to switch train mode.".format(
342 343 344
                    self.dygraph_function
                )
            )
345 346 347
        self._training = True

    def eval(self):
348 349
        if (
            isinstance(self._class_instance, layers.Layer)
350
            and self._class_instance.training is True
351
        ):
352 353 354
            raise RuntimeError(
                "Failed to switch eval mode. {} is a Layer's method, "
                "please use Layer.eval() to switch eval mode.".format(
355 356 357
                    self.dygraph_function
                )
            )
358
        self._training = False
359 360 361 362 363 364

    def __get__(self, instance, owner):
        """
        Overrides this method to parse the class instance and call bound method correctly.

        For example:
365

366 367 368 369
            '''
            class Net(Layer):
                def __init__(self):
                    pass
370

371 372 373 374 375 376 377
                @paddle.jit.to_static
                def forward(self, x, y):
                    return x + y

            net = Net()
            out = net(x, y)
            '''
378

379 380
        In above case, `net(x, y)` will call `net.forward(x, y)` firstly that is a bound method
        of `Net` instance. After decorated by `@paddle.jit.to_static`, it will firstly to call `__get__`
381
        to parse the class instance correctly instead of the `StaticFunction` instance.
382
        """
383 384 385
        if instance not in self._descriptor_cache:
            if instance is None:
                return self
386
            # Note(Aurelius84): To construct new instance of StaticFunction when we
387 388 389 390 391 392 393 394
            # first encouter the bound function of layer and cache it.
            new_static_layer = self._clone()
            new_static_layer._class_instance = instance
            self._descriptor_cache[instance] = new_static_layer

        return self._descriptor_cache[instance]

    def _clone(self):
395 396 397
        return self.__class__(
            self._dygraph_function, self._input_spec, **self._kwargs
        )
398 399

    def __call__(self, *args, **kwargs):
400
        """
401 402 403 404
        Supports to call the returned instance with input `args` and `kwargs` directly.

        Args:
            *args(tuple): tuple of all input arguments from original decorated function.
405
            **kwargs(dict): dict of all input keyward arguments from original decorated function.
406 407 408

        Return:
            Outputs of decorated function.
409
        """
410 411
        if self._property:
            return self._call_dygraph_function(*args, **kwargs)
412

413
        # 1. call dygraph function directly if not enable `declarative`
414
        if not self._program_trans.enable_to_static:
415 416 417 418
            # NOTE(liym27):
            # Here calls `warnings.warn` but not `logging_utils.warn` because by default warnings.warn(message)
            # will show up **only once**. StaticFunction.__call__ will run many times, it is appropriate to
            # display this warning message only once.
419
            logging_utils.warn(
420 421
                "The decorator '@paddle.jit.to_static' does NOT work when setting ProgramTranslator.enable to False. "
                "We will just return dygraph output. If you would like to get static graph output, please call API "
422 423
                "ProgramTranslator.enable(True)"
            )
424 425
            return self._call_dygraph_function(*args, **kwargs)

J
Jiabin Yang 已提交
426
        if not _non_static_mode():
427 428
            raise RuntimeError(
                "Failed to run the callable object {} decorated by '@paddle.jit.to_static', "
429
                "because it is NOT in dynamic mode. Please disable the static mode to enter dynamic mode with the "
430
                "following API: paddle.disable_static().".format(
431 432 433
                    self.dygraph_function
                )
            )
434

435 436
        # 2. trace ops from dygraph layers and cache the generated program.
        args, kwargs = self._function_spec.unified_args_and_kwargs(args, kwargs)
437

438 439
        try:
            concrete_program, partial_program_layer = self.get_concrete_program(
440 441
                *args, **kwargs, is_train=self._is_train_mode()
            )
442 443 444
            # 3. synchronize self.training attribute.
            if isinstance(self._class_instance, layers.Layer):
                partial_program_layer.training = self._class_instance.training
445 446
            else:
                partial_program_layer.training = self._training
447

448 449 450
            partial_program_layer._cuda_graph_capture_mode = (
                self._cuda_graph_capture_mode
            )
451 452
            partial_program_layer._cuda_graph_pool_id = self._cuda_graph_pool_id

453
            # 4. return outputs.
454 455 456 457 458 459 460
            try:
                return partial_program_layer(args)
            except Exception as e:
                if not hasattr(e, error.ERROR_DATA):
                    # runtime error
                    error.attach_error_data(e, in_runtime=True)
                    raise
461
        except Exception as e:
462
            error_data = getattr(e, error.ERROR_DATA, None)
463
            if error_data:
464
                error_data.raise_new_exception()
465
            else:
466 467
                logging_utils.warn(
                    "Please file an issue at 'https://github.com/PaddlePaddle/Paddle/issues'"
468 469
                    " if you can't handle this {} yourself.".format(type(e))
                )
470
                raise e
471

472 473
    def _is_train_mode(self):
        if self._class_instance is not None:
474 475 476
            if not hasattr(self._class_instance, 'training'):
                raise TypeError(
                    "When using 'to_static' to convert method of a class, "
477 478
                    "please ensure the class inherits from nn.Layer"
                )
479 480 481 482
            return self._class_instance.training
        else:
            return self._training

483 484 485 486 487 488
    def _call_dygraph_function(self, *args, **kwargs):
        """
        Calls dygraph function directly and returns the outputs.

        Args:
            *args(tuple): tuple of all input arguments from original decorated function.
489
            **kwargs(dict): dict of all input keyward arguments from original decorated function.
490 491 492 493 494 495

        Return:
            Outputs of dygraph function.
        """
        if self._class_instance is not None:
            dygraph_function = self._dygraph_function.__get__(
496 497
                self._class_instance
            )
498 499 500 501 502
        else:
            dygraph_function = self._dygraph_function

        return dygraph_function(*args, **kwargs)

503 504 505 506 507 508 509 510 511
    def _raise_when_property(self):
        """raise RuntimeError when property=True

        Raises:
            RuntimeError: can not call this func when property=True
        """
        if self.is_property:
            raise RuntimeError("Can not call the func when property=True.")

512 513 514 515 516 517 518 519 520 521 522
    def get_concrete_program(self, *args, **kwargs):
        """
        Returns traced concrete program and inner executable partial layer.

        Args:
            *args(tuple): input arguments values or InputSpec
            **kwargs(dict) : input kwargs values.

        Returns:
            Traced ConcreteProgram and executable translated Layer.
        """
523
        self._raise_when_property()
524

525
        with_hook = kwargs.get("with_hook", False)
526
        is_train = kwargs.get("is_train", True)
527 528 529 530
        if "is_train" in kwargs:
            kwargs.pop("is_train")
        if "with_hook" in kwargs:
            kwargs.pop("with_hook")
531 532
        # 1. unify args/kwargs and replace Tensor with InputSpec
        if len(args) != len(self._function_spec.args_name):
533
            args, kwargs = self._function_spec.unified_args_and_kwargs(
534 535 536 537 538 539
                args, kwargs
            )
        (
            input_args_with_spec,
            input_kwargs_with_spec,
        ) = self._function_spec.args_to_input_spec(args, kwargs)
540 541

        # 2. generate cache key
542 543 544 545 546 547 548 549 550
        cache_key = CacheKey(
            self._function_spec,
            input_args_with_spec,
            input_kwargs_with_spec,
            self._class_instance,
            **self._kwargs,
            with_hook=with_hook,
            is_train=is_train
        )
551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581

        # 3. check whether hit the cache or build a new program for the input arguments
        concrete_program, partial_program_layer = self._program_cache[cache_key]
        return concrete_program, partial_program_layer

    def get_traced_count(self):
        """
        Returns the number of traced programs for the decorated function.
        """
        return len(self._program_cache)

    @property
    def code(self):
        """
        Returns the source code of transformed static function for debugging.
        """
        static_func = convert_to_static(self._dygraph_function)
        source_code = func_to_source_code(static_func)
        return source_code

    @property
    def dygraph_function(self):
        """
        Returns the original decorated function.
        """
        return self._dygraph_function

    @property
    def concrete_program(self):
        """
        Returns recent ConcreteProgram instance of decorated function.
A
Aurelius84 已提交
582 583 584 585 586 587 588 589 590 591 592 593 594

        Examples:
            .. code-block:: python

                import paddle
                from paddle.jit import to_static
                from paddle.static import InputSpec

                paddle.disable_static()

                def foo(x, y):
                    z = x + y
                    return z
595

A
Aurelius84 已提交
596 597 598 599 600 601 602 603
                # usage 1:
                decorated_foo = to_static(foo, input_spec=[InputSpec([10], name='x'), InputSpec([10], name='y')])
                print(decorated_foo.concrete_program)

                # usage 2:
                decorated_foo = to_static(foo)
                out_foo = decorated_foo(paddle.rand([10]), paddle.rand([10]))
                print(decorated_foo.concrete_program)
604
        """
605 606
        return self.concrete_program_specify_input_spec(input_spec=None)

607 608 609
    def concrete_program_specify_input_spec(
        self, input_spec=None, with_hook=False
    ):
610 611 612
        """
        Returns recent ConcreteProgram instance of decorated function while
        specifying input_spec. If the self._function_spec already has
613
        input_spec, it will check the compatibility of input input_spec and
614 615 616 617 618 619 620
        the self._function_spec.input_spec. If input input_spec=None, then
        this method uses self._function_spec.input_spec

        args:
            input_spec (list[InputSpec], optional): Describes the input of
                the translate function.
        """
621
        self._raise_when_property()
622 623 624 625 626
        # if specific the `input_spec`, the length of program_cache will always 1,
        # else, return the last one.
        cached_program_len = len(self._program_cache)
        # If specific `input_spec`, apply convertion from dygraph layers into static Program.
        if cached_program_len == 0:
C
Chen Weihang 已提交
627 628 629
            desired_input_spec = input_spec
            if self._function_spec.input_spec is not None:
                if input_spec is not None and not input_specs_compatible(
630 631
                    flatten(input_spec), flatten(self._function_spec.input_spec)
                ):
632
                    raise ValueError(
633 634 635 636
                        "The `input_spec`: {} used to construct concrete_program is conflict with the `input_spec`: {} in `@paddle.jit.to_static`".format(
                            input_spec, self._function_spec.input_spec
                        )
                    )
C
Chen Weihang 已提交
637 638 639
                # NOTE(chenweihang): we should always translated program based on the `input_spec`
                # decorated on forward if it is valid
                desired_input_spec = self._function_spec.input_spec
640 641
                if input_spec is not None:
                    logging_utils.warn(
642 643 644 645
                        "\n\nYou have specified `input_spec` both in function definition (higher priority) and `paddle.jit.save` (will be ignored.)\n\n\t Using: {}\n\n\t Ignore: {}\n".format(
                            desired_input_spec, input_spec
                        )
                    )
646

647
            has_input_spec = desired_input_spec is not None
A
Aurelius84 已提交
648
            if has_input_spec:
C
Chen Weihang 已提交
649
                concrete_program, _ = self.get_concrete_program(
650 651
                    *desired_input_spec,
                    with_hook=with_hook,
652 653
                    is_train=self._is_train_mode()
                )
654
                return concrete_program
655
            else:
A
Aurelius84 已提交
656
                raise ValueError(
657 658 659 660
                    "No valid transformed program for {}.\n\t    Please specific `input_spec` in `@paddle.jit.to_static` or feed input tensor to call the decorated function at once.\n".format(
                        self._function_spec
                    )
                )
661 662 663 664 665 666
        elif with_hook:
            cache_key = self._program_cache._recent_cache_key
            cache_key.kwargs["with_hook"] = True
            concrete_program, _ = self._program_cache[cache_key]
            return concrete_program

667 668
        # If more than one programs have been cached, return the recent converted program by default.
        elif cached_program_len > 1:
669
            logging_utils.warn(
670 671 672 673
                "Current {} has more than one cached programs: {}, the last traced progam will be return by default.".format(
                    self._function_spec, cached_program_len
                )
            )
674

675 676 677 678
        cache_key, (
            concrete_program,
            partial_layer,
        ) = self._program_cache.last()
679
        return concrete_program
680

681 682 683
    def rollback(self):
        """
        Rollback into original dygraph functions for current class instance.
684

685 686 687 688 689 690 691 692 693 694
        Returns:
            Function or Method

        Example::
            .. code-block:: python

                import paddle

                class Net(paddle.nn.Layer):
                    def __init__(self):
695
                        super().__init__()
696 697 698 699 700 701 702 703 704 705 706

                    def forward(self, x, flag=True):
                        if flag:
                            out = x + 1
                        else:
                            out = x - 1
                        return out

                x = paddle.randn([10, 1], 'float32')
                net = paddle.jit.to_static(Net())  # convert into static mode
                out = net(x)
707

708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723
                net.forward.rollback()  # rollback into dygraph mode
                out = net(x)
        """

        def rollback_impl(class_instance):
            for name, func in class_instance._original_funcs.items():
                setattr(class_instance, name, func.__get__(class_instance))

            for sublayer in class_instance.sublayers(include_self=False):
                rollback_impl(sublayer)

        if self._class_instance is None:
            return self._dygraph_function

        # only rollback sub-functions on path of top _dygraph_function
        func_name = self._dygraph_function.__name__
724 725 726 727 728
        assert (
            func_name in self._class_instance._original_funcs
        ), "Not Found function '{}' in class '{}'.".format(
            func_name, self._class_instance.__name__
        )
729
        func = self._class_instance._original_funcs[func_name]
730 731 732
        setattr(
            self._class_instance, func_name, func.__get__(self._class_instance)
        )
733 734 735 736 737 738

        for sublayer in self._class_instance.sublayers(include_self=False):
            rollback_impl(sublayer)

        return getattr(self._class_instance, func_name)

739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754
    def __deepcopy__(self, memo):
        """
        Customized behavior for copy.deepcopy, return original decorated function instead
        of a new StaticFunction Object. StaticFunction itself is not copyable becuase it's
        associated with class_instance.

        We add __deepcopy__ here only for the following usage:

        Example::
            .. code-block:: python

                import copy
                import paddle

                class Net(paddle.nn.Layer):
                    def __init__(self):
755
                        super().__init__()
756 757 758 759 760 761 762 763 764 765 766 767

                    def forward(self, x, flag=True):
                        if flag:
                            out = x + 1
                        else:
                            out = x - 1
                        return out

                x = paddle.randn([10, 1], 'float32')
                net = paddle.jit.to_static(Net())  # convert into static mode

                copy_net = copy.deepcopy(net)      # deepcopy a new net without @to_static
768

769 770 771 772 773 774
        Please attention that original 'net' will unwrap @to_static and rollback into simple Layer.
        """
        if self._class_instance is not None:
            net_name = type(self._class_instance).__name__
            logging_utils.log(
                level=-1,
775 776 777 778 779
                msg="Not recommend to deepcopy '{}' decorated with @to_static, it has side effect that will"
                " rollback into original state before @to_static. Please deepcopy '{}' before applying @to_static.".format(
                    net_name, net_name
                ),
            )
780
            self.rollback()
781 782 783
            return self._dygraph_function.__get__(
                memo[id(self._class_instance)]
            )
784 785 786
        else:
            return self._dygraph_function

787 788 789 790 791
    @property
    def inputs(self):
        """
        Returns input tensors of recent converted static program.
        """
792
        self._raise_when_property()
793 794
        concrete_program = self.concrete_program
        inputs = [
795 796
            var
            for var in flatten(concrete_program.inputs)
797 798 799
            if isinstance(var, framework.Variable)
        ]
        return inputs
800

801
    @property
802 803 804 805
    def outputs(self):
        """
        Returns output tensors of recent converted static program.
        """
806
        self._raise_when_property()
807 808
        concrete_program = self.concrete_program
        outputs = [
809 810
            var
            for var in flatten(concrete_program.outputs)
811 812 813 814
            if isinstance(var, framework.Variable)
        ]

        return outputs
815

816
    @property
817 818 819 820
    def main_program(self):
        """
        Returns recent converted static main program.
        """
821
        self._raise_when_property()
822 823 824
        concrete_program = self.concrete_program
        main_program = concrete_program.main_program
        return main_program
825

826 827 828
    @property
    def program_cache(self):
        return self._program_cache
829

830 831 832
    @property
    def function_spec(self):
        return self._function_spec
833 834


835 836 837 838 839 840 841 842 843 844
def _verify_init_in_dynamic_mode(class_instance):
    """
    Verifies the instance is initialized in dynamic mode.
    """
    if isinstance(class_instance, layers.Layer):
        if not class_instance._init_in_dynamic_mode:
            raise RuntimeError(
                " `paddle.jit.to_static` is only available in dynamic mode. Please call `paddle.disable_static()` before "
                "initializing your Layer class `{}` . Because parameters of Layer class should be initialized firstly "
                "in dynamic mode while applying transformation.".format(
845 846 847
                    class_instance
                )
            )
848 849


850
class HookHelper:
851 852 853 854 855 856 857 858 859
    """
    Only For converting pre/post hooks operation in outermost layer while jit.save.
    Because hooks in sublayer have been processed automatically.
    """

    def __init__(self, func, class_instance, with_hook=False):
        self.func = func
        self.class_instance = class_instance
        self.with_hook = with_hook
860 861 862 863 864
        self.need_apply_hook = (
            with_hook
            and isinstance(self.class_instance, layers.Layer)
            and getattr(func, "__name__") == "forward"
        )
865 866 867 868 869

    def apply_pre_hooks(self, inputs):
        """
        Apply _forward_pre_hooks from outermost layer
        """
870 871
        if not self.need_apply_hook:
            return inputs
872 873 874 875 876 877

        inputs = inputs[1:]
        for forward_pre_hook in self.class_instance._forward_pre_hooks.values():
            hook_result = forward_pre_hook(self.class_instance, inputs)
            if hook_result is not None:
                if not isinstance(hook_result, tuple):
878
                    hook_result = (hook_result,)
879 880 881 882 883 884 885 886
                inputs = hook_result

        return [self.class_instance] + list(inputs)

    def apply_post_hooks(self, inputs, outputs):
        """
        Apply _forward_post_hooks from outermost layer
        """
887 888
        if not self.need_apply_hook:
            return outputs
889 890

        inputs = inputs[1:]
891 892 893 894 895 896
        for (
            forward_post_hook
        ) in self.class_instance._forward_post_hooks.values():
            hook_result = forward_post_hook(
                self.class_instance, inputs, outputs
            )
897 898 899 900 901 902 903
            if hook_result is not None:
                outputs = hook_result

        inputs.insert(0, self.class_instance)
        return outputs


904
class ConcreteProgram:
905 906

    __slots__ = [
907 908 909 910 911 912 913
        'inputs',
        'outputs',
        'main_program',
        "startup_program",
        "parameters",
        "function",
        'kwargs',
914 915
    ]

916 917 918 919 920 921 922 923 924 925
    def __init__(
        self,
        inputs,
        outputs,
        parameters,
        function,
        main_program,
        startup_program=None,
        **kwargs
    ):
926 927 928
        self.inputs = inputs
        self.outputs = outputs
        self.main_program = main_program
929
        self.startup_program = startup_program
930
        self.parameters = parameters
931
        self.function = function
932
        self.kwargs = kwargs
933 934 935

    @staticmethod
    @switch_to_static_graph
936 937 938
    def from_func_spec(
        func_spec, input_spec, input_kwargs_spec, class_instance, **kwargs
    ):
939
        """
940 941
        Builds the main_program with specialized inputs and returns outputs
        of program as fetch_list.
942 943 944

        Args:
            func_spec(FunctionSpec): A FunctionSpec instance for decorated function.
945
            input_spec(list[InputSpec]):
946
        """
947 948 949
        # verify the instance is initialized in imperative mode.
        _verify_init_in_dynamic_mode(class_instance)

950
        # Transforms dygraph function into static function and caches it.
951
        dygraph_function = func_spec.dygraph_function
952
        static_func = convert_to_static(dygraph_function)
953
        # apply pre\post hook for outermost layer
954 955 956
        hook_helper = HookHelper(
            dygraph_function, class_instance, kwargs.get("with_hook", False)
        )
957

958 959
        main_program, startup_program = framework.Program(), framework.Program()
        # Note: The random seed should be synchronized into cached program
960
        # if set in `fluid.dygraph_guard` because some ops rely on it, such as
961
        # `fluid.layers.dropout`.
962
        main_program.random_seed = framework.default_main_program().random_seed
963 964 965
        startup_program.random_seed = (
            framework.default_startup_program().random_seed
        )
966

967
        from paddle.fluid.dygraph.base import _switch_declarative_mode_guard_
968

969
        with framework.program_guard(main_program, startup_program):
970 971
            with _switch_declarative_mode_guard_(is_declarative=True):
                # 1. Adds `fluid.data` layers for input if needed
972
                static_inputs = func_spec.to_static_inputs_with_spec(
973 974
                    input_spec, main_program
                )
975
                _kwargs = func_spec.to_static_inputs_with_spec(
976 977
                    input_kwargs_spec, main_program
                )
978
                if class_instance:
979 980 981
                    static_inputs = tuple(
                        [class_instance] + list(static_inputs)
                    )
982

983
                # 2. Gets all ParamBases and buffered VarBases in the function
984
                all_parameters_and_buffers = _extract_indeed_params_buffers(
985 986
                    class_instance
                )
987 988

                # 3. Builds program only once and returns the output Variables.
989 990 991
                with param_guard(
                    get_parameters(class_instance, False)
                ), param_guard(get_buffers(class_instance, False)):
992
                    try:
993 994
                        # only for jit.save, do nothing while train and eval process
                        inputs = hook_helper.apply_pre_hooks(static_inputs)
995 996
                        if _kwargs:
                            outputs = static_func(*inputs, **_kwargs)
997 998
                        else:
                            outputs = static_func(*inputs)
999
                        outputs = hook_helper.apply_post_hooks(inputs, outputs)
1000 1001
                    except BaseException as e:
                        # NOTE: If e is raised in compile time, e should be attached to ERROR_DATA here.
1002
                        error.attach_error_data(e)
1003 1004 1005
                        error_data = getattr(e, error.ERROR_DATA, None)
                        if error_data:
                            error_data.raise_new_exception()
1006 1007
                        raise

1008
                if outputs is not None:
1009 1010 1011 1012
                    need_wrap_into_list = (
                        not isinstance(outputs, (tuple, list))
                        or len(outputs) == 1
                    )
1013 1014
                    if need_wrap_into_list:
                        outputs = [outputs]
1015

1016 1017
        main_program = update_op_callstack_with_origin_info(main_program)

1018 1019 1020 1021 1022 1023 1024 1025 1026
        return ConcreteProgram(
            inputs=static_inputs,
            outputs=outputs,
            parameters=all_parameters_and_buffers,
            function=dygraph_function,
            main_program=main_program,
            startup_program=startup_program,
            **kwargs
        )
1027 1028


1029 1030 1031 1032 1033 1034
def _extract_indeed_params_buffers(class_instance):
    """
    To filter not initialzed buffers.
    """
    params = list(get_parameters(class_instance).values())
    buffers = list(get_buffers(class_instance).values())
1035
    buffers = [buffer for buffer in buffers if len(buffer.shape) != 0]
1036 1037 1038 1039

    return params + buffers


1040
class ProgramCache:
1041 1042 1043
    """
    Wrapper class for the program functions defined by dygraph function.
    """
1044

1045
    def __init__(self):
1046
        # {hash_id : (concrete_program, partial_layer)}
1047
        self._caches = collections.OrderedDict()
1048
        # trace mostly recent used program
1049
        self._recent_key = None
1050
        self._recent_cache_key = None
1051

1052 1053 1054
    def _build_once(self, cache_key):
        concrete_program = ConcreteProgram.from_func_spec(
            func_spec=cache_key.function_spec,
1055 1056
            input_spec=cache_key.input_args_with_spec,
            input_kwargs_spec=cache_key.input_kwargs_with_spec,
1057
            class_instance=cache_key.class_instance,
1058 1059
            **cache_key.kwargs
        )
1060
        return concrete_program, partial_program_from(concrete_program)
1061

1062
    def __getitem__(self, item):
1063
        if not isinstance(item, CacheKey):
1064 1065 1066 1067
            raise ValueError(
                'type(item) should be CacheKey, but received %s'
                % type_name(item)
            )
1068
        item_id = hash(item)
1069
        self._recent_cache_key = item
1070
        self._recent_key = item_id
1071 1072
        if item_id not in self._caches:
            self._caches[item_id] = self._build_once(item)
1073 1074 1075
            # Note: raise warnings if number of traced program is more than `max_tracing_count`
            current_tracing_count = len(self._caches)
            if current_tracing_count > MAX_TRACED_PROGRAM_COUNT:
1076
                logging_utils.warn(
1077
                    "Current traced program number: {} > `max_tracing_count`:{}. Too much cached programs will bring expensive overhead. "
1078 1079 1080 1081
                    "The reason may be: (1) passing tensors with different shapes, (2) passing python objects instead of tensors.".format(
                        current_tracing_count, MAX_TRACED_PROGRAM_COUNT
                    )
                )
1082

1083
        return self._caches[item_id]
1084

1085
    def get_program(self, item):
1086
        if not isinstance(item, CacheKey):
1087
            raise ValueError(
1088 1089 1090
                "Input item's type should be FunctionSpec, but received %s"
                % type_name(item)
            )
1091 1092
        item_id = hash(item)
        if item_id not in self._caches:
1093
            raise RuntimeError(
1094
                "Failed to find program for input item, please decorate input function by `@paddle.jit.to_static`."
1095
            )
1096
        return self._caches[item_id]
1097

1098
    def last(self):
1099 1100 1101
        assert (
            len(self._caches) >= 1
        ), "No valid cached program in ProgramCache."
1102 1103
        assert self._recent_key is not None
        return self._recent_key, self._caches[self._recent_key]
1104

1105 1106 1107 1108
    def __len__(self):
        return len(self._caches)

    def concrete_programs(self):
1109
        return [cp for key, (cp, _) in self._caches.items()]
1110

1111

1112 1113
def synchronized(func):
    func.__lock__ = threading.Lock()
1114

1115 1116 1117
    def lock_func(*args, **kwargs):
        with func.__lock__:
            return func(*args, **kwargs)
1118

1119
    return lock_func
1120 1121


1122
class ProgramTranslator:
1123
    """
1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135
    Class to translate dygraph function into static graph function. The object
    of this class is a singleton.

    Args:
        None.

    Returns:
        ProgramTranslator: the singleton object.

    Examples:
        .. code-block:: python

1136
            import paddle
1137

1138 1139 1140
            # Two methods get same object because ProgramTranslator is a singleton
            paddle.jit.ProgramTranslator()
            paddle.jit.ProgramTranslator.get_instance()
1141

1142 1143
    """

1144
    _singleton_lock = threading.Lock()
1145 1146 1147 1148 1149 1150
    _instance = None

    @synchronized
    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = object.__new__(cls, *args, **kwargs)
1151
            cls._instance._initialized = False
1152 1153 1154 1155 1156
        return cls._instance

    @classmethod
    def get_instance(cls):
        if cls._instance is None:
1157 1158
            with cls._singleton_lock:
                cls._instance = cls()
1159 1160 1161 1162 1163
        return cls._instance

    @classmethod
    def reset(cls):
        if cls._instance is not None:
1164
            cls._instance._initialized = False
1165 1166
            cls._instance.__init__()

1167
    def __init__(self):
1168
        # To make sure that calls __init__ only once.
1169
        if self._initialized:
1170
            return
1171 1172
        self._initialized = True
        self._program_cache = ProgramCache()
1173
        self.enable_to_static = True
1174

1175
    def enable(self, enable_to_static):
1176
        """
1177
        Enable or disable the converting from imperative to static graph by
1178 1179 1180
        ProgramTranslator globally.

        Args:
1181
            enable_to_static (bool): True or False to enable or disable converting to static.
1182 1183 1184 1185 1186 1187 1188

        Returns:
            None.

        Examples:
            .. code-block:: python

1189
                import paddle
1190 1191


1192 1193 1194 1195 1196 1197 1198
                @paddle.jit.to_static
                def func(x):
                    if paddle.mean(x) > 0:
                        x_v = x - 1
                    else:
                        x_v = x + 1
                    return x_v
1199

1200 1201 1202 1203 1204 1205

                prog_trans = paddle.jit.ProgramTranslator()
                prog_trans.enable(False)

                x = paddle.ones([1, 2])
                # ProgramTranslator is disabled so the func is run in dygraph
1206
                print(func(x))  # [[0. 0.]]
L
liym27 已提交
1207

1208
        """
1209 1210 1211 1212 1213 1214
        check_type(
            enable_to_static,
            "enable_to_static",
            bool,
            "ProgramTranslator.enable",
        )
1215
        self.enable_to_static = enable_to_static
1216

1217 1218
    def get_output(self, dygraph_func, *args, **kwargs):
        """
1219
        Returns the output dygraph Tensor for dygraph function. The dygraph
1220
        function will be translated into static graph function so the under
1221
        beneath numerical result will be calculated by static graph mode.
1222 1223 1224

        Args:
            dygraph_func (callable): the dygraph function.
1225 1226
            *args (tuple): the input argument of dygraph_func.
            **kwargs (dict): the input argument of dygraph_func.
1227 1228

        Returns:
1229
            Tensor or tuple of Tensors: the dygraph Tensor containing digital result.
1230 1231 1232 1233

        Examples:
            .. code-block:: python

1234 1235
                import paddle

1236 1237

                def func(x):
1238
                    if paddle.mean(x) > 0:
1239 1240 1241 1242 1243 1244
                        x_v = x - 1
                    else:
                        x_v = x + 1
                    return x_v


1245 1246 1247 1248
                prog_trans = paddle.jit.ProgramTranslator()

                x = paddle.ones([1, 2])
                x_v = prog_trans.get_output(func, x)
1249
                print(x_v)  # [[0. 0.]]
1250

1251
        """
1252 1253 1254
        assert callable(
            dygraph_func
        ), "Input dygraph_func is not a callable in ProgramTranslator.get_output"
1255

1256
        if not self.enable_to_static:
1257 1258
            # Here calls `warnings.warn` but not `logging_utils.warn` because by default warnings.warn(message)
            # will show up **only once**.
1259
            logging_utils.warn(
1260 1261 1262 1263
                "The ProgramTranslator.get_output doesn't work when setting ProgramTranslator.enable to False. "
                "We will just return dygraph output. "
                "Please call ProgramTranslator.enable(True) if you would like to get static output."
            )
1264
            return dygraph_func(*args, **kwargs)
1265
        try:
1266
            function_spec = FunctionSpec(dygraph_func)
1267
            cache_key = CacheKey.from_func_and_args(
1268 1269 1270 1271 1272
                function_spec,
                args,
                kwargs,
                getattr(dygraph_func, '__self__', None),
            )
1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288
            _, partial_program_layer = self._program_cache[cache_key]

            if args and isinstance(args[0], layers.Layer):
                # Synchronize self.training attribute.
                partial_program_layer.training = args[0].training
                args = args[1:]
            try:
                return partial_program_layer(args)
            except BaseException as e:
                # NOTE:
                # 1. If e is raised in compile time, e should have been attached to ERROR_DATA before;
                # 2. If e raised in runtime, e should be attached to ERROR_DATA here.
                if not hasattr(e, error.ERROR_DATA):
                    # runtime error
                    error.attach_error_data(e, in_runtime=True)
                raise
1289
        except BaseException as e:
1290 1291 1292 1293 1294 1295
            error_data = getattr(e, error.ERROR_DATA, None)
            if error_data:
                error_data.raise_new_exception()
            else:
                logging_utils.warn(
                    "Please file an issue at 'https://github.com/PaddlePaddle/Paddle/issues'"
1296 1297
                    " if you can't handle this {} yourself.".format(type(e))
                )
1298
                raise e
1299 1300 1301

    def get_func(self, dygraph_func):
        """
1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312
        Returns a callable function which converts imperative dygraph APIs of
        the input dygraph_func into declarative net-building APIs, which means
        it doesn't return immediate digital result as get_output does.
        Users should handle Program and Executor by themselves.

        Args:
            dygraph_func (callable): the dygraph function.

        Returns:
            callable: converting imperative dygraph APIs into declarative
            net-building APIs.
1313 1314 1315 1316

        Examples:
            .. code-block:: python

1317 1318
                import paddle

1319 1320

                def func(x):
1321
                    if paddle.mean(x) > 0:
1322 1323 1324 1325 1326 1327
                        x_v = x - 1
                    else:
                        x_v = x + 1
                    return x_v


1328
                prog_trans = paddle.jit.ProgramTranslator()
1329 1330 1331
                static_func = prog_trans.get_func(func)
                print(callable(static_func)) # True

1332
        """
1333 1334 1335
        assert callable(
            dygraph_func
        ), "Input dygraph_func is not a callable in ProgramTranslator.get_func"
1336

1337
        if not self.enable_to_static:
1338
            logging_utils.warn(
1339 1340 1341
                "The ProgramTranslator.get_func doesn't work when setting ProgramTranslator.enable to False. We will "
                "just return dygraph output. Please call ProgramTranslator.enable(True) if you would like to get static output."
            )
1342
            return dygraph_func
1343

1344
        static_func = convert_to_static(dygraph_func)
1345 1346
        return static_func

1347 1348
    def get_program(self, dygraph_func, *args, **kwargs):
        """
1349
        Returns the translated static program and input/output Tensors from
1350 1351 1352 1353
        dygraph function. The users can use the program to run by executor.

        Args:
            dygraph_func (callable): the dygraph function.
1354 1355
            *args (tuple): the input argument of dygraph_func.
            **kwargs (dict): the input argument of dygraph_func.
1356 1357 1358

        Returns:
            tuple of (main_program, startup_program, inputs, outputs) whose
1359
            types are (Program, Program, list of Tensors, list of Tensors).
1360 1361
            main_program: the converted main program.
            startup_program: the converted startup program.
1362 1363
            inputs: list of input Tensors which need to be fed.
            outputs: list of output Tensors which users can fetch.
1364 1365 1366 1367

        Examples:
            .. code-block:: python

1368 1369
                import paddle

1370 1371

                def func(x):
1372
                    if paddle.mean(x) > 0:
1373 1374 1375 1376 1377 1378
                        x_v = x - 1
                    else:
                        x_v = x + 1
                    return x_v


1379 1380
                prog_trans = paddle.jit.ProgramTranslator()
                x = paddle.ones([1, 2])
1381 1382
                main_prog, start_prog, inputs, outputs = prog_trans.get_program(func, x)
                print([i.name for i in inputs])
1383
                # [u'generated_tensor_0'] the feed input Tensor name representing x
1384
                print([o.name for o in outputs])
1385
                # [u'_generated_var_4'] the fetch output Tensor name representing x_v
1386

1387
        """
1388 1389 1390
        assert callable(
            dygraph_func
        ), "Input dygraph_func is not a callable in ProgramTranslator.get_program"
1391

1392
        if not self.enable_to_static:
1393
            logging_utils.warn(
1394 1395 1396 1397
                "The ProgramTranslator.get_program doesn't work when setting ProgramTranslator.enable to False."
                "We will just return dygraph output. "
                "Please call ProgramTranslator.enable(True) if you would like to get static output."
            )
1398
            return dygraph_func(*args, **kwargs)
1399

1400
        function_spec = FunctionSpec(dygraph_func)
1401
        cache_key = CacheKey.from_func_and_args(
1402 1403
            function_spec, args, kwargs, getattr(dygraph_func, '__self__', None)
        )
1404 1405
        concrete_program, partial_program_layer = self._program_cache[cache_key]

1406 1407
        # Note: concrete_program hold all input/output infos include non-Variable
        input_vars = [
1408 1409
            var
            for var in concrete_program.inputs
1410 1411 1412
            if isinstance(var, framework.Variable)
        ]
        output_vars = [
1413 1414
            var
            for var in concrete_program.outputs
1415 1416 1417
            if isinstance(var, framework.Variable)
        ]

1418 1419 1420 1421 1422 1423
        return (
            concrete_program.main_program,
            concrete_program.startup_program,
            input_vars,
            output_vars,
        )
1424

1425 1426
    def get_code(self, dygraph_func):
        """
1427 1428 1429 1430 1431 1432
        Returns the translated static function string code from dygraph function.

        Args:
            dygraph_func (callable): the dygraph function.

        Returns:
1433 1434 1435 1436 1437
            str: the string code of translated static function.

        Examples:
            .. code-block:: python

1438 1439 1440 1441 1442 1443 1444 1445 1446
                import paddle


                def func(x):
                    if paddle.mean(x) > 0:
                        x_v = x - 1
                    else:
                        x_v = x + 1
                    return x_v
1447 1448


1449
                prog_trans = paddle.jit.ProgramTranslator()
1450

1451 1452
                code = prog_trans.get_code(func)
                print(type(code)) # <class 'str'>
1453

1454
        """
1455 1456 1457
        assert callable(
            dygraph_func
        ), "Input dygraph_func is not a callable in ProgramTranslator.get_code"
1458
        # Gets AST from dygraph function
1459 1460 1461

        unwrap_func = unwrap(dygraph_func)
        raw_code = inspect.getsource(unwrap_func)
1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472
        code = textwrap.dedent(raw_code)
        root = gast.parse(code)

        # Transform AST
        dygraph_to_static = DygraphToStaticAst()
        root_wrapper = dygraph_to_static.get_static_ast(root)

        # Get source_code
        source_code = ast_to_source_code(root_wrapper.node)
        return source_code

1473
    def get_program_cache(self):
1474
        """
1475 1476 1477 1478 1479 1480 1481 1482 1483
        Returns the ProgramCache instance. This method is used by PaddlePaddle
        developers to manage program cache in ProgramTranslator. Normal users
        don't have to call this method.

        Returns:
            ProgramCache: ProgramCache instance of ProgramTranslator.

        Examples:
            .. code-block:: python
1484

1485
                import paddle
1486

1487
                prog_trans = paddle.jit.ProgramTranslator()
1488 1489
                prog_cache = prog_trans.get_program_cache()

1490
        """
1491
        return self._program_cache