compiler.py 15.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
#   Copyright (c) 2018 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.

import multiprocessing
import os
import six
X
polish  
Xin Pan 已提交
18
import sys
19
from .. import compat as cpt
X
Xin Pan 已提交
20
from . import framework
S
sneaxiy 已提交
21
from .framework import cuda_places, cpu_places
22 23 24

from . import core

X
Xin Pan 已提交
25 26
__all__ = ['CompiledProgram', 'ExecutionStrategy', 'BuildStrategy']

27 28
ExecutionStrategy = core.ParallelExecutor.ExecutionStrategy
BuildStrategy = core.ParallelExecutor.BuildStrategy
F
flame 已提交
29 30
InferNativeConfig = core.NativeConfig
InferAnalysisConfig = core.AnalysisConfig
31 32 33 34 35 36 37 38


def _place_obj(place):
    p = core.Place()
    p.set_place(place)
    return p


39 40
def _is_pserver_mode(main_program):
    main = main_program if main_program \
C
chengduo 已提交
41
        else framework.default_main_program()
42 43 44 45 46 47
    for op in main.global_block().ops:
        if op.type in ["send", "recv"]:
            return True
    return False


C
chengduo 已提交
48 49 50 51 52 53 54 55
def _has_backward_op(graph):
    for node in graph.nodes():
        if node.is_op() and node.op() is not None and \
                node.op().type().endswith("_grad"):
            return True
    return False


56 57 58 59 60 61 62 63 64
def _prune_feed_ops(program):
    # prune the feed ops in the program.
    pop_idx = []
    for i, op in enumerate(program.global_block().ops):
        if op.type == "feed": pop_idx.append(i)
    for index in pop_idx[::-1]:
        program.global_block()._remove_op(index)


X
polish  
Xin Pan 已提交
65
class CompiledProgram(object):
X
polish  
Xin Pan 已提交
66
    """
X
Xin Pan 已提交
67
    Compiles to Graph for execution.
X
polish  
Xin Pan 已提交
68

X
Xin Pan 已提交
69 70 71 72
    1. Users first create the program with layers.
    2. Optionally, users use CompiledProgram to optimize the program before run.
    3. The original program or CompiledProgram is run by executor.

X
polish  
Xin Pan 已提交
73 74 75 76
    The CompiledProgram is used to transform a program for various
    optimizations, for example.
      * Pre-compute some logic once so that each run is faster.
      * Transform the program so that it can run in multiple devices.
77 78
      * Transform the program for optimized inference or distributed
        training. **Note that: this part is not finished.**
X
polish  
Xin Pan 已提交
79 80

    Example:
X
Xin Pan 已提交
81
        .. code-block:: python
82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104

          import paddle.fluid as fluid
          import paddle.fluid.compiler as compiler
          import numpy
          import os

          place = fluid.CUDAPlace(0) # fluid.CPUPlace()
          exe = fluid.Executor(place)

          data = fluid.layers.data(name='X', shape=[1], dtype='float32')
          hidden = fluid.layers.fc(input=data, size=10)
          loss = fluid.layers.mean(hidden)
          fluid.optimizer.SGD(learning_rate=0.01).minimize(loss)

          fluid.default_startup_program().random_seed=1
          exe.run(fluid.default_startup_program())
          compiled_prog = compiler.CompiledProgram(
                   fluid.default_main_program())

          x = numpy.random.random(size=(10, 1)).astype('float32')
          loss_data, = exe.run(compiled_prog,
                               feed={"X": x},
                               fetch_list=[loss.name])
X
polish  
Xin Pan 已提交
105 106

    Args:
X
Xin Pan 已提交
107 108 109 110 111
        program_or_graph (Graph|Program): If it's Program, it will be first
            lowered to a graph for further optimizations. If it's a graph
            (potentially optimized before), it will be directly used for
            further optimizations. Note: graph is only supported when compiled
            with with_data_parallel option.
C
chengduo 已提交
112 113 114 115
        build_strategy(BuildStrategy): build_strategy is used to
            build the graph with the specified options.
            For more information, please refer to fluid.BuildStrategy.
            Default None.
X
polish  
Xin Pan 已提交
116 117
    """

C
chengduo 已提交
118
    def __init__(self, program_or_graph, build_strategy=None):
X
Xin Pan 已提交
119 120
        if isinstance(program_or_graph, core.Graph):
            self._graph = program_or_graph
121
            # don't not create a new program here.
X
Xin Pan 已提交
122 123
            self._program = None
        elif isinstance(program_or_graph, framework.Program):
124
            _prune_feed_ops(program_or_graph)
X
Xin Pan 已提交
125 126 127 128 129 130
            self._graph = core.Graph(program_or_graph.desc)
            self._program = program_or_graph
        else:
            raise ValueError("Wrong program_to_graph type: %s" %
                             type(program_or_graph))

X
polish  
Xin Pan 已提交
131 132 133
        self._scope = None
        self._place = None
        self._executor = None
134 135
        self._compiled = False
        self._is_data_parallel = False
F
flame 已提交
136
        self._is_inference = False
C
chengduo 已提交
137 138 139 140 141
        self._loss_name = None
        self._share_vars_from = None
        self._places = None
        self._build_strategy = build_strategy
        self._exec_strategy = None
142

X
Xin Pan 已提交
143 144 145 146
    def with_data_parallel(self,
                           loss_name=None,
                           build_strategy=None,
                           exec_strategy=None,
S
sneaxiy 已提交
147 148
                           share_vars_from=None,
                           places=None):
X
Xin Pan 已提交
149 150
        """Configs the program to run in data parallel way.

151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188
        Example:
            .. code-block:: python

              import paddle.fluid as fluid
              import paddle.fluid.compiler as compiler
              import numpy
              import os

              use_cuda = True
              place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()

              # NOTE: If you use CPU to run the program, you need
              # to specify the CPU_NUM, otherwise, fluid will use
              # all the number of the logic core as the CPU_NUM,
              # in that case, the batch size of the input should be
              # greater than CPU_NUM, if not, the process will be
              # failed by an exception.
              if not use_cuda:
                  os.environ['CPU_NUM'] = str(2)

              exe = fluid.Executor(place)

              data = fluid.layers.data(name='X', shape=[1], dtype='float32')
              hidden = fluid.layers.fc(input=data, size=10)
              loss = fluid.layers.mean(hidden)
              fluid.optimizer.SGD(learning_rate=0.01).minimize(loss)

              fluid.default_startup_program().random_seed=1
              exe.run(fluid.default_startup_program())
              compiled_prog = compiler.CompiledProgram(
                       fluid.default_main_program()).with_data_parallel(
                                loss_name=loss.name)

              x = numpy.random.random(size=(10, 1)).astype('float32')
              loss_data, = exe.run(compiled_prog,
                                   feed={"X": x},
                                   fetch_list=[loss.name])

X
Xin Pan 已提交
189 190 191
        Args:
            loss_name (str): The loss name must set in training. Default None.
            build_strategy(BuildStrategy): build_strategy is used to
C
chengduo 已提交
192
                build the graph with the specified options.
X
Xin Pan 已提交
193
                For more information, please refer to fluid.BuildStrategy.
C
chengduo 已提交
194 195 196
                Note that, if you set build_strategy in the argument list when
                creating CompiledProgram and calling with_data_parallel,
                the build_strategy in CompiledProgram will be overwritten by the latter.
X
Xin Pan 已提交
197 198 199 200 201 202
                Default None.
            exec_strategy(ExecutionStrategy): exec_strategy is used to
                to select the a way to execute the graph, for example how many
                threads are used, how many iterations to clean up the temp
                variables. For more information, please refer
                to fluid.ExecutionStrategy. Default None.
S
sneaxiy 已提交
203
            share_vars_from(CompiledProgram): If provided, this CompiledProgram
X
Xin Pan 已提交
204 205 206
                will share variables from `share_vars_from`. `share_vars_from`
                must be run by the executor before this CompiledProgram so that
                vars are ready.
S
sneaxiy 已提交
207
            places(list(CUDAPlace)|list(CPUPlace)|None): If provided, only compile
S
sneaxiy 已提交
208 209 210
                program in the given places. Otherwise, the places used when compiled 
                is determined by the Executor, and the places used are controlled 
                by environment variables: FLAGS_selected_gpus or CUDA_VISIBLE_DEVICES
S
sneaxiy 已提交
211 212 213
                if using GPU; or CPU_NUM if using CPU. For example, if you want to 
                run on GPU 0 and 1, set places=[fluid.CUDAPlace(0), fluid.CUDAPlace(1)].
                If you want to run on 2 CPU cores, set places=[fluid.CPUPlace()]*2.  
S
sneaxiy 已提交
214

X
Xin Pan 已提交
215 216 217
        Returns:
            self
        """
218
        assert not self._is_data_parallel, "Already compiled with parallel."
X
Xin Pan 已提交
219
        assert not self._is_inference, "Cannot compile both data parallel and inference"
220
        self._is_data_parallel = True
C
chengduo 已提交
221 222 223 224 225
        # FIXME(zcd): Currently, the build_strategy can be set during creating
        # CompiledProgram or calling with_data_parallel, and it may be confusing,
        # but in the long run, we should set up build_strategy only when creating
        # CompiledProgram, and exec_strategy should be deprecated.
        if build_strategy is not None: self._build_strategy = build_strategy
226 227
        self._exec_strategy = exec_strategy
        self._loss_name = loss_name
X
polish  
Xin Pan 已提交
228
        self._share_vars_from = share_vars_from
C
chengduo 已提交
229 230 231 232 233 234 235 236 237
        self._places = places

        if _has_backward_op(self._graph):
            assert self._loss_name is not None, "The loss_name should be set here."

        if self._places is not None:
            if not isinstance(self._places, (list, tuple)):
                self._places = [self._places]

238 239
        return self

F
flame 已提交
240
    def _with_inference_optimize(self, config):
F
flame 已提交
241 242 243 244 245 246 247
        """ Add inference optimize

        Args:
            config: instance of `NativeConfig` or `AnalysisConfig` to create predictor
        Returns:
            self
        """
X
Xin Pan 已提交
248
        assert not self._is_data_parallel, "Cannot compile both data parallel and inference"
X
Xin Pan 已提交
249 250
        assert not self._is_inference, "Already compiled with inference"

F
flame 已提交
251 252 253 254 255 256 257
        assert any([
            isinstance(config, InferNativeConfig),
            isinstance(config, InferAnalysisConfig)
        ])
        self._is_inference = True
        self._infer_config = config
        return self
X
polish  
Xin Pan 已提交
258

F
flame 已提交
259
    def _with_distributed(self):
X
polish  
Xin Pan 已提交
260 261
        raise NotImplementedError()

C
chengduo 已提交
262
    def _compile_data_parallel(self, places, use_cuda=False, scope=None):
X
polish  
Xin Pan 已提交
263
        if self._share_vars_from:
264
            if scope:
X
polish  
Xin Pan 已提交
265
                sys.stderr.write("share_vars_from is set, scope is ignored.\n")
C
chengduo 已提交
266 267 268
            if not self._is_data_parallel:
                raise ValueError(
                    "Currently, only data parallel mode need share_vars_from.")
X
polish  
Xin Pan 已提交
269 270 271 272 273 274 275 276 277
            if not self._share_vars_from._is_data_parallel:
                raise ValueError("share_vars_from is not data parallel. Cannot "
                                 "share vars from it.")
            if self._share_vars_from._executor is None:
                raise ValueError(
                    "share_vars_from is not compiled and run, so there is no "
                    "var to share.")
            self._local_scopes = self._share_vars_from._executor.local_scopes()
        else:
278
            assert scope is not None, ""
X
polish  
Xin Pan 已提交
279
            self._local_scopes = []
280

C
chengduo 已提交
281 282 283 284 285 286 287 288 289 290
        assert isinstance(places, tuple) or isinstance(places, list), \
            "Currently , The places type only should be list or tuple, \n" \
            "but the input type is {}.".format(type(places))

        if self._build_strategy is None:
            self._build_strategy = BuildStrategy()
        self._build_strategy.is_distribution = _is_pserver_mode(self._program)

        if self._exec_strategy is None:
            self._exec_strategy = ExecutionStrategy()
S
sneaxiy 已提交
291
        self._exec_strategy.use_cuda = use_cuda
292 293 294 295 296

        if self._exec_strategy.num_threads == 0:
            if self._exec_strategy.use_cuda:
                # Experiments on se-resnext shows that too many threads hurt
                # performance. Worth tunning for other models in the future.
C
chengduo 已提交
297
                self._exec_strategy.num_threads = len(places) * 4
298
            else:
C
chengduo 已提交
299 300 301 302 303 304
                self._exec_strategy.num_threads = len(places) * 2

        if self._build_strategy.num_trainers > 1:
            assert self._is_data_parallel, \
                "If you use multi-trainer to train the model, you should use "\
                "the data parallel model, i.e. calling with_data_parallel function."
305

X
Xin Pan 已提交
306 307
        # TODO(wuyi): trainer endpoings should be passed in through
        # build_strategy, not program.xxx.
308
        # TODO(gongwb): let user to set them once.
X
Xin Pan 已提交
309 310 311
        if self._program and self._build_strategy.num_trainers > 1 and \
                self._program._trainers_endpoints:
            tps = self._program._trainers_endpoints
D
dzhwinter 已提交
312

313
            assert self._build_strategy.num_trainers == len(
X
Xin Pan 已提交
314 315 316
                tps), "num_trainers == len(end_points)"
            self._build_strategy.trainers_endpoints = tps

317 318
        if self._program:
            self._build_strategy.nccl_comm_num = self._program._nccl_comm_num
319 320
            self._build_strategy.use_hierarchical_allreduce = self._program._use_hierarchical_allreduce
            self._build_strategy.hierarchical_allreduce_inter_nranks = self._program._hierarchical_allreduce_inter_nranks
321

Q
qingqing01 已提交
322 323 324
        if self._build_strategy.sync_batch_norm:
            self._build_strategy.enable_sequential_execution = True

X
Xin Pan 已提交
325
        self._persistable_vars = []
Z
Zhen Wang 已提交
326 327 328 329
        for node in self._graph.nodes():
            if node.is_var() and node.var() is not None and node.var().persistable() and \
                    node.var().type() != core.VarDesc.VarType.RAW:
                self._persistable_vars.append(cpt.to_text(node.name()))
330

C
chengduo 已提交
331 332
        places = list(map(_place_obj, places))

Y
Yan Xu 已提交
333 334 335 336 337 338 339 340 341 342 343
        # ParallelExecutor would broadcast all the parameters during initializing.
        # The parameters of each process should be in the same ordered for the data-parallelism
        # distributed training to keep the broadcast correct.
        self._persistable_vars = list(set(self._persistable_vars))
        self._persistable_vars.sort()

        return core.ParallelExecutor(
            places, self._persistable_vars,
            cpt.to_text(self._loss_name)
            if self._loss_name else six.u(''), self._scope, self._local_scopes,
            self._exec_strategy, self._build_strategy, self._graph)
344

F
flame 已提交
345 346 347
    def _compile_inference(self):
        return core.create_paddle_predictor(self._infer_config)

348
    def _compile(self, scope, place):
X
Xin Pan 已提交
349 350 351 352 353 354 355 356 357 358
        """Compile the program based on the configs.

        Args:
            scope: The variables (resources) that are associated with
               this compiled program.
            place: The location that the compiled program will be run on.

        Returns:
            self
        """
359
        if self._compiled:
X
polish  
Xin Pan 已提交
360 361
            if scope and self._scope != scope:
                raise ValueError("Cannot compile with different scope")
S
sneaxiy 已提交
362
            if place and not self._place._equals(place):
X
polish  
Xin Pan 已提交
363
                raise ValueError("Cannot compile with different place")
364
            return self
X
fix  
Xin Pan 已提交
365
        self._compiled = True
366 367 368

        self._scope = scope
        self._place = place
C
chengduo 已提交
369 370

        if self._is_inference:
F
flame 已提交
371
            self._executor = self._compile_inference()
372
        else:
C
chengduo 已提交
373 374 375 376 377 378 379 380
            if self._is_data_parallel:
                self._places = self._get_places(self._place, self._places)
            else:
                self._places = [self._place]
            self._executor = self._compile_data_parallel(
                use_cuda=isinstance(self._place, core.CUDAPlace),
                scope=self._scope,
                places=self._places)
381
        return self
C
chengduo 已提交
382 383 384 385 386 387 388 389 390 391 392 393

    def _get_places(self, place, place_list):
        has_set_place = (place_list is not None)
        if has_set_place:
            for p in place_list:
                assert p._type() == place._type(), \
                    "Place type not match. You may set the wrong type of places"
        else:
            place_list = cuda_places() if isinstance(
                place, core.CUDAPlace) else cpu_places()
        assert place_list, "no place for execution"
        return place_list