jit.py 14.0 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
__all__ = [
18
    'TracedLayer', 'dygraph_to_static_code', 'dygraph_to_static_func',
19 20
    'dygraph_to_static_output'
]
21

22
import warnings
23 24

from ..wrapped_decorator import wrap_decorator
25
from .base import program_desc_tracing_guard, switch_to_static_graph
26
from .dygraph_to_static import ProgramTranslator, convert_to_static
27
from .layers import Layer
28 29 30 31
from paddle.fluid import core
from paddle.fluid.framework import Program, Block, Variable, _dygraph_tracer, dygraph_only, _dygraph_guard, _current_expected_place, in_dygraph_mode
from paddle.fluid.executor import Executor, scope_guard
from paddle.fluid.compiler import CompiledProgram
32 33 34 35 36 37 38 39 40 41 42 43


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


def _extract_vars(inputs, result_list):
    if isinstance(inputs, Variable):
44
        result_list.append(inputs)
45 46 47 48 49 50 51 52 53 54 55 56

    if isinstance(inputs, (list, tuple)):
        for var in inputs:
            _extract_vars(var, result_list)


def extract_vars(inputs):
    result_list = []
    _extract_vars(inputs, result_list)
    return result_list


57
def _dygraph_to_static_code_(dygraph_func):
58 59 60
    def __impl__(*args, **kwargs):
        if in_dygraph_mode():
            warnings.warn(
61
                "The decorator 'dygraph_to_static_code' doesn't work in dygraph mode."
62 63
                " Please use it in static mode.")
            return dygraph_func(*args, **kwargs)
64 65
        program_translator = ProgramTranslator()
        return program_translator.get_code(dygraph_func)
66 67 68 69

    return __impl__


70
dygraph_to_static_code = wrap_decorator(_dygraph_to_static_code_)
71 72


73
def _dygraph_to_static_func_(dygraph_func):
74 75 76
    def __impl__(*args, **kwargs):
        if in_dygraph_mode():
            warnings.warn(
77
                "The decorator 'dygraph_to_static_func' doesn't work in dygraph mode."
78 79
                " Please use it in static mode.")
            return dygraph_func(*args, **kwargs)
80 81 82 83 84
        program_translator = ProgramTranslator()
        static_func = program_translator.get_func(dygraph_func)
        return static_func(*args, **kwargs)

    return __impl__
85

86

87
dygraph_to_static_func = wrap_decorator(_dygraph_to_static_func_)
88

89 90 91 92 93

def _dygraph_to_static_output_(dygraph_func):
    def __impl__(*args, **kwargs):
        program_translator = ProgramTranslator()
        return program_translator.get_output(dygraph_func, *args, **kwargs)
94

95
    return __impl__
96 97


98 99 100
dygraph_to_static_output = wrap_decorator(_dygraph_to_static_output_)


101
@dygraph_only
Z
Zeng Jinle 已提交
102 103 104 105 106
def _trace(layer,
           inputs,
           feed_prefix='feed_',
           fetch_prefix='fetch_',
           tmp_prefix='t_'):
107
    assert isinstance(layer, Layer)
108 109 110 111 112 113 114 115 116

    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):
117
        original_outputs = layer(*inputs)
118 119 120 121
        if not isinstance(original_outputs, (list, tuple)):
            outputs = [original_outputs]
        else:
            outputs = original_outputs
122
        out_vars = [var for var in outputs]
123

124
        program_desc, feed_names, fetch_names, parameters = tracer.create_program_desc(
Z
Zeng Jinle 已提交
125
            var_list, feed_prefix, out_vars, fetch_prefix, tmp_prefix)
126 127 128 129 130
        tracer.reset()

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

131
    return original_outputs, program, feed_names, fetch_names, parameters
132 133 134 135


class TracedLayer(object):
    """
136 137 138 139 140
    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.
141 142 143 144

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

    All TracedLayer objects should not be created by constructor and should
147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
    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

        self._place = _current_expected_place()

        self._scope = core.Scope()
        for p in parameters:
163
            src_tensor = p.value().get_tensor()
164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186
            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):
        """
187
        This method is the only allowed method to create TracedLayer object.
188 189 190 191
        It would call the :code:`layer(*inputs)` method to run the dygraph
        model and convert it into a static graph model.

        Args:
192 193
            layer (dygraph.Layer): the layer object to be traced.
            inputs (list(Variable)): the input variables of the layer object.
194 195

        Returns:
196
            tuple: A tuple of 2 items, whose the first item is the output of
197
            :code:`layer(*inputs)` , and the second item is the created
198
            TracedLayer object.
199

200
        Examples:
201 202 203
            .. code-block:: python:

                import paddle.fluid as fluid
204
                from paddle.fluid.dygraph import Linear, to_variable, TracedLayer
205 206 207
                import numpy as np

                class ExampleLayer(fluid.dygraph.Layer):
208 209 210
                    def __init__(self):
                        super(ExampleLayer, self).__init__()
                        self._fc = Linear(3, 10)
211 212 213 214 215

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

                with fluid.dygraph.guard():
216
                    layer = ExampleLayer()
217 218 219
                    in_np = np.random.random([2, 3]).astype('float32')
                    in_var = to_variable(in_np)
                    out_dygraph, static_layer = TracedLayer.trace(layer, inputs=[in_var])
220 221 222 223 224 225 226 227 228

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

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

                    # save the static graph model for inference
                    static_layer.save_inference_model(dirname='./saved_infer_model')
229
        """
230 231
        outs, prog, feed, fetch, parameters = _trace(layer, inputs)
        traced = TracedLayer(prog, parameters, feed, fetch)
232 233 234 235 236 237 238
        return outs, traced

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

        Args:
239
            build_strategy (BuildStrategy, optional): build strategy of
240 241 242 243 244 245 246 247 248 249 250
                :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:

                import paddle.fluid as fluid
251
                from paddle.fluid.dygraph import Linear, to_variable, TracedLayer
252 253 254
                import numpy as np

                class ExampleLayer(fluid.dygraph.Layer):
255 256 257
                    def __init__(self):
                        super(ExampleLayer, self).__init__()
                        self._fc = Linear(3, 10)
258 259 260 261 262

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

                with fluid.dygraph.guard():
263
                    layer = ExampleLayer()
264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296
                    in_np = np.random.random([2, 3]).astype('float32')
                    in_var = to_variable(in_np)

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

                    build_strategy = fluid.BuildStrategy()
                    build_strategy.enable_inplace = True

                    exec_strategy = fluid.ExecutionStrategy()
                    exec_strategy.num_threads = 2

                    static_layer.set_strategy(build_strategy=build_strategy, exec_strategy=exec_strategy)
                    out_static_graph = static_layer([in_var])
        """
        assert self._compiled_program is None, "Cannot set strategy after run"
        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):
297
                feed_dict[name] = x.value().get_tensor()
298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319
        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
    def save_inference_model(self, dirname, feed=None, fetch=None):
        """
320 321
        Save the TracedLayer to a model for inference. The saved
        inference model can be loaded by C++ inference APIs.
322 323

        Args:
324
            dirname (str): the directory to save the inference model.
325
            feed (list[int], optional): the input variable indices of the saved
326
                inference model. If None, all input variables of the
327 328 329 330 331 332 333 334
                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:
335
            None
336 337 338 339 340

        Examples:
            .. code-block:: python:

                import paddle.fluid as fluid
341
                from paddle.fluid.dygraph import Linear, to_variable, TracedLayer
342 343 344
                import numpy as np

                class ExampleLayer(fluid.dygraph.Layer):
345 346 347
                    def __init__(self):
                        super(ExampleLayer, self).__init__()
                        self._fc = Linear(3, 10)
348 349 350 351

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

352 353 354
                save_dirname = './saved_infer_model'
                in_np = np.random.random([2, 3]).astype('float32')

355
                with fluid.dygraph.guard():
356
                    layer = ExampleLayer()
357 358
                    in_var = to_variable(in_np)
                    out_dygraph, static_layer = TracedLayer.trace(layer, inputs=[in_var])
359
                    static_layer.save_inference_model(save_dirname, feed=[0], fetch=[0])
360 361

                place = fluid.CPUPlace()
362 363
                exe = fluid.Executor(place)
                program, feed_vars, fetch_vars = fluid.io.load_inference_model(save_dirname,
364
                                                    exe)
365 366 367

                fetch, = exe.run(program, feed={feed_vars[0]: in_np}, fetch_list=fetch_vars)
                print(fetch.shape) # (2, 10)
368
        """
369
        from paddle.fluid.io import save_inference_model
370 371 372 373 374

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

375
            return [all_vars[idx] for idx in partial_vars]
376 377 378 379 380 381 382 383 384 385

        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)

386
            save_inference_model(
387 388 389 390 391
                dirname=dirname,
                feeded_var_names=feeded_var_names,
                target_vars=target_vars,
                executor=self._exe,
                main_program=self._program.clone())