optimizer.py 323.4 KB
Newer Older
1
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13
# 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.
14 15

from __future__ import print_function
16

17
import numpy as np
18
import six
19
import os
20
import logging
21
from collections import defaultdict
22

23
import paddle
Q
Qiao Longfei 已提交
24
from paddle.fluid.distribute_lookup_table import find_distributed_lookup_table
25
from paddle.fluid.framework import Program, Variable, Parameter, name_scope, default_main_program, default_startup_program, device_guard
26

27 28
from . import framework
from . import layers
29
from . import unique_name
30
from .backward import append_backward, _some_in_set_, _append_grad_suffix_, _get_no_grad_set_name
31
from .clip import GradientClipBase, GradientClipByNorm, error_clip_callback, append_gradient_clip_ops, ClipGradByGlobalNorm
32 33 34
from .framework import program_guard
from .initializer import Constant
from .layer_helper import LayerHelper
S
sneaxiy 已提交
35
from .layers import ops
36
from .dygraph import base as imperative_base
37
from .dygraph import no_grad
38
from .dygraph.learning_rate_scheduler import LearningRateDecay, _LearningRateEpochDecay
39 40 41
from paddle.fluid import core
from paddle.fluid.layers import tensor
from functools import reduce
42
from functools import cmp_to_key
43
from .wrapped_decorator import signature_safe_contextmanager
M
mapingshuo 已提交
44
from .. import compat as cpt
45
import warnings
W
wanghuancoder 已提交
46
from paddle import _C_ops
Z
zyfncg 已提交
47
from ..fluid.framework import _in_legacy_dygraph, in_dygraph_mode
48

49
__all__ = [
50 51 52 53
    'SGD', 'Momentum', 'Adagrad', 'Adam', 'Adamax', 'Dpsgd', 'DecayedAdagrad',
    'Ftrl', 'SGDOptimizer', 'MomentumOptimizer', 'AdagradOptimizer',
    'AdamOptimizer', 'AdamaxOptimizer', 'DpsgdOptimizer',
    'DecayedAdagradOptimizer', 'RMSPropOptimizer', 'FtrlOptimizer', 'Adadelta',
Z
Zeng Jinle 已提交
54
    'AdadeltaOptimizer', 'ModelAverage', 'LarsMomentum',
55 56
    'LarsMomentumOptimizer', 'LambOptimizer', 'ExponentialMovingAverage',
    'PipelineOptimizer', 'LookaheadOptimizer', 'RecomputeOptimizer'
57
]
Q
Qiao Longfei 已提交
58 59 60 61 62 63


class Optimizer(object):
    """Optimizer Base class.

    Define the common interface of an optimizer.
64 65
    User should not use this class directly,
    but need to use one of it's implementation.
Q
Qiao Longfei 已提交
66 67
    """

68
    @imperative_base.no_grad
69 70 71 72
    def __init__(self,
                 learning_rate,
                 parameter_list=None,
                 regularization=None,
73
                 grad_clip=None,
74 75
                 flatten_param_grads=False,
                 align_size=-1,
76
                 name=None):
77 78 79 80 81 82
        """
        Args:
            flatten_param_grads (bool, optional): Whether to flatten all the parameters and grads. 
                If true, the parameters and gradients will be coalesce to contiguous mempry, 
                and the grad_clip ops / optimizer ops will be fuse to one operator.
        """
83
        # Because of the loop import, so place it in the function body
84
        from paddle.optimizer.lr import LRScheduler
H
hong 已提交
85 86
        self._parameter_list = list(
            parameter_list) if parameter_list is not None else None
87
        self._name = name
J
Jiabin Yang 已提交
88
        if framework._non_static_mode():
89
            if not isinstance(learning_rate,
90
                              (float, LearningRateDecay, LRScheduler)):
M
minqiyang 已提交
91
                raise TypeError(
92
                    "learning rate should be float or LRScheduler, got %s here"
M
minqiyang 已提交
93
                    % type(learning_rate))
94
            if self._parameter_list is None:
95 96 97
                raise AttributeError(
                    "parameter_list argument given to the Optimizer should not be None in dygraph mode."
                )
98 99 100 101 102 103 104 105
            if regularization is not None:
                for param in self._parameter_list:
                    if param.regularizer is not None:
                        logging.info(
                            "If regularizer of a Parameter has been set by 'fluid.ParamAttr' or 'fluid.WeightNormParamAttr' already. "
                            "The Regularization[%s] in Optimizer will not take effect, and it will only be applied to other Parameters!"
                            % regularization.__str__())
                        break
M
minqiyang 已提交
106
        else:
107
            if not isinstance(learning_rate,
108
                              (float, framework.Variable, LRScheduler)):
M
minqiyang 已提交
109
                raise TypeError(
110
                    "learning rate should be float or LRScheduler, got %s here"
111
                    % type(learning_rate))
M
minqiyang 已提交
112

113 114 115 116 117
        if grad_clip is not None:
            if not isinstance(grad_clip, GradientClipBase):
                raise TypeError(
                    "'grad_clip' should be an instance of GradientClipBase's derived class"
                )
D
dzhwinter 已提交
118
        self.regularization = regularization
119
        self._grad_clip = grad_clip
120
        self._learning_rate = learning_rate
121 122
        self._flatten_param_grads = flatten_param_grads
        self._align_size = align_size
L
Leo Chen 已提交
123

D
dzhwinter 已提交
124
        self._dtype = None
L
Leo Chen 已提交
125 126 127 128
        # Infer the dtype form parameter
        if self._parameter_list:
            self._dtype = self._parameter_list[0].dtype

129
        # each program should have a independent learning rate
130
        # program -> Variable(learning_rate)
Q
qiaolongfei 已提交
131
        self._learning_rate_map = dict()
132
        if isinstance(self._learning_rate, framework.Variable):
133 134
            self._learning_rate_map[
                framework.default_main_program()] = self._learning_rate
135 136 137 138 139
        # Dictionary of accumulators. Some optimizer subclasses need to
        # allocate and manage extra variables associated with the parameters
        # to train. These variables are called accumulators.
        # {accum_name : { paramter_name : accumulator_for_parameter, ...}, ...}
        self._accumulators = defaultdict(lambda: dict())
140 141
        # global_accumulator dict, {accum_name : acc_variable, ...}
        self._global_accumulators = {}
142
        self.helper = LayerHelper(self.__class__.__name__)
143
        self._opti_name_list = []
H
hong 已提交
144
        self._accumulators_holder = {}
145
        self._param_device_map = dict()
146 147
        # NOTE(zhiqiu): sometimes we want to add some variables(Tenosr) to the optimizer for a specific optimization,
        # for example, we want to pass 'found_inf' to adam optimizer so it can skip update when found_inf is True.
148
        # And these variables should not be the parameters of Optimizer's construnctor (because not commonly used).
149 150
        # Use _auxiliary_vars together with _set_auxiliary_var/_get_auxiliary_var to achieve that.
        self._auxiliary_vars = dict()
H
hong 已提交
151 152 153 154

    @framework.dygraph_only
    def state_dict(self):
        '''
T
tianshuo78520a 已提交
155 156
        Get state dict information from optimizer. It contain all the variable used by optimizer. For Adam optimizer, contains beta1, beta2, momentum etc. If LearningRateDecay have been used, global_step will be include in state dict.
        If the optimizer never be called(minimize function), the state_dict is empty.
H
hong 已提交
157 158 159

        Args: None
        Return:
T
tianshuo78520a 已提交
160
            state_dict(dict) : dict contains all the variable used by optimizer
H
hong 已提交
161 162 163 164 165
        
        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
166 167 168 169 170 171

                with fluid.dygraph.guard():
                    emb = fluid.dygraph.Embedding([10, 10])

                    adam = fluid.optimizer.Adam(0.001, parameter_list=emb.parameters())
                    state_dict = adam.state_dict()
H
hong 已提交
172 173

        '''
174
        from paddle.optimizer.lr import LRScheduler
H
hong 已提交
175 176 177 178
        state_dict = {}
        for k, v in self._accumulators.items():
            for para_name, var_tmp in v.items():
                state_dict[var_tmp.name] = var_tmp
179 180
        for k, v in self._global_accumulators.items():
            state_dict[v.name] = v
H
hong 已提交
181
        # global step if use lr decay
182
        if isinstance(self._learning_rate, LRScheduler):
183 184
            state_dict["LR_Scheduler"] = self._learning_rate.state_dict()
            return state_dict
H
hong 已提交
185
        if isinstance(self._learning_rate, LearningRateDecay):
186 187 188 189
            state_dict["LR_Scheduler"] = self._learning_rate.state_dict()

            if not isinstance(self._learning_rate, _LearningRateEpochDecay):
                var_tmp = None
190 191 192
                var_temp = framework._varbase_creator(None,
                                                      name='global_step',
                                                      dtype='int32')
193

194 195 196 197
                tensor.fill_constant([1],
                                     "int32",
                                     self._learning_rate.step_num,
                                     out=var_temp)
H
hong 已提交
198

199
                state_dict['global_step'] = var_temp
H
hong 已提交
200 201 202
        return state_dict

    @framework.dygraph_only
203
    def set_state_dict(self, state_dict):
H
hong 已提交
204
        '''
T
tianshuo78520a 已提交
205
        Load optimizer state dict. For Adam optimizer, contains beta1, beta2, momentum etc. If LearningRateDecay have been used, global_step will be changed.
H
hong 已提交
206 207 208 209 210 211 212 213

        Args: 
            state_dict(dict) : Dict contains all the Variable needed by optimizer
        Return:
            None
        
        Examples:
            .. code-block:: python
214

215 216
                import paddle
                import paddle.fluid as fluid
217 218 219

                paddle.disable_static()

220
                emb = paddle.nn.Embedding(10, 10)
221

222
                state_dict = emb.state_dict()
223
                fluid.save_dygraph(state_dict, "paddle_dy")
224

225
                scheduler = paddle.optimizer.lr.NoamDecay(	
226 227 228 229
                    d_model=0.01, warmup_steps=100, verbose=True)
                adam = paddle.optimizer.Adam(
                    learning_rate=scheduler,
                    parameters=emb.parameters())
230
                state_dict = adam.state_dict()
231
                fluid.save_dygraph(state_dict, "paddle_dy")
232

233
                para_state_dict, opti_state_dict = fluid.load_dygraph("paddle_dy")
H
hong 已提交
234
        '''
235 236
        from paddle.optimizer.lr import LRScheduler
        if isinstance(self._learning_rate, LRScheduler):
237
            self._learning_rate.set_dict(state_dict["LR_Scheduler"])
H
hong 已提交
238 239

        if isinstance(self._learning_rate, LearningRateDecay):
240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261
            self._learning_rate.set_dict(state_dict["LR_Scheduler"])

            if not isinstance(self._learning_rate, _LearningRateEpochDecay):
                assert 'global_step' in state_dict, \
                        'Global step not in state dict, Dygraph use LearningRateDecay, global_step must in state_dict'
                global_step = state_dict['global_step']

                if isinstance(global_step, Variable):
                    step_np = global_step
                    step_np = np.array(step_np.value().get_tensor())
                    assert step_np.shape == (1,),  \
                            "global step shape is (1,), the shape is {}".format( step_np.shape )

                    self._learning_rate.step_num = int(step_np[0])
                elif isinstance(global_step, np.ndarray):
                    assert global_step.shape == (1,),  \
                            "global step shape is (1,), the shape is {}".format( global_step.shape )
                    self._learning_rate.step_num = global_step[0]
                else:
                    raise RuntimeError(
                        "Type not supprt, value in state dict must be [VarBase, Variable, numpy], the type is ",
                        type(global_step))
H
hong 已提交
262

263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279
        def _load_state_para(state_dict, param):
            var = param.value()
            tensor = var.get_tensor()
            model_np = np.array(tensor)
            load_para = state_dict[param.name]
            if isinstance(load_para, Variable):
                load_para_np = load_para.numpy()
            elif isinstance(load_para, core.VarBase):
                load_para_np = load_para.numpy()
            elif isinstance(load_para, np.ndarray):
                load_para_np = load_para
            else:
                raise RuntimeError("State dict type {} not supprt".format(
                    str(type(load_para))))

            assert model_np.shape == load_para_np.shape,  \
                                        "Parameter shape not match, Dygraph Parameter [ {} ] need tensor with shape {} but load tensor with shape {}".format(
280
                                                param.name, model_np.shape, load_para_np.shape)
281 282 283

            assert model_np.dtype == load_para_np.dtype, \
                                        "Parameter dtype not match, Dygraph Parameter [ {} ] need tensor with dtype {}  but load tensor with dtype {}".format(
284
                                            param.name, model_np.dtype, load_para_np.dtype)
285 286 287

            tensor.set(load_para_np, framework._current_expected_place())

H
hong 已提交
288 289 290 291 292
        self._accumulators_holder = state_dict
        for k, v in self._accumulators.items():
            for para_name, var_tmp in v.items():
                assert var_tmp.name in state_dict, \
                        "optimizer variable {} not found".format( var_tmp.name )
293
                _load_state_para(state_dict, var_tmp)
H
hong 已提交
294

295 296 297 298
        for k, v in self._global_accumulators.items():
            assert v.name in state_dict, \
                        "optimizer variable {} not found".format( v.name )
            _load_state_para(state_dict, v)
299

300 301 302
    # [aliases] Compatible with old method names
    set_dict = set_state_dict

303 304
    def get_opti_var_name_list(self):
        return self._opti_name_list
Q
Qiao Longfei 已提交
305

306 307 308 309 310 311 312 313 314
    def _set_auxiliary_var(self, key, val):
        self._auxiliary_vars[key] = val

    def _get_auxiliary_var(self, key):
        if key in self._auxiliary_vars:
            return self._auxiliary_vars[key]
        else:
            return None

Q
Qiao Longfei 已提交
315
    def _create_global_learning_rate(self):
316 317
        from paddle.optimizer.lr import LRScheduler
        if isinstance(self._learning_rate, LRScheduler):
318 319 320 321 322 323 324 325 326 327 328 329 330 331
            lr_var = self._global_learning_rate()
            # only create global lr_var once
            if not isinstance(lr_var, framework.Variable):
                lr_name = unique_name.generate('learning_rate')
                self._learning_rate._var_name = lr_name
                lr_var = self.helper.create_global_variable(
                    name=lr_name,
                    shape=[1],
                    persistable=True,
                    stop_gradient=True,
                    dtype='float32' if self._dtype is None else self._dtype)
                main_prog = framework.default_main_program()
                main_prog.lr_sheduler = self._learning_rate
                main_prog.lr_var = lr_var
332 333
                self._learning_rate_map[
                    framework.default_main_program()] = lr_var
334 335 336 337 338 339

            lr_value = float(self._learning_rate())
            self.helper.set_variable_initializer(
                lr_var, initializer=Constant(value=lr_value))
            return

340 341 342
        if imperative_base.enabled():
            # create learning rate Variable
            if isinstance(self._learning_rate, float):
M
minqiyang 已提交
343 344 345 346 347 348 349 350 351 352 353 354
                lr = self._global_learning_rate()

                if isinstance(lr, framework.Variable):
                    return
                else:
                    self._learning_rate_map[framework.default_main_program(
                    )] = layers.create_global_var(
                        name=unique_name.generate("learning_rate"),
                        shape=[1],
                        value=float(self._learning_rate),
                        dtype='float32' if self._dtype is None else self._dtype,
                        persistable=True)
355
            # get learning rate Variable from LearningRateDecay
M
minqiyang 已提交
356
            elif isinstance(self._learning_rate, LearningRateDecay):
357 358
                self._learning_rate_map[
                    framework.default_main_program()] = self._learning_rate()
359
            else:
Q
qiaolongfei 已提交
360
                raise TypeError(
361 362
                    "optimizer's learning rate must be float or LearningRateDecay"
                )
363
        else:
364 365 366 367
            lr = self._global_learning_rate()

            if isinstance(lr, framework.Variable):
                return
M
minqiyang 已提交
368 369 370 371 372 373
            else:
                if not isinstance(self._learning_rate, float):
                    raise TypeError(
                        "learning rate variable is create outside optimizer,"
                        "can not create new learning rate variable for new program"
                    )
Q
Qiao Longfei 已提交
374

375
            # create learning rate in the current main program
376 377 378 379 380 381 382
            self._learning_rate_map[
                framework.default_main_program()] = layers.create_global_var(
                    name=unique_name.generate("learning_rate"),
                    shape=[1],
                    value=float(self._learning_rate),
                    dtype='float32' if self._dtype is None else self._dtype,
                    persistable=True)
383

384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445
    @framework.dygraph_only
    def set_lr(self, value):
        """
        :api_attr: imperative
        
        Set the value of the learning rate manually in the optimizer. If the optimizer use LearningRateDecay,
        this API cannot be invoked, because it will lead to conflict.

        Args:
            value (float|Variable): the value of learning rate

        Returns:
            None
          
        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
                        
                with fluid.dygraph.guard():
                    linear = fluid.dygraph.nn.Linear(10, 10)

                    adam = fluid.optimizer.Adam(0.1, parameter_list=linear.parameters())

                    # set learning rate manually by python float value
                    lr_list = [0.2, 0.3, 0.4, 0.5, 0.6]
                    for i in range(5):
                        adam.set_lr(lr_list[i])
                        lr = adam.current_step_lr()
                        print("current lr is {}".format(lr))
                    # Print:
                    #    current lr is 0.2
                    #    current lr is 0.3
                    #    current lr is 0.4
                    #    current lr is 0.5
                    #    current lr is 0.6


                    # set learning rate manually by framework Variable
                    lr_var = fluid.layers.create_global_var(
                        shape=[1], value=0.7, dtype='float32')
                    adam.set_lr(lr_var)
                    lr = adam.current_step_lr()
                    print("current lr is {}".format(lr))
                    # Print:
                    #    current lr is 0.7



        """
        if not isinstance(value, (framework.Variable, float)):
            raise TypeError(
                "The type of 'value' in optimizer.set_lr must be (float, Variable), but received %s."
                % (type(value)))
        if isinstance(self._learning_rate, LearningRateDecay):
            raise RuntimeError(
                "optimizer's learning rate can't be LearningRateDecay when invoke this API, because this will lead to conflict."
            )
        if isinstance(value, float):
            self._learning_rate = value
            current_lr = self._global_learning_rate()
            if current_lr is not None:
J
Jiabin Yang 已提交
446
                if framework._non_static_mode():
447 448
                    _C_ops.fill_constant(current_lr, 'value', float(value),
                                         'dtype', current_lr.dtype, 'shape',
449 450 451 452
                                         list(current_lr.shape))
                else:
                    global_block = framework.default_main_program(
                    ).global_block()
453 454 455 456 457 458 459 460
                    global_block.append_op(type='fill_constant',
                                           outputs={'Out': [current_lr]},
                                           attrs={
                                               'dtype': current_lr.dtype,
                                               'shape': list(current_lr.shape),
                                               'value': float(value)
                                           },
                                           stop_gradient=True)
461 462 463 464 465
        else:
            assert len(value.shape) == 1 and value.shape[
                0] == 1, "optimizer's learning rate must be 1-D Tensor with shape[1]"
            self._learning_rate_map[framework.default_main_program()] = value

466 467 468
    @framework.dygraph_only
    def current_step_lr(self):
        """
469
        :api_attr: imperative
470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514
        
        Get current step learning rate. The return value is all the same When LearningRateDecay is not used,
        otherwise return the step learning rate.

        Returns:
            float: The learning rate of the current step.

        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
                import numpy as np

                # example1: LearningRateDecay is not used, return value is all the same
                with fluid.dygraph.guard():
                    emb = fluid.dygraph.Embedding([10, 10])
                    adam = fluid.optimizer.Adam(0.001, parameter_list = emb.parameters())
                    lr = adam.current_step_lr()
                    print(lr) # 0.001

                # example2: PiecewiseDecay is used, return the step learning rate
                with fluid.dygraph.guard():
                    inp = np.random.uniform(-0.1, 0.1, [10, 10]).astype("float32")
                    linear = fluid.dygraph.nn.Linear(10, 10)
                    inp = fluid.dygraph.to_variable(inp)
                    out = linear(inp)
                    loss = fluid.layers.reduce_mean(out)
                    
                    bd = [2, 4, 6, 8]
                    value = [0.2, 0.4, 0.6, 0.8, 1.0]
                    adam = fluid.optimizer.Adam(fluid.dygraph.PiecewiseDecay(bd, value, 0),
                                           parameter_list=linear.parameters())

                    # first step: learning rate is 0.2
                    np.allclose(adam.current_step_lr(), 0.2, rtol=1e-06, atol=0.0) # True

                    # learning rate for different steps
                    ret = [0.2, 0.2, 0.4, 0.4, 0.6, 0.6, 0.8, 0.8, 1.0, 1.0, 1.0, 1.0]
                    for i in range(12):
                        adam.minimize(loss)
                        lr = adam.current_step_lr()
                        np.allclose(lr, ret[i], rtol=1e-06, atol=0.0) # True

        """
        current_lr = self._global_learning_rate()
515
        if isinstance(current_lr, framework.Variable):
516 517 518 519
            return self._global_learning_rate().numpy()[0]

        if isinstance(self._learning_rate, float):
            return self._learning_rate
520 521 522
        elif isinstance(self._learning_rate, _LearningRateEpochDecay):
            step_lr = self._learning_rate()
            return step_lr.numpy()[0]
523 524 525 526 527 528 529
        else:
            step_lr = self._learning_rate.step()
            if isinstance(step_lr, (float, int)):
                return step_lr
            else:
                return step_lr.numpy()[0]

Y
yuyang18 已提交
530
    def _global_learning_rate(self, program=None):
Q
Qiao Longfei 已提交
531 532 533 534
        """
        get global decayed learning rate
        :return:
        """
535 536
        if program is None:
            program = framework.default_main_program()
Q
qiaolongfei 已提交
537
        return self._learning_rate_map.get(program, None)
Q
Qiao Longfei 已提交
538

Q
Qiao Longfei 已提交
539 540 541 542 543
    def _append_optimize_op(self, block, param_and_grad):
        """ append optimize operator to block and return all the added optimize_op
        """
        raise NotImplementedError()

544 545 546 547
    def _create_param_lr(self, param_and_grad):
        # create learning rate variable for every parameter
        param = param_and_grad[0]
        param_lr = param.optimize_attr['learning_rate']
W
Wu Yi 已提交
548 549
        if type(param_lr) == Variable:
            return param_lr
Q
qiaolongfei 已提交
550
        else:
W
Wu Yi 已提交
551
            if param_lr == 1.0:
Y
yuyang18 已提交
552
                return self._global_learning_rate()
W
Wu Yi 已提交
553
            else:
X
Xin Pan 已提交
554 555 556
                with default_main_program()._lr_schedule_guard(
                        is_with_opt=True), framework.name_scope(
                            'scale_with_param_lr'):
557
                    return self._global_learning_rate() * param_lr
558 559 560 561 562 563 564

    def _create_accumulators(self, block, parameters):
        """Create all accumulators needed by the parameters

        Args:
            block: the block in which the loss variable is present
            parameters: list of parameter variables for the optimizer
Q
Qiao Longfei 已提交
565
        """
566 567
        pass

568
    def _finish_update(self, block, parameters_and_grads):
569 570 571 572 573 574 575 576
        """Finish any custom updates needed
           before completing an optimization step

        Args:
            block: the block in which the loss variable is present
            parameters: list of parameter variables for the optimizer

        Returns:
Q
qiaolongfei 已提交
577
            None
578 579 580
        """
        pass

581 582 583 584 585
    def _add_accumulator(self,
                         name,
                         param,
                         dtype=None,
                         fill_value=0.0,
586
                         shape=None,
587
                         type=None,
588
                         device=None):
589 590 591 592 593 594 595 596 597
        """Utility function to add an accumulator for a parameter

        Args:
            block: the block in which the loss variable is present
            name: name of the accumulator
            param: parameter variable for which accumulator is to be added
            dtype: data type of the accumulator variable
            fill_value: value to initialize the accumulator variable
        """
W
whs 已提交
598 599
        if self._name is not None:
            name = self._name + "_" + name
600 601
        if (name in self._accumulators
                and param.name in self._accumulators[name]):
J
Jiabin Yang 已提交
602
            if framework._non_static_mode():
X
polish  
Xin Pan 已提交
603
                return self._accumulators[name][param.name]
604 605 606
            raise Exception(
                "Accumulator {} already exists for parameter {}".format(
                    name, param.name))
607 608
        if shape == None:
            shape = param.shape
Q
Qiao Longfei 已提交
609
        assert isinstance(self.helper, LayerHelper)
610 611 612 613 614

        var_name = param.name + "_" + name
        var_name = unique_name.generate(var_name)
        self._opti_name_list.append(var_name)

Q
Qiao Longfei 已提交
615
        var = self.helper.create_global_variable(
616
            name=var_name,
Q
Qiao Longfei 已提交
617
            persistable=True,
F
fengjiayi 已提交
618
            dtype=dtype or param.dtype,
619
            type=core.VarDesc.VarType.LOD_TENSOR
620 621
            if framework._non_static_mode() else
            (param.type if type is None else type),
H
hong 已提交
622 623
            shape=shape,
            belong_to_optimizer=True)
624 625 626 627 628
        if device is None:
            device = self._get_device_for_param(param.name)
        with device_guard(device):
            self.helper.set_variable_initializer(
                var, initializer=Constant(value=float(fill_value)))
H
hong 已提交
629

J
Jiabin Yang 已提交
630
        if framework._non_static_mode():
H
hong 已提交
631 632 633 634 635
            if len(self._accumulators_holder) > 0:
                assert var_name in self._accumulators_holder, \
                        "Optimizer set error, {} should in state dict".format( var_name )
                var.set_value(self._accumulators_holder[var_name])

Q
Qiao Longfei 已提交
636
        self._accumulators[name][param.name] = var
637
        return var
638

639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659
    def _add_global_accumulator(self,
                                name,
                                dtype=None,
                                fill_value=0.0,
                                shape=None,
                                type=None,
                                device=None):
        """Utility function to add a global accumulator for all parameters in the model

        Args:
            block: the block in which the loss variable is present
            name: name of the accumulator
            dtype: data type of the accumulator variable
            fill_value: value to initialize the accumulator variable
            shape: the shape of the accumulator
            type: the variable type of the accumulator
            device: the target place of the accumulator
        """
        if self._name is not None:
            name = self._name + "_" + name
        if (name in self._global_accumulators):
J
Jiabin Yang 已提交
660
            if framework._non_static_mode():
661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683
                return self._global_accumulators[name]
            raise Exception("Global accumulator {} already exists".format(name))
        if shape == None:
            shape = [1]  # most case, global accumulator is of shape [1]
        assert isinstance(self.helper, LayerHelper)

        var_name = name
        var_name = unique_name.generate(var_name)
        self._opti_name_list.append(var_name)

        var = self.helper.create_global_variable(
            name=var_name,
            persistable=True,
            dtype=dtype if dtype else self._dtype,
            type=type,
            shape=shape,
            belong_to_optimizer=True)
        if device is None:
            device = 'cpu'
        with device_guard(device):
            self.helper.set_variable_initializer(
                var, initializer=Constant(value=float(fill_value)))

J
Jiabin Yang 已提交
684
        if framework._non_static_mode():
685 686 687 688 689 690 691 692
            if len(self._accumulators_holder) > 0:
                assert var_name in self._accumulators_holder, \
                        "Optimizer set error, {} should in state dict".format( var_name )
                var.set_value(self._accumulators_holder[var_name])

        self._global_accumulators[name] = var
        return var

693 694 695 696 697 698 699 700
    def _get_accumulator(self, name, param):
        """Utility function to fetch an accumulator for a parameter

        Args:
            name: name of the accumulator
            param: parameter variable for which accumulator is to be fetched

        Returns:
701
            accumulator variable
702
        """
W
whs 已提交
703 704
        if self._name is not None:
            name = self._name + "_" + name
705 706 707 708 709
        if (name not in self._accumulators
                or param.name not in self._accumulators[name]):
            raise Exception(
                "Accumulator {} does not exist for parameter {}".format(
                    name, param.name))
710 711
        return self._accumulators[name][param.name]

712 713 714 715 716 717 718 719 720 721 722 723 724 725 726
    def _get_global_accumulator(self, name):
        """Utility function to fetch a global accumulator

        Args:
            name: name of the accumulator

        Returns:
            accumulator variable
        """
        if self._name is not None:
            name = self._name + "_" + name
        if (name not in self._global_accumulators):
            raise Exception("Global accumulator {} does not exist".format(name))
        return self._global_accumulators[name]

727 728 729 730 731 732 733 734 735 736 737 738
    def _update_param_device_map(self, parameters_and_grads, target_block):
        for param_and_grad in parameters_and_grads:
            if param_and_grad[0].trainable is True:
                param_name = param_and_grad[0].name
                ops = target_block.ops
                device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName(
                )
                for op in ops:
                    input_arg_names = op.input_arg_names
                    if param_name in input_arg_names:
                        self._param_device_map[param_name] = op.attr(
                            device_attr_name)
739
                        break
740 741 742 743 744 745 746

    def _get_device_for_param(self, param_name):
        device = None
        if param_name in self._param_device_map:
            device = self._param_device_map[param_name]
        return device

747
    def _create_optimization_pass(self, parameters_and_grads):
Q
Qiao Longfei 已提交
748 749 750
        """Add optimization operators to update gradients to variables.

        Args:
Q
qiaolongfei 已提交
751
          parameters_and_grads(list(tuple(Variable, Variable))):
752
            a list of (variable, gradient) pair to update.
Q
Qiao Longfei 已提交
753 754

        Returns:
755
          return_op_list: a list of operators that will complete one step of
756 757 758
            optimization. This will include parameter update ops, global step
            update ops and any other custom ops required by subclasses to manage
            their internal state.
Q
Qiao Longfei 已提交
759
        """
760 761 762 763 764
        # This is a default implementation of create_optimization_pass that
        # can be shared by most optimizers. This implementation assumes that
        # the subclass will implement the _append_optimize_op method and the
        #  _initialize_tensors method. The subclass can extend the
        # _create_accumulators method if it needs to create accumulators
765
        # for parameters and extend _finish_update method to add custom ops.
766

767
        # Allways called under program_guard use global block as loss block
768 769 770
        # But if current block is in control flow, append optimize op in the
        # grad block of current block

771
        global_block = framework.default_main_program().global_block()
772 773 774 775 776 777 778 779 780
        target_block = global_block
        current_block = framework.default_main_program().current_block()
        if current_block.idx != global_block.idx:
            assert current_block.backward_block_idx != -1, \
                "current block is not global_block, but it doesn't have backward block."
            target_block = framework.default_main_program().blocks[
                current_block.backward_block_idx]

        start = len(target_block.ops)
781

782
        self._update_param_device_map(parameters_and_grads, target_block)
C
chengduo 已提交
783
        self._create_accumulators(
784
            target_block,
C
chengduo 已提交
785
            [p[0] for p in parameters_and_grads if p[0].trainable])
786 787
        self._create_global_learning_rate()

J
Jiabin Yang 已提交
788
        if framework._non_static_mode():
789 790 791
            for param_and_grad in parameters_and_grads:
                if param_and_grad[1] is None:
                    continue
792 793
                if param_and_grad[0].trainable is True:
                    self._append_optimize_op(target_block, param_and_grad)
794 795 796 797 798 799 800
        else:
            for param_and_grad in parameters_and_grads:
                if param_and_grad[1] is None:
                    continue
                with param_and_grad[0].block.program._optimized_guard(
                        param_and_grad), name_scope("optimizer"):
                    if param_and_grad[0].trainable is True:
801 802
                        device = self._get_device_for_param(
                            param_and_grad[0].name)
803 804 805
                        with device_guard(device):
                            optimize_op = self._append_optimize_op(
                                target_block, param_and_grad)
806 807 808

        # Get custom finish ops for subclasses
        # FIXME: Need to fix this once we figure out how to handle dependencies
809
        self._finish_update(target_block, parameters_and_grads)
810

811 812
        end = len(target_block.ops)
        return target_block._slice_ops(start, end)
813 814

    def _process_distribute_lookuptable(self, param_grads):
Q
Qiao Longfei 已提交
815 816 817 818 819 820 821 822 823
        """
        Because distribute lookup table only support SGD optimizer for now, not support
        other optimizer and regularization, so we should find the table parameter out,
        and avoid to add regularization and other op for it, and add sgd optimize op
        for it independently.
        :param param_grads(list((Var, Var))): list of (param, grad) pair.
        :param loss: the loss variable.
        :param startup_program: the startup program
        """
824 825
        program = framework.default_main_program()
        global_block = framework.default_main_program().global_block()
Q
Qiao Longfei 已提交
826 827 828 829 830 831 832 833 834 835 836 837 838 839 840
        table_name = find_distributed_lookup_table(program)
        table_param = None
        table_grad = None
        new_param_grads = []
        for p, g in param_grads:
            if p.name == table_name:
                if table_param is not None:
                    raise RuntimeError(
                        "multi dist table var found, only support one now!")
                table_param = p
                table_grad = g
            else:
                new_param_grads.append((p, g))
        sgd_op = None
        if table_param is not None:
841 842 843 844 845 846 847 848 849 850 851 852 853
            param_and_grad = [table_param, table_grad]
            with table_param.block.program._optimized_guard(param_and_grad), \
                    framework.name_scope("optimizer"):
                self._create_global_learning_rate()
                # create the optimize op
                sgd_op = global_block.append_op(
                    type='sgd',
                    inputs={
                        "Param": table_param,
                        "Grad": table_grad,
                        "LearningRate": self._create_param_lr(param_and_grad)
                    },
                    outputs={"ParamOut": param_and_grad[0]})
Q
Qiao Longfei 已提交
854 855
        return new_param_grads, (table_param, table_grad), sgd_op

856 857 858 859 860 861 862
    def backward(self,
                 loss,
                 startup_program=None,
                 parameter_list=None,
                 no_grad_set=None,
                 callbacks=None):
        """
863
        The first part of ``minimize``, do auto-diff to append backward operations for
864 865 866
        the current program.

        Args:
867 868 869 870
            loss (Variable): ``loss`` variable to run optimizations.
            startup_program (Program, optional): :ref:`api_fluid_Program` for
                initializing parameters in ``parameter_list``. The default value
                is None, at this time :ref:`api_fluid_default_startup_program` will be used.
H
hong 已提交
871
            parameter_list (Iterable, optional): Iterable of ``Variable`` or ``Variable.name`` to update
872 873
                to minimize ``loss``. The default value is None, at this time all parameters
                will be updated.
874
            no_grad_set (set, optional): Set of ``Variable``  or ``Variable.name`` that don't need
875 876 877
                to be updated. The default value is None.
            callbacks (list, optional): list of callable objects to run when appending backward
                operator for one parameter. The default value is None.
M
minqiyang 已提交
878

879
        Return:
880 881
            list: list of (param, grad) variable pairs, param is ``Parameter``,
                grad is the gradient value corresponding to the parameter.
M
minqiyang 已提交
882

883
        Examples:
884
            See examples in ``apply_gradients``.
885
        """
886
        act_no_grad_set = None
J
Jiabin Yang 已提交
887
        if framework._non_static_mode():
888
            pass
L
Leo Chen 已提交
889 890
        else:
            act_no_grad_set = self._get_no_grad_set(loss, no_grad_set)
G
gongweibao 已提交
891

L
Leo Chen 已提交
892 893 894 895
        # Infer dtype by loss if None
        if self._dtype is None:
            self._dtype = loss.dtype

J
Jiabin Yang 已提交
896
        if framework._non_static_mode():
897 898 899
            parameter_list = parameter_list if parameter_list \
                else self._parameter_list

C
chengduo 已提交
900
            params_grads = []
901
            for param in parameter_list:
C
chengduo 已提交
902 903
                if not param.trainable:
                    continue
904
                if param._grad_ivar() is not None:
C
chengduo 已提交
905
                    # create gradient variable
906
                    grad_var = param._grad_ivar()
C
chengduo 已提交
907
                    params_grads.append((param, grad_var))
908
        else:
C
chengduo 已提交
909 910 911 912 913
            if callbacks is None:
                callbacks = [error_clip_callback]
            else:
                assert (isinstance(callbacks, list))
            program = loss.block.program
C
chengduo 已提交
914 915
            assert len(loss.shape) == 1 and loss.shape[0] == 1, \
                "The loss.shape should be (1L,), but the current loss.shape is {}. " \
916
                "Maybe that you should call paddle.mean to process the current loss.".format(
C
chengduo 已提交
917
                    loss.shape)
918 919
            parameter_list = parameter_list if parameter_list \
                else self._parameter_list
C
chengduo 已提交
920 921
            with program_guard(program, startup_program):
                params_grads = append_backward(loss, parameter_list,
922
                                               act_no_grad_set, callbacks)
C
chengduo 已提交
923
        return params_grads
924

925 926 927 928 929 930
    def _create_regularization_of_grad(self, param, grad, regularization=None):
        """ Create and add backward regularization Operators
    
        Function helper of append_regularization_ops.
        """
        # If no gradient or no regularization is specified,  then we don't need to do anything
931 932 933 934
        if grad is None or (
            (not hasattr(param, 'regularizer') or
             (hasattr(param, 'regularizer') and param.regularizer is None))
                and regularization is None):
935 936 937 938 939 940 941 942 943 944
            return grad
        regularization_term = None
        if hasattr(param, 'regularizer') and param.regularizer is not None:
            # Add variable for regularization term in grad block
            regularization_term = param.regularizer(param, grad, grad.block)
        elif regularization is not None:
            regularization_term = regularization(param, grad, grad.block)

        assert regularization_term is not None

J
Jiabin Yang 已提交
945
        if framework._non_static_mode():
W
wanghuancoder 已提交
946
            return _C_ops.sum([grad, regularization_term])
947

948 949 950 951 952 953 954 955 956 957 958 959 960 961 962
        new_grad = grad
        if grad.type == core.VarDesc.VarType.SELECTED_ROWS:
            # FIXME(zcd): If the grad is SELECTED_ROWS, after regularization,
            # the grad's type and name will be changed. But the gradient's name
            # is used in ParallelExecutor Reduce mode, so I add a flag for
            # the new_grad here.
            new_grad = grad.block.create_var(
                name=grad.name + core.kNewGradSuffix(),
                dtype=param.dtype,
                shape=param.shape,
                lod_level=param.lod_level,
                type=core.VarDesc.VarType.LOD_TENSOR)

        inputs = {"X": [grad, regularization_term]}
        outputs = {"Out": [new_grad]}
963
        grad.block.append_op(type='sum', inputs=inputs, outputs=outputs)
964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990

        return new_grad

    def append_regularization_ops(self,
                                  parameters_and_grads,
                                  regularization=None):
        r"""Create and add backward regularization Operators
    
        Creates and adds backward regularization operators in the BlockDesc.
        This will add gradients of the regularizer function to the gradients
        of the parameters and return these modified gradients. This is the
        same as implementing weight decay in optimizers for regularization.
    
        Args:
            parameters_and_grads: A list of (parameters, gradients) pairs
                                  that need to be regularized.
            regularization: A global regularizer. If the parameter is not
                            set. It will be applied with regularizer.
    
        Returns:
            list[(Variable, Variable)]: list of (parameters, gradients) \
            pair with the regularized gradient
    
        Raises:
            Exception: Unknown regularization type
        """
        params_and_grads = []
J
Jiabin Yang 已提交
991
        if framework._non_static_mode():
992
            for param, grad in parameters_and_grads:
993 994
                new_grad = self._create_regularization_of_grad(
                    param, grad, regularization)
995 996 997 998 999
                params_and_grads.append((param, new_grad))
        else:
            repeate_regularizer = False
            with framework.name_scope('regularization'):
                for param, grad in parameters_and_grads:
1000 1001 1002
                    if not repeate_regularizer and getattr(
                            param, 'regularizer',
                            None) is not None and regularization is not None:
1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013
                        repeate_regularizer = True
                        logging.info(
                            "If regularizer of a Parameter has been set by 'fluid.ParamAttr' or 'fluid.WeightNormParamAttr' already. "
                            "The Regularization[%s] in Optimizer will not take effect, and it will only be applied to other Parameters!"
                            % regularization.__str__())
                    with param.block.program._optimized_guard([param, grad]):
                        new_grad = self._create_regularization_of_grad(
                            param, grad, regularization)
                        params_and_grads.append((param, new_grad))
        return params_and_grads

1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053
    def flatten_param_grads(self, params_grads):
        need_flatten_params = []
        need_flatten_grads = []
        for p, g in params_grads:
            if g is None:
                continue
            g.persistable = True
            if getattr(p, 'need_clip', True) is False or getattr(
                    p, 'regularizer', None) is not None:
                warnings.warn(
                    "flatten_param_grads=True will be discarded since paramter '{}''s need_clip is False or "
                    "the regularizer is set".format(p.name))
                self._flatten_param_grads = False
                return params_grads

            need_flatten_params.append(p)
            need_flatten_grads.append(g)

        shape = [np.prod(p.shape) for p in need_flatten_params]
        block = need_flatten_params[0].block

        flatten_param = self.helper.create_global_variable(
            name='flatten_param',
            persistable=True,
            dtype=need_flatten_params[0].dtype,
            shape=[np.sum(shape)],
            belong_to_optimizer=True)

        flatten_param.trainable = True
        flatten_param.optimize_attr = need_flatten_params[0].optimize_attr
        flatten_param.regularizer = need_flatten_params[0].regularizer

        flatten_grad = self.helper.create_global_variable(
            name='flatten_grad',
            persistable=True,
            dtype=need_flatten_grads[0].dtype,
            shape=[np.sum(shape)],
            belong_to_optimizer=True)

        with program_guard(default_main_program()):
1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065
            block.append_op(type="coalesce_tensor",
                            inputs={"Input": need_flatten_params},
                            outputs={
                                "Output": need_flatten_params,
                                "FusedOutput": flatten_param
                            },
                            attrs={
                                "copy_data": True,
                                "use_align": True,
                                "align_size": self._align_size,
                                "dtype": need_flatten_params[0].dtype
                            })
1066

1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078
            block.append_op(type="coalesce_tensor",
                            inputs={"Input": need_flatten_grads},
                            outputs={
                                "Output": need_flatten_grads,
                                "FusedOutput": flatten_grad
                            },
                            attrs={
                                "copy_data": True,
                                "use_align": True,
                                "align_size": self._align_size,
                                "dtype": need_flatten_grads[0].dtype
                            })
1079 1080 1081

        #NOTE(zhiqiu): the initializer should be set after coalesce_tensor op,
        # so the shape of flatten_param and flatten_grad will be inferred.
1082 1083 1084 1085
        self.helper.set_variable_initializer(flatten_param,
                                             initializer=Constant(0.0))
        self.helper.set_variable_initializer(flatten_grad,
                                             initializer=Constant(0.0))
1086 1087 1088

        return [(flatten_param, flatten_grad)]

1089 1090 1091 1092 1093 1094 1095
    def apply_gradients(self, params_grads):
        """
        Second part of `minimize`, appending optimization operators for
        given `params_grads` pairs.

        Args:
            params_grads (list): list of (param, grad) pair to do optimization.
M
minqiyang 已提交
1096

1097 1098
        Returns:
            list: A list of operators appended to the current program.
M
minqiyang 已提交
1099

1100 1101 1102
        Examples:
            .. code-block:: python

1103
                import paddle.fluid as fluid
1104 1105 1106 1107 1108 1109 1110 1111 1112
                loss = network()
                optimizer = fluid.optimizer.SGD(learning_rate=0.1)
                params_grads = optimizer.backward(loss)
                # you may append operations for params_grads here
                # ...
                optimizer.apply_gradients(params_grads)
        """
        params_grads = sorted(params_grads, key=lambda x: x[0].name)

1113 1114 1115 1116 1117 1118
        # NOTE(zhiqiu): currently, only support ClipGradByGlobalNorm and without regularization.
        if self._flatten_param_grads and self.regularization is None:
            if self._grad_clip == None or isinstance(self._grad_clip,
                                                     ClipGradByGlobalNorm):
                params_grads = self.flatten_param_grads(params_grads)

1119
        # 'optimizer(grad_clip)' or 'set_gradient_clip'
1120 1121 1122 1123
        if self._grad_clip is not None:
            params_grads = self._grad_clip(params_grads)
        else:
            params_grads = append_gradient_clip_ops(params_grads)
1124 1125

        # Add regularization if any
1126 1127
        params_grads = self.append_regularization_ops(params_grads,
                                                      self.regularization)
1128 1129 1130 1131

        optimize_ops = self._create_optimization_pass(params_grads)
        return optimize_ops

C
chengduo 已提交
1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143
    def apply_optimize(self, loss, startup_program, params_grads):
        """
        Second part of `minimize`, appending optimization operators for
        given `params_grads` pairs.
        Args:
            loss (Variable): loss variable to run optimizations.
            startup_program (Program): startup_program for initializing parameters
                in `parameter_list`.
            params_grads (list): list of (param, grad) pair to do optimization.
        Returns:
            list: A list of operators appended to the current program.
        """
J
Jiabin Yang 已提交
1144
        if framework._non_static_mode():
C
chengduo 已提交
1145 1146
            with program_guard(framework.default_main_program(),
                               framework.default_startup_program()):
1147 1148
                if self._grad_clip is not None:
                    params_grads = self._grad_clip(params_grads)
1149 1150
                params_grads = self.append_regularization_ops(
                    params_grads, self.regularization)
C
chengduo 已提交
1151 1152 1153 1154 1155 1156 1157
                optimize_ops = self._create_optimization_pass(params_grads)
        else:
            program = loss.block.program
            with program_guard(program, startup_program):
                optimize_ops = self.apply_gradients(params_grads)
        return optimize_ops

G
gongweibao 已提交
1158
    def _get_no_grad_set(self, loss, no_grad_set=None):
1159
        no_grad_set = _get_no_grad_set_name(no_grad_set)
G
gongweibao 已提交
1160 1161 1162 1163 1164 1165 1166 1167
        parameters = loss.block.program.global_block().all_parameters()
        param_no_trainable = set(
            [param.name for param in parameters if param.trainable is False])
        # If the parameter is no trainable, it should not have a gradient.
        no_grad_set.update(param_no_trainable)

        return no_grad_set

1168 1169 1170 1171
    @framework.dygraph_only
    def clear_gradients(self):
        """
        Clear the gradients of all optimized parameters for model.
1172 1173

        If not, new gradient will accumulat on previous gradient.
1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200
        
        Returns:
            None
        
        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
                import numpy as np

                with fluid.dygraph.guard():
                    value = np.arange(26).reshape(2, 13).astype("float32")
                    a = fluid.dygraph.to_variable(value)
                    linear = fluid.Linear(13, 5, dtype="float32")
                    # This can be any optimizer supported by dygraph.
                    adam = fluid.optimizer.Adam(learning_rate = 0.01, 
                                                parameter_list = linear.parameters())
                    out = linear(a)
                    out.backward()
                    adam.minimize(out)
                    adam.clear_gradients()

        """
        for p in self._parameter_list:
            if p.trainable:
                p.clear_gradient()

1201
    @imperative_base.no_grad
Q
Qiao Longfei 已提交
1202 1203
    def minimize(self,
                 loss,
1204
                 startup_program=None,
Q
Qiao Longfei 已提交
1205
                 parameter_list=None,
1206
                 no_grad_set=None):
1207
        """
1208
        Add operations to minimize ``loss`` by updating ``parameter_list``.
M
minqiyang 已提交
1209

1210
        Args:
1211 1212 1213 1214
            loss (Variable): A ``Variable`` containing the value to minimize.
            startup_program (Program, optional): :ref:`api_fluid_Program` for
                initializing parameters in ``parameter_list``. The default value
                is None, at this time :ref:`api_fluid_default_startup_program` will be used.
H
hong 已提交
1215
            parameter_list (Iterable, optional): Iterable of ``Variable`` or ``Variable.name`` to update
1216 1217
                to minimize ``loss``. The default value is None, at this time all parameters
                will be updated.
1218
            no_grad_set (set, optional): Set of ``Variable``  or ``Variable.name`` that don't need
1219
                to be updated. The default value is None.
Q
Qiao Longfei 已提交
1220

1221
        Returns:
1222 1223 1224
            tuple: tuple (optimize_ops, params_grads), A list of operators appended
            by minimize and a list of (param, grad) variable pairs, param is
            ``Parameter``, grad is the gradient value corresponding to the parameter.
1225 1226 1227
            The returned tuple can be passed to ``fetch_list`` in ``Executor.run()`` to 
            indicate program pruning. If so, the program will be pruned by ``feed`` and 
            ``fetch_list`` before run, see details in ``Executor``.
1228 1229 1230

        Examples:
            Please refer to the example of current Optimizer.
Q
Qiao Longfei 已提交
1231
        """
C
chengduo 已提交
1232
        assert isinstance(loss, Variable), "The loss should be an Variable."
1233

1234 1235
        parameter_list = parameter_list if parameter_list \
            else self._parameter_list
1236

1237 1238 1239 1240
        params_grads = self.backward(loss,
                                     startup_program=startup_program,
                                     parameter_list=parameter_list,
                                     no_grad_set=no_grad_set)
1241

1242 1243 1244
        optimize_ops = self.apply_optimize(loss,
                                           startup_program=startup_program,
                                           params_grads=params_grads)
M
minqiyang 已提交
1245

Q
Qiao Longfei 已提交
1246
        return optimize_ops, params_grads
Q
Qiao Longfei 已提交
1247 1248 1249


class SGDOptimizer(Optimizer):
1250
    r"""
Q
qiaolongfei 已提交
1251 1252 1253 1254 1255 1256
    Optimizer of the stochastic gradient descent algorithm.

    .. math::

        param\_out = param - learning\_rate * grad

1257 1258 1259
    Parameters:
        learning_rate (float|Variable): The learning rate used to update parameters. \
            Can be a float value or a Variable with one float value as data element.
H
hong 已提交
1260
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1261 1262
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
1263 1264 1265 1266 1267
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
1268 1269 1270 1271
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
1272 1273
        name (str, optional): This parameter is used by developers to print debugging information. \
            For details, please refer to :ref:`api_guide_Name`. Default is None.
Q
qiaolongfei 已提交
1274 1275 1276 1277

    Examples:
        .. code-block:: python

1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302
            import paddle
            import paddle.fluid as fluid
            import numpy as np

            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
                x = fluid.layers.data(name='x', shape=[13], dtype='float32')
                y = fluid.layers.data(name='y', shape=[1], dtype='float32')
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
                cost = fluid.layers.square_error_cost(input=y_predict, label=y)
                avg_cost = fluid.layers.mean(cost)

                sgd_optimizer = fluid.optimizer.SGD(learning_rate=0.001)
                sgd_optimizer.minimize(avg_cost)

                fetch_list = [avg_cost]
                train_reader = paddle.batch(
                    paddle.dataset.uci_housing.train(), batch_size=1)
                feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
                exe = fluid.Executor(place)
                exe.run(fluid.default_startup_program())
                for data in train_reader():
                    exe.run(main, feed=feeder.feed(data), fetch_list=fetch_list)

Q
Qiao Longfei 已提交
1303 1304
    """

1305 1306 1307 1308
    def __init__(self,
                 learning_rate,
                 parameter_list=None,
                 regularization=None,
1309
                 grad_clip=None,
1310
                 multi_precision=False,
1311
                 name=None):
Q
Qiao Longfei 已提交
1312
        assert learning_rate is not None
1313 1314 1315 1316 1317
        super(SGDOptimizer, self).__init__(learning_rate=learning_rate,
                                           parameter_list=parameter_list,
                                           regularization=regularization,
                                           grad_clip=grad_clip,
                                           name=name)
Q
Qiao Longfei 已提交
1318
        self.type = "sgd"
1319
        self._use_mkldnn = False
1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330
        self._multi_precision = multi_precision
        self._master_weights = {}

    def _create_master_weight(self, param):
        if param.name in self._master_weights:
            var = self._master_weights[param.name]
        else:
            assert isinstance(self.helper, LayerHelper)

            var_name = param.name + "_fp32_master"
            var_name = unique_name.generate(var_name)
1331 1332 1333 1334 1335
            var = layers.create_global_var(name=var_name,
                                           shape=param.shape,
                                           value=0,
                                           dtype='float32',
                                           persistable=True)
1336
            block = self.helper.startup_program.global_block()
1337 1338 1339 1340 1341 1342 1343
            block.append_op(type="cast",
                            inputs={"X": [param]},
                            outputs={"Out": [var]},
                            attrs={
                                "in_dtype": param.dtype,
                                "out_dtype": core.VarDesc.VarType.FP32
                            })
1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359 1360 1361
            self._master_weights[param.name] = var
        return var

    def _create_accumulators(self, block, parameters):
        assert isinstance(block, framework.Block)
        if isinstance(parameters, dict):
            parameters = self._update_param_group(parameters)

        # Create accumulator tensors for first and second moments
        for p in parameters:
            if self._multi_precision and p.dtype == core.VarDesc.VarType.FP16:
                master_p = self._create_master_weight(p)
                continue
            if p.dtype == core.VarDesc.VarType.FP16 and not self._multi_precision:
                warnings.warn(
                    "Accumulating with FP16 in optimizer can lead to poor accuracy or slow convergence."
                    "Consider using multi_precision=True option of the Adam optimizer."
                )
Q
Qiao Longfei 已提交
1362

1363
    @no_grad
1364
    def _append_optimize_op(self, block, param_and_grad):
1365 1366 1367 1368 1369 1370

        find_master = self._multi_precision and param_and_grad[
            0].dtype == core.VarDesc.VarType.FP16
        master_weight = (self._master_weights[param_and_grad[0].name]
                         if find_master else None)

1371
        lr = self._create_param_lr(param_and_grad)
Z
zyfncg 已提交
1372
        if in_dygraph_mode():
1373 1374
            _C_ops.final_state_sgd_(param_and_grad[0], lr, param_and_grad[1],
                                    master_weight, find_master)
Z
zyfncg 已提交
1375 1376
            return None
        if _in_legacy_dygraph():
1377 1378
            _C_ops.sgd(param_and_grad[0], lr, param_and_grad[1], master_weight,
                       param_and_grad[0], master_weight)
1379
            return None
1380

1381
        assert isinstance(block, framework.Block)
Q
Qiao Longfei 已提交
1382
        # create the optimize op
1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396
        inputs = {
            "Param": param_and_grad[0],
            "Grad": param_and_grad[1],
            "LearningRate": lr
        }

        outputs = {"ParamOut": param_and_grad[0]}

        attrs = {"multi_precision": find_master}

        if find_master:
            inputs["MasterParam"] = master_weight
            outputs["MasterParamOut"] = master_weight

1397 1398 1399 1400 1401
        sgd_op = block.append_op(type=self.type,
                                 inputs=inputs,
                                 outputs=outputs,
                                 attrs=attrs,
                                 stop_gradient=True)
Q
Qiao Longfei 已提交
1402 1403

        return sgd_op
1404 1405 1406


class MomentumOptimizer(Optimizer):
1407
    r"""
Q
qiaolongfei 已提交
1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420

    Simple Momentum optimizer with velocity state

    This optimizer has a flag for Nestrov Momentum.

    The update equations are as follows:

    .. math::

        & velocity = mu * velocity + gradient

        & if (use\_nesterov):

1421
        &\quad   param = param - (gradient + mu * velocity) * learning\_rate
Q
qiaolongfei 已提交
1422 1423 1424

        & else:

Q
qiaolongfei 已提交
1425
        &\quad   param = param - learning\_rate * velocity
Q
qiaolongfei 已提交
1426

1427 1428 1429 1430
    Parameters:
        learning_rate (float|Variable): The learning rate used to update parameters. \
            Can be a float value or a Variable with one float value as data element.
        momentum (float): Momentum factor
H
hong 已提交
1431
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1432 1433
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
1434
        use_nesterov (bool, optional): Enables Nesterov momentum, default is false.
1435 1436 1437 1438 1439
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
1440 1441 1442 1443
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
1444 1445
        name (str, optional): This parameter is used by developers to print debugging information. \
            For details, please refer to :ref:`api_guide_Name`. Default is None.
Q
qiaolongfei 已提交
1446 1447 1448 1449

    Examples:
        .. code-block:: python

1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474
            import paddle
            import paddle.fluid as fluid
            import numpy as np

            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
                x = fluid.layers.data(name='x', shape=[13], dtype='float32')
                y = fluid.layers.data(name='y', shape=[1], dtype='float32')
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
                cost = fluid.layers.square_error_cost(input=y_predict, label=y)
                avg_cost = fluid.layers.mean(cost)

                moment_optimizer = fluid.optimizer.MomentumOptimizer(learning_rate=0.001, momentum=0.9)
                moment_optimizer.minimize(avg_cost)

                fetch_list = [avg_cost]
                train_reader = paddle.batch(
                    paddle.dataset.uci_housing.train(), batch_size=1)
                feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
                exe = fluid.Executor(place)
                exe.run(fluid.default_startup_program())
                for data in train_reader():
                    exe.run(main, feed=feeder.feed(data), fetch_list=fetch_list)

1475 1476 1477
    """
    _velocity_acc_str = "velocity"

X
Xin Pan 已提交
1478 1479 1480
    def __init__(self,
                 learning_rate,
                 momentum,
1481
                 parameter_list=None,
X
Xin Pan 已提交
1482 1483
                 use_nesterov=False,
                 regularization=None,
1484
                 grad_clip=None,
X
Xin Pan 已提交
1485
                 name=None):
1486 1487
        assert learning_rate is not None
        assert momentum is not None
1488 1489 1490 1491 1492
        super(MomentumOptimizer, self).__init__(learning_rate=learning_rate,
                                                parameter_list=parameter_list,
                                                regularization=regularization,
                                                grad_clip=grad_clip,
                                                name=name)
1493 1494
        self.type = "momentum"
        self._momentum = momentum
1495
        self._use_nesterov = bool(use_nesterov)
1496 1497 1498 1499 1500

    def _create_accumulators(self, block, parameters):
        assert isinstance(block, framework.Block)

        for p in parameters:
Q
Qiao Longfei 已提交
1501
            self._add_accumulator(self._velocity_acc_str, p)
1502 1503 1504 1505 1506 1507

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)

        velocity_acc = self._get_accumulator(self._velocity_acc_str,
                                             param_and_grad[0])
1508
        lr = self._create_param_lr(param_and_grad)
1509
        master_weight = None
J
Jiabin Yang 已提交
1510
        if framework._non_static_mode():
1511 1512 1513 1514 1515
            _, _, _ = _C_ops.momentum(param_and_grad[0], param_and_grad[1],
                                      velocity_acc, lr, master_weight,
                                      param_and_grad[0], velocity_acc,
                                      master_weight, 'mu', self._momentum,
                                      'use_nesterov', self._use_nesterov)
1516
            return None
1517

1518
        attrs = {"mu": self._momentum, "use_nesterov": self._use_nesterov}
1519 1520 1521 1522
        inputs = {
            "Param": [param_and_grad[0]],
            "Grad": [param_and_grad[1]],
            "Velocity": [velocity_acc],
1523
            "LearningRate": [lr]
1524 1525 1526 1527 1528 1529
        }

        outputs = {
            "ParamOut": [param_and_grad[0]],
            "VelocityOut": [velocity_acc]
        }
1530
        # create the momentum optimize op
1531 1532 1533 1534 1535
        momentum_op = block.append_op(type=self.type,
                                      inputs=inputs,
                                      outputs=outputs,
                                      attrs=attrs,
                                      stop_gradient=True)
1536 1537

        return momentum_op
1538 1539


1540
class DGCMomentumOptimizer(Optimizer):
1541
    r"""
1542
	:api_attr: Static Graph
S
swtkiwi 已提交
1543

1544
    DGC (Deep Gradient Compression) Momentum Optimizer. Original paper is https://arxiv.org/abs/1712.01887
1545

G
gongweibao 已提交
1546
    DGC reduces the communication bandwidth by sending only the important gradients (sparse update):\
1547 1548
        only gradients larger than a threshold are transmitted.

G
gongweibao 已提交
1549
    To avoid losing information, DGC accumulates the rest of the gradients locally.
1550 1551 1552

    Eventually, these gradients become large enough to be transmitted.

1553
    Thus, DGC sends the large gradients immediately but eventually sends all of the gradients over time.
1554

G
gongweibao 已提交
1555
    To ensure no loss of accuracy, DGC employs momentum correction and local gradient clipping on top of the gradient sparsification to maintain model performance.
1556 1557 1558 1559

    DGC also uses momentum factor masking and warmup training to overcome the staleness problem caused by reduced communication.

    This optimizer will do two things:
1560

1561 1562
        1. Compress the gradient by get TopK import value from tensor \
            and use it for allreduce to reduce network bandwidth.
1563

1564
        2. Call momentum to optimize the cost.
1565 1566

    Args:
1567 1568
        learning_rate (float|Variable): The learning rate used to update parameters. \
            It can be a float value or a Variable with one float value as a data element.
1569
        momentum (float): Momentum factor.
G
gongweibao 已提交
1570
        rampup_begin_step (int): The beginning step from which gradient compression is implemented.
1571 1572 1573 1574 1575 1576 1577
        rampup_step (int): Time steps used in sparsity warm-up periods. Default is 1.
            For example, if the sparsity is [0.75, 0.9375, 0.984375, 0.996, 0.999], and the rampup_step is 100, \
                it will use 0.75 at 0~19 steps, and 0.9375 at 20~39 steps, and so on. \
                And when reach sparsity array ends, it will use 0.999 then and after.
        sparsity (list[float]): Get top important element from gradient tensor, the ratio is (1 - current sparsity). \
            Default is [0.999]. For example, if the sparsity is [0.99, 0.999], \
                the top [1%, 0.1%] important element will be transmitted.
H
hong 已提交
1578
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1579 1580
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
1581
        use_nesterov (bool): Enables Nesterov momentum. True means use Nesterov. Default is False.
1582 1583 1584 1585 1586
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
1587 1588 1589
        grad_clip (GradientClipByNorm, optional): Gradient cliping strategy. ``DGCMomentumOptimizer`` only support 
            :ref:`api_fluid_clip_GradientClipByNorm` , and if not, it will raise TypeError. Default None, 
            meaning there is no gradient clipping.
1590 1591
        name (str, optional): This parameter is used by developers to print debugging information. \
            For details, please refer to :ref:`api_guide_Name`. Default is None.
1592 1593 1594 1595

    Examples:
        .. code-block:: python

1596
            import paddle.fluid as fluid
1597
            optimizer = fluid.optimizer.DGCMomentumOptimizer(
G
gongweibao 已提交
1598 1599 1600 1601 1602
                        learning_rate=0.0001,
                        momentum=0.9,
                        rampup_step=1000,
                        rampup_begin_step=1252,
                        sparsity=[0.999, 0.999])
1603 1604

    """
1605 1606
    _u_velocity_acc_str = "_dgc_u_"
    _v_velocity_acc_str = "_dgc_v_"
1607 1608 1609 1610 1611 1612 1613

    def __init__(self,
                 learning_rate,
                 momentum,
                 rampup_begin_step,
                 rampup_step=1,
                 sparsity=[0.999],
1614
                 parameter_list=None,
1615 1616 1617
                 use_nesterov=False,
                 num_trainers=None,
                 regularization=None,
1618
                 grad_clip=None,
1619
                 name=None):
J
Jiabin Yang 已提交
1620
        if framework._non_static_mode():
Z
zhongpu 已提交
1621
            raise Exception("In dygraph, don't support DGCMomentumOptimizer.")
1622 1623 1624 1625

        assert core.is_compiled_with_cuda(), \
            "Paddle is not compiled with CUDA. DGC is only support GPU for now."

1626 1627
        assert learning_rate is not None
        assert momentum is not None
1628 1629 1630 1631 1632 1633
        super(DGCMomentumOptimizer,
              self).__init__(learning_rate=learning_rate,
                             parameter_list=parameter_list,
                             regularization=regularization,
                             grad_clip=grad_clip,
                             name=name)
1634 1635 1636
        self.type = "dgc_momentum"
        self._momentum = momentum
        self._use_nesterov = bool(use_nesterov)
1637

1638
        assert rampup_begin_step >= 0, "rampup_begin_step must >= 0"
1639
        self._rampup_begin_step = rampup_begin_step
1640 1641
        self._rampup_step = rampup_step
        self._sparsity = sparsity
1642

1643
        self._rampup_begin_step_var = None
1644
        self._global_step_var = None
1645

1646 1647 1648 1649 1650 1651 1652 1653 1654
        self._dgc_clip_norm = None
        if grad_clip is not None:
            if not isinstance(grad_clip, GradientClipByNorm):
                raise TypeError(
                    "The type of grad_clip should be 'GradientClipByNorm', because DGCMomentumOptimizer only support GradientClipByNorm"
                )
            assert isinstance(
                num_trainers, int
            ), "The type of num_trainers should be 'int', but received %s" % type(
J
Jiangxinz 已提交
1655
                num_trainers)
1656
            assert num_trainers > 0, "The value of num_trainers should be greater than 0!"
1657 1658

            self._num_trainers = num_trainers
1659
            self._dgc_clip_norm = grad_clip.clip_norm * (num_trainers**-0.5)
1660

1661 1662
        self.regular_type, self.regular_coeff = self._get_regularization_param(
            self.regularization)
1663

1664 1665 1666
    def _get_regularization_param(self, regularization):
        regular_type = 0
        regular_coeff = 0.0
1667

1668 1669
        if regularization is not None:
            regular_coeff = regularization._regularization_coeff
1670
            from .regularizer import L1Decay, L2Decay
1671 1672 1673 1674
            if isinstance(regularization, L1Decay):
                regular_type = 1
            elif isinstance(regularization, L2Decay):
                regular_type = 2
1675 1676
            else:
                assert False, 'regularization must be None|L1Decay|L2Deacy'
1677
        return regular_type, regular_coeff
1678

1679 1680 1681 1682 1683 1684 1685 1686 1687 1688 1689
    def _is_use_dgc(self, param_var, grad_var):
        var_numel = abs(reduce(lambda x, y: x * y, param_var.shape))
        if var_numel < 16384 or \
           param_var.type == core.VarDesc.VarType.SELECTED_ROWS  or \
           grad_var.type == core.VarDesc.VarType.SELECTED_ROWS  or  \
               param_var.dtype != core.VarDesc.VarType.FP32 :
            return False
        return True

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)
1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700 1701 1702 1703 1704
        velocity_acc = self._get_accumulator(self._u_velocity_acc_str,
                                             param_and_grad[0])
        assert velocity_acc is not None

        inputs = {
            "Param": param_and_grad[0],
            "Grad": param_and_grad[1],
            "Velocity": velocity_acc,
            "LearningRate": self._create_param_lr(param_and_grad),
        }
        outputs = {
            "ParamOut": param_and_grad[0],
            "VelocityOut": velocity_acc,
        }
        attrs = {"mu": self._momentum, "use_nesterov": self._use_nesterov}
1705 1706

        if not self._is_use_dgc(param_and_grad[0], param_and_grad[1]):
1707 1708 1709
            type = "momentum"
        else:
            type = "dgc_momentum"
1710 1711 1712 1713 1714
            inputs.update({
                "current_step": self._global_step_var,
                "nranks": self._nranks_var
            })
            outputs.update({'Grad_out': param_and_grad[1]})
1715
            attrs.update({"rampup_begin_step": float(self._rampup_begin_step)})
1716 1717

        # create the dgc momentum optimize op
1718 1719 1720 1721 1722
        dgc_momentum_op = block.append_op(type=type,
                                          inputs=inputs,
                                          outputs=outputs,
                                          attrs=attrs,
                                          stop_gradient=True)
1723 1724
        return dgc_momentum_op

1725 1726 1727 1728 1729
    def _add_auto_increment_var(self, counter_name, begin, step=1):
        helper = LayerHelper('global_step_counter')
        counter, is_new_var = helper.create_or_get_global_variable(
            name=counter_name, dtype='float32', shape=[1], persistable=True)
        if is_new_var:
1730 1731 1732 1733
            helper.set_variable_initializer(counter,
                                            initializer=Constant(
                                                value=float(begin - 1),
                                                force_cpu=True))
1734 1735 1736 1737 1738 1739 1740 1741 1742 1743
            helper.main_program.global_block()._prepend_op(
                type='increment',
                inputs={'X': [counter]},
                outputs={'Out': [counter]},
                attrs={'step': float(step)},
                stop_gradient=True)
            counter.stop_gradient = True

        return counter

1744 1745 1746 1747 1748
    def _add_nranks_var(self, name, value=-1):
        helper = LayerHelper('global_step_counter')
        counter, is_new_var = helper.create_or_get_global_variable(
            name=name, dtype='float32', shape=[1], persistable=True)
        if is_new_var:
1749 1750 1751 1752
            helper.set_variable_initializer(counter,
                                            initializer=Constant(
                                                value=float(value),
                                                force_cpu=True))
1753 1754 1755 1756
            counter.stop_gradient = True

        return counter

1757 1758 1759 1760 1761 1762
    def _append_dgc_ops(self, param_and_grads):
        main_program = default_main_program()
        main_program._enable_dgc = True

        # step counter
        self._global_step_var = self._add_auto_increment_var(
G
gongweibao 已提交
1763
            counter_name=core.dgc.kDGCCounterName(), begin=0)
1764

1765 1766
        self._nranks_var = self._add_nranks_var(name=core.dgc.kDGCNRanksName(),
                                                value=-1)
1767

1768 1769 1770 1771 1772
        # rampup begin step var for all_reduce_op_handle
        self._rampup_begin_step_var = tensor.create_global_var(
            shape=[1],
            dtype=core.VarDesc.VarType.FP32,
            persistable=True,
G
gongweibao 已提交
1773
            name=core.dgc.kDGCRampUpBeginStepName(),
1774 1775 1776
            value=self._rampup_begin_step * 1.0,
            force_cpu=True)

1777 1778
        self.helper = LayerHelper(self.__class__.__name__)

1779
        for param_var, grad_var in param_and_grads:
1780 1781 1782
            # reuse velocity in dgc_op and dgc_momentum_op
            u_var = self._add_accumulator(self._u_velocity_acc_str, param_var)

1783
            if not self._is_use_dgc(param_var, grad_var):
1784 1785
                continue

1786
            v_var = self._add_accumulator(self._v_velocity_acc_str, param_var)
1787

1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810
            k_var = tensor.create_global_var(shape=[1],
                                             dtype=param_var.dtype,
                                             persistable=True,
                                             name=param_var.name +
                                             core.dgc.kDGCKName(),
                                             value=0.0,
                                             force_cpu=True)

            encoded_var = tensor.create_global_var(shape=[1],
                                                   dtype=param_var.dtype,
                                                   persistable=True,
                                                   name=param_var.name +
                                                   core.dgc.kDGCEncodedName(),
                                                   value=0.0,
                                                   force_cpu=False)

            gather_var = tensor.create_global_var(shape=[1],
                                                  dtype=param_var.dtype,
                                                  persistable=True,
                                                  name=param_var.name +
                                                  core.dgc.kDGCGatherName(),
                                                  value=0.0,
                                                  force_cpu=False)
1811

1812 1813 1814 1815 1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830
            # del back oprolevarname
            op_maker = core.op_proto_and_checker_maker
            backward = core.op_proto_and_checker_maker.OpRole.Backward
            for op in main_program.global_block().ops:
                if not self._is_the_backward_op(op):
                    continue

                var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
                if param_var.name not in var_attr:
                    continue

                var_attr.remove(param_var.name)
                var_attr.remove(grad_var.name)
                if len(var_attr) > 1:
                    op._set_attr(op_maker.kOpRoleVarAttrName(), var_attr)
                else:
                    op._remove_attr(op_maker.kOpRoleVarAttrName())

            clip_var = grad_var
1831 1832
            if self._dgc_clip_norm is not None:
                clip_var = self._append_clip_norm(grad_var, self._dgc_clip_norm)
1833
            self._dgc_op(param_var, clip_var, grad_var, u_var, v_var, k_var,
1834
                         encoded_var, gather_var)
1835 1836 1837 1838 1839 1840 1841 1842 1843 1844 1845 1846 1847 1848 1849

    def _is_the_backward_op(self, op):
        op_maker = core.op_proto_and_checker_maker
        backward = core.op_proto_and_checker_maker.OpRole.Backward
        if op_maker.kOpRoleVarAttrName() in op.attr_names and \
                int(op.all_attrs()[op_maker.kOpRoleAttrName()]) == int(backward):
            return True
        return False

    def _clip_by_norm(self, x, max_norm, name=None):
        args = {'x': x, 'max_norm': max_norm, 'name': name}

        helper = LayerHelper("dgc_clip_by_norm_op", **args)

        if name is None:
1850 1851
            name = unique_name.generate_with_ignorable_key(".".join(
                [helper.name, 'tmp']))
1852

1853 1854 1855 1856 1857 1858 1859 1860 1861 1862 1863 1864 1865 1866 1867
        out = helper.create_variable(type=x.type,
                                     name=name,
                                     dtype=x.dtype,
                                     persistable=False)

        helper.append_op(type="dgc_clip_by_norm",
                         inputs={
                             "X": x,
                             "current_step": self._global_step_var
                         },
                         attrs={
                             "max_norm": max_norm,
                             "rampup_begin_step": float(self._rampup_begin_step)
                         },
                         outputs={"Out": out})
1868 1869 1870 1871
        return out

    def _append_clip_norm(self, grad_var, clip_norm):
        with grad_var.block.program._backward_role_guard():
1872 1873 1874
            return self._clip_by_norm(x=grad_var,
                                      max_norm=clip_norm,
                                      name=grad_var.name)
1875 1876

    def _dgc_op(self, param_var, clip_var, grad_var, u_var, v_var, k_var,
1877
                encoded_var, gather_var):
1878 1879
        block = framework.default_main_program().global_block()
        op_maker = core.op_proto_and_checker_maker
1880

1881 1882 1883 1884 1885 1886 1887
        regular_type = self.regular_type
        regular_coeff = self.regular_coeff
        # The regularizer of the Parameters have higher priority
        if param_var.regularizer is not None:
            regular_type, regular_coeff = self._get_regularization_param(
                param_var.regularizer)

1888 1889 1890 1891 1892 1893 1894 1895 1896 1897 1898 1899 1900 1901 1902 1903 1904 1905 1906 1907 1908 1909 1910 1911 1912 1913 1914 1915 1916 1917 1918 1919 1920 1921
        dgc_op = block.append_op(type="dgc",
                                 inputs={
                                     "U": u_var,
                                     "V": v_var,
                                     "Grad": clip_var,
                                     "Param": param_var,
                                     "current_step": self._global_step_var,
                                     "nranks": self._nranks_var,
                                 },
                                 outputs={
                                     "U_out": u_var,
                                     "V_out": v_var,
                                     "EncodeGrad": encoded_var,
                                     "k": k_var,
                                     "Grad_out": grad_var,
                                     "GatherBuff": gather_var,
                                 },
                                 attrs={
                                     "m":
                                     self._momentum,
                                     "sparsity":
                                     self._sparsity,
                                     "use_nesterov":
                                     self._use_nesterov,
                                     "rampup_begin_step":
                                     float(self._rampup_begin_step),
                                     "rampup_step":
                                     float(self._rampup_step),
                                     "regular_coeff":
                                     float(regular_coeff),
                                     "regular_type":
                                     int(regular_type),
                                 },
                                 stop_gradient=True)
1922 1923 1924 1925 1926 1927

        backward = op_maker.OpRole.Backward
        dgc_op._set_attr(op_maker.kOpRoleAttrName(), backward)
        dgc_op._set_attr(op_maker.kOpRoleVarAttrName(),
                         [param_var.name, grad_var.name])

1928
    @imperative_base.no_grad
1929
    def apply_gradients(self, params_grads):
1930 1931 1932 1933 1934
        # Note: since we can't use all_reduce_op now,
        # dgc_op should be the last op of one grad.
        # Maybe need a grad allreduce pass.
        self._append_dgc_ops(params_grads)

1935 1936 1937 1938 1939 1940
        params_grads = sorted(params_grads, key=lambda x: x[0].name)
        params_grads, table_param_and_grad, table_optimize_op = \
            self._process_distribute_lookuptable(params_grads)

        not_dgc_params_grads = []
        dgc_params_grads = []
1941
        # DGC clip and regularization in optimizer.backward
1942 1943 1944 1945 1946 1947
        for param, grad in params_grads:
            if not self._is_use_dgc(param, grad):
                not_dgc_params_grads.append((param, grad))
            else:
                dgc_params_grads.append((param, grad))

1948
        # 'optimizer(grad_clip)' or 'set_gradient_clip'
1949 1950 1951 1952 1953
        if self._grad_clip is not None:
            not_dgc_params_grads = self._grad_clip(not_dgc_params_grads)
        else:
            not_dgc_params_grads = append_gradient_clip_ops(
                not_dgc_params_grads)
1954

1955 1956
        not_dgc_params_grads = self.append_regularization_ops(
            not_dgc_params_grads, self.regularization)
1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967

        params_grads = not_dgc_params_grads + dgc_params_grads
        params_grads = sorted(params_grads, key=lambda x: x[0].name)

        optimize_ops = self._create_optimization_pass(params_grads)
        if table_optimize_op is not None:
            optimize_ops.append(table_optimize_op)
            params_grads.append(table_param_and_grad)

        return optimize_ops

1968

1969
class LarsMomentumOptimizer(Optimizer):
1970
    r"""
1971 1972 1973 1974 1975 1976 1977 1978 1979
    Momentum optimizer with LARS support

    The update equations are as follows:

    .. math::

        & local\_learning\_rate = learning\_rate * lars\_coeff * \\
          \\frac{||param||}{||gradient|| + lars\_weight\_decay * ||param||}

1980
        & velocity = mu * velocity + local\_learning\_rate * (gradient + lars\_weight\_decay * param + epsilon)
1981 1982 1983

        & param = param - velocity

1984 1985 1986 1987 1988 1989
    Parameters:
        learning_rate (float|Variable): The learning rate used to update parameters. \
            Can be a float value or a Variable with one float value as data element. \
            momentum (float): momentum factor
        lars_coeff (float): Defines how much we trust the layer to change its weights.
        lars_weight_decay (float): Weight decay coefficient for decaying using LARS.
H
hong 已提交
1990
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1991 1992
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
1993 1994 1995 1996 1997
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
1998 1999 2000 2001
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2002 2003
        name (str, optional): This parameter is used by developers to print debugging information. \
            For details, please refer to :ref:`api_guide_Name`. Default is None.
2004 2005
        exclude_from_weight_decay (list[str], optional): Name string of layers which will be exclude from lars weight decay. Default is None.
        epsilon (float, optional): Epsilon to avoid Division by Zero when calculate local lr. Default is 0.
2006 2007 2008
        multi_precision (bool, optional): Whether to use multi-precision during weight updating.
        rescale_grad (float, optional): Multiply the gradient with `rescale_grad` \
            before updating. Often choose to be `1.0/batch_size`.
2009
        
2010 2011 2012
    Examples:
        .. code-block:: python

2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027 2028
            import paddle.fluid as fluid
            import numpy as np

            np_inp = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
            inp = fluid.layers.data(
                name="inp", shape=[2, 2], append_batch_size=False)
            out = fluid.layers.fc(inp, size=3)
            out = fluid.layers.reduce_sum(out)
            optimizer = fluid.optimizer.LarsMomentumOptimizer(learning_rate=0.001, momentum=0.9)
            optimizer.minimize(out)

            exe = fluid.Executor(fluid.CPUPlace())
            exe.run(fluid.default_startup_program())
            exe.run(
                feed={"inp": np_inp},
                fetch_list=[out.name])
2029 2030 2031 2032 2033 2034 2035 2036
    """
    _velocity_acc_str = "velocity"

    def __init__(self,
                 learning_rate,
                 momentum,
                 lars_coeff=0.001,
                 lars_weight_decay=0.0005,
2037
                 parameter_list=None,
2038
                 regularization=None,
2039
                 grad_clip=None,
2040 2041
                 name=None,
                 exclude_from_weight_decay=None,
2042 2043 2044
                 epsilon=0,
                 multi_precision=False,
                 rescale_grad=1.0):
2045 2046
        assert learning_rate is not None
        assert momentum is not None
2047 2048 2049 2050 2051 2052
        super(LarsMomentumOptimizer,
              self).__init__(learning_rate=learning_rate,
                             parameter_list=parameter_list,
                             regularization=regularization,
                             grad_clip=grad_clip,
                             name=name)
2053 2054 2055 2056
        self.type = "lars_momentum"
        self._momentum = momentum
        self._lars_coeff = float(lars_coeff)
        self._lars_weight_decay = float(lars_weight_decay)
2057 2058 2059 2060 2061
        self._epsilon = float(epsilon)
        if exclude_from_weight_decay is None:
            self._exclude_from_weight_decay = []
        else:
            self._exclude_from_weight_decay = exclude_from_weight_decay
2062 2063 2064 2065 2066
        self._multi_precision = multi_precision
        self._rescale_grad = float(rescale_grad)
        self._master_weights = {}

    def _create_master_weight(self, param):
2067 2068 2069 2070
        if param.name in self._master_weights:
            var = self._master_weights[param.name]
        else:
            assert isinstance(self.helper, LayerHelper)
2071

2072 2073
            var_name = param.name + '_fp32_master'
            var_name = unique_name.generate(var_name)
2074 2075 2076 2077 2078
            var = layers.create_global_var(name=var_name,
                                           shape=param.shape,
                                           value=0,
                                           dtype='float32',
                                           persistable=True)
2079
            block = self.helper.startup_program.global_block()
2080 2081 2082 2083 2084 2085 2086
            block.append_op(type="cast",
                            inputs={"X": [param]},
                            outputs={"Out": [var]},
                            attrs={
                                "in_dtype": param.dtype,
                                "out_dtype": core.VarDesc.VarType.FP32
                            })
2087
            self._master_weights[param.name] = var
2088 2089 2090 2091 2092 2093 2094 2095 2096 2097 2098 2099 2100 2101 2102 2103
        return var

    def _get_accumulator(self, name, param):
        """Utility function to fetch an accumulator for a parameter
        Args:
            name: name of the accumulator
            param: parameter variable for which accumulator is to be fetched
        Returns:
            accumulator variable for the parameter
        """
        if self._name is not None:
            name = self._name + "_" + name
        find_master = self._multi_precision and param.dtype == core.VarDesc.VarType.FP16
        target_param = self._master_weights[
            param.name] if find_master else param
        target_name = target_param.name
2104 2105 2106 2107 2108
        if (name not in self._accumulators
                or target_name not in self._accumulators[name]):
            raise Exception(
                "Accumulator {} does not exist for parameter {}".format(
                    name, target_name))
2109
        return self._accumulators[name][target_name]
2110 2111 2112 2113 2114

    def _create_accumulators(self, block, parameters):
        assert isinstance(block, framework.Block)

        for p in parameters:
2115 2116 2117 2118 2119 2120 2121 2122 2123
            if self._multi_precision and p.dtype == core.VarDesc.VarType.FP16:
                master_p = self._create_master_weight(p)
                self._add_accumulator(self._velocity_acc_str, master_p)
                continue
            if p.dtype == core.VarDesc.VarType.FP16 and not self._multi_precision:
                warnings.warn(
                    "Accumulating with FP16 in optimizer can lead to poor accuracy or slow convergence."
                    "Consider using multi_precision=True option of the Lars optimizer."
                )
2124 2125 2126 2127
            self._add_accumulator(self._velocity_acc_str, p)

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)
2128 2129 2130 2131 2132 2133 2134 2135
        _lars_weight_decay = self._lars_weight_decay
        param_name = param_and_grad[0].name
        if len(self._exclude_from_weight_decay) > 0:
            for name in self._exclude_from_weight_decay:
                if name in param_name:
                    _lars_weight_decay = 0.0
                    break

2136 2137
        velocity_acc = self._get_accumulator(self._velocity_acc_str,
                                             param_and_grad[0])
2138 2139 2140 2141 2142 2143 2144 2145 2146
        lr = self._create_param_lr(param_and_grad)

        find_master = self._multi_precision and param_and_grad[
            0].dtype == core.VarDesc.VarType.FP16
        master_weight = (self._master_weights[param_and_grad[0].name]
                         if find_master else None)

        attrs = {
            "mu": self._momentum,
2147
            "lars_coeff": self._lars_coeff,
L
limingshu 已提交
2148
            "lars_weight_decay": [_lars_weight_decay],
2149
            "multi_precision": find_master,
L
limingshu 已提交
2150
            "epsilon": self._epsilon,
2151 2152 2153 2154 2155 2156 2157 2158 2159 2160 2161 2162 2163 2164 2165 2166
            "rescale_grad": self._rescale_grad
        }

        inputs = {
            "Param": param_and_grad[0],
            "Grad": param_and_grad[1],
            "Velocity": velocity_acc,
            "LearningRate": lr
        }

        outputs = {"ParamOut": param_and_grad[0], "VelocityOut": velocity_acc}

        if find_master:
            inputs["MasterParam"] = master_weight
            outputs["MasterParamOut"] = master_weight

J
Jiabin Yang 已提交
2167
        if framework._non_static_mode():
D
duanboqiang 已提交
2168 2169 2170 2171 2172 2173
            tmp, tmp2 = _C_ops.lars_momentum(
                [param_and_grad[0]], [param_and_grad[1]], [velocity_acc], [lr],
                [param_and_grad[0]], [velocity_acc], "mu", self._momentum,
                "lars_coeff", self._lars_coeff, "lars_weight_decay",
                [_lars_weight_decay], "multi_precision", find_master, "epsilon",
                self._epsilon, "rescale_grad", self._rescale_grad)
2174 2175
        else:
            # create the momentum optimize op
2176 2177 2178 2179 2180
            momentum_op = block.append_op(type=self.type,
                                          inputs=inputs,
                                          outputs=outputs,
                                          attrs=attrs,
                                          stop_gradient=True)
2181

2182
            return momentum_op
2183 2184


2185
class AdagradOptimizer(Optimizer):
2186
    r"""
2187 2188
    The Adaptive Gradient optimizer (Adagrad for short) can adaptively assign
    different learning rates to individual parameters.
Q
qiaolongfei 已提交
2189

2190
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2191 2192 2193 2194 2195 2196 2197

    .. math::

        moment\_out &= moment + grad * grad

        param\_out &= param - \\frac{learning\_rate * grad}{\sqrt{moment\_out} + \epsilon}

2198 2199 2200 2201 2202 2203
    Related paper: `Adaptive Subgradient Methods for Online Learning and
    Stochastic Optimization <http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf>`_.

    The original paper does not have the ``epsilon`` attribute. It is added here
    in our implementation as also proposed `Per-parameter adaptive learning rate
    methods <http://cs231n.github.io/neural-networks-3/#ada>`_
Q
qiaolongfei 已提交
2204 2205 2206
    for numerical stability to avoid the division by zero error.

    Args:
2207 2208 2209 2210
        learning_rate (float|Variable): The learning rate used to update ``Parameter``.
            It can be a float value or a ``Variable`` with a float type.
        epsilon (float, optional): A small float value for numerical stability.
            The default value is 1e-06.
H
hong 已提交
2211
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2212 2213
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2214 2215 2216 2217 2218
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
2219 2220 2221 2222
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2223 2224 2225 2226 2227
        name (str, optional): Normally there is no need for user to set this property.
            For more information, please refer to :ref:`api_guide_Name`.
            The default value is None.
        initial_accumulator_value (float, optional): Initial value for moment accumulator.
            The default value is 0.0.
Q
qiaolongfei 已提交
2228 2229 2230 2231

    Examples:
        .. code-block:: python

2232
            import numpy as np
2233
            import paddle.fluid as fluid
2234 2235

            np_inp = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
2236
            inp = fluid.data(name="inp", shape=[2, 2])
2237 2238
            out = fluid.layers.fc(inp, size=3)
            out = fluid.layers.reduce_sum(out)
2239
            optimizer = fluid.optimizer.AdagradOptimizer(learning_rate=0.2)
2240 2241 2242 2243 2244 2245 2246
            optimizer.minimize(out)

            exe = fluid.Executor(fluid.CPUPlace())
            exe.run(fluid.default_startup_program())
            exe.run(
                feed={"inp": np_inp},
                fetch_list=[out.name])
2247 2248 2249
    """
    _moment_acc_str = "moment"

X
Xin Pan 已提交
2250 2251 2252
    def __init__(self,
                 learning_rate,
                 epsilon=1.0e-6,
2253
                 parameter_list=None,
X
Xin Pan 已提交
2254
                 regularization=None,
2255
                 grad_clip=None,
2256
                 name=None,
X
xuezhong 已提交
2257
                 initial_accumulator_value=0.0):
2258 2259
        assert learning_rate is not None
        assert epsilon is not None
2260 2261 2262 2263 2264
        super(AdagradOptimizer, self).__init__(learning_rate=learning_rate,
                                               parameter_list=parameter_list,
                                               regularization=regularization,
                                               grad_clip=grad_clip,
                                               name=name)
2265 2266
        self.type = "adagrad"
        self._epsilon = epsilon
2267
        self.initial_accumulator_value = initial_accumulator_value
2268 2269 2270 2271 2272

    def _create_accumulators(self, block, parameters):
        assert isinstance(block, framework.Block)

        for p in parameters:
2273 2274 2275
            self._add_accumulator(self._moment_acc_str,
                                  p,
                                  fill_value=self.initial_accumulator_value)
2276 2277 2278 2279 2280 2281

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)

        moment_acc = self._get_accumulator(self._moment_acc_str,
                                           param_and_grad[0])
C
caozhou 已提交
2282 2283 2284 2285 2286 2287 2288
        if in_dygraph_mode():
            _C_ops.final_state_adagrad_(param_and_grad[0], param_and_grad[1],
                                        moment_acc,
                                        self._create_param_lr(param_and_grad),
                                        self._epsilon)
            return None
        elif _in_legacy_dygraph():
2289 2290 2291 2292
            _C_ops.adagrad(param_and_grad[0], param_and_grad[1], moment_acc,
                           self._create_param_lr(param_and_grad),
                           param_and_grad[0], moment_acc, "epsilon",
                           self._epsilon)
C
caozhou 已提交
2293
            return None
2294 2295 2296 2297 2298 2299 2300 2301 2302 2303 2304 2305 2306 2307 2308 2309
        else:
            # Create the adagrad optimizer op
            adagrad_op = block.append_op(
                type=self.type,
                inputs={
                    "Param": param_and_grad[0],
                    "Grad": param_and_grad[1],
                    "Moment": moment_acc,
                    "LearningRate": self._create_param_lr(param_and_grad)
                },
                outputs={
                    "ParamOut": param_and_grad[0],
                    "MomentOut": moment_acc
                },
                attrs={"epsilon": self._epsilon},
                stop_gradient=True)
2310

2311
            return adagrad_op
2312 2313 2314


class AdamOptimizer(Optimizer):
2315
    r"""
T
tianshuo78520a 已提交
2316
    The Adam optimizer uses an optimization described at the end
2317 2318 2319 2320 2321
    of section 2 of `Adam paper <https://arxiv.org/abs/1412.6980>`_ ,
    it can dynamically adjusts the learning rate of each parameter using
    the 1st moment estimates and the 2nd moment estimates of the gradient.
    
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2322 2323 2324 2325 2326 2327 2328 2329 2330 2331 2332 2333 2334 2335

    .. math::

        t & = t + 1

        moment\_1\_out & = {\\beta}_1 * moment\_1 + (1 - {\\beta}_1) * grad

        moment\_2\_out & = {\\beta}_2 * moment\_2 + (1 - {\\beta}_2) * grad * grad

        learning\_rate & = learning\_rate * \\
                          \\frac{\sqrt{1 - {\\beta}_2^t}}{1 - {\\beta}_1^t}

        param\_out & = param - learning\_rate * \\frac{moment\_1}{\sqrt{moment\_2} + \epsilon}

2336 2337
    Related paper: `Adam: A Method for Stochastic Optimization <https://arxiv.org/abs/1412.6980>`_

Q
qiaolongfei 已提交
2338
    Args:
2339 2340
        learning_rate (float|Variable, optional): The learning rate used to update ``Parameter``.
            It can be a float value or a ``Variable`` with a float type. The default value is 0.001.
2341 2342
        beta1 (float|Variable, optional): The exponential decay rate for the 1st moment estimates.
            It should be a float number or a Variable with shape [1] and data type as float32.
2343
            The default value is 0.9.
2344 2345
        beta2 (float|Variable, optional): The exponential decay rate for the 2nd moment estimates.
            It should be a float number or a Variable with shape [1] and data type as float32.
2346
            The default value is 0.999.
2347 2348
        epsilon (float|Tensor, optional): A small float value for numerical stability.
            It should be a float number or a Variable with shape [1] and data type as float32.
2349
            The default value is 1e-08.
H
hong 已提交
2350
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2351 2352
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2353 2354 2355 2356 2357
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
2358 2359 2360 2361
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2362 2363 2364 2365 2366 2367 2368 2369 2370 2371
        name (str, optional): Normally there is no need for user to set this property.
            For more information, please refer to :ref:`api_guide_Name`.
            The default value is None.
        lazy_mode (bool, optional): The official Adam algorithm has two moving-average accumulators.
            The accumulators are updated at every step. Every element of the two moving-average
            is updated in both dense mode and sparse mode. If the size of parameter is very large,
            then the update may be very slow. The lazy mode only update the element that has
            gradient in current mini-batch, so it will be much more faster. But this mode has
            different semantics with the original Adam algorithm and may lead to different result.
            The default value is False.
2372 2373
        use_global_beta_pow (bool, optional): Whether to use global beta_pow. If true, Adam will use global beta_pow 
            for whole model instead of creating beta_pow for each parameter. Default is false.
2374 2375 2376
        flatten_param_grads (bool, optional): Whether to flatten all parameters and gradients. Default is false.
        align_size (int, optional): The alignment size when flatten parameters and gradients. Default is -1, which means
            use same align_size as allocator. 
Q
qiaolongfei 已提交
2377 2378 2379 2380

    Examples:
        .. code-block:: python

2381 2382 2383 2384 2385 2386
            import paddle
            import paddle.fluid as fluid

            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
2387 2388
                x = fluid.data(name='x', shape=[None, 13], dtype='float32')
                y = fluid.data(name='y', shape=[None, 1], dtype='float32')
2389 2390 2391 2392 2393 2394 2395 2396 2397 2398 2399 2400 2401 2402 2403
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
                cost = fluid.layers.square_error_cost(input=y_predict, label=y)
                avg_cost = fluid.layers.mean(cost)

                adam_optimizer = fluid.optimizer.AdamOptimizer(0.01)
                adam_optimizer.minimize(avg_cost)

                fetch_list = [avg_cost]
                train_reader = paddle.batch(
                    paddle.dataset.uci_housing.train(), batch_size=1)
                feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
                exe = fluid.Executor(place)
                exe.run(fluid.default_startup_program())
                for data in train_reader():
                    exe.run(main, feed=feeder.feed(data), fetch_list=fetch_list)
Q
qiaolongfei 已提交
2404

2405 2406 2407 2408 2409 2410 2411 2412 2413 2414 2415 2416 2417 2418 2419 2420 2421
        .. code-block:: python

            # Adam with beta1/beta2 as Variable
            import paddle
            import paddle.fluid as fluid
            import paddle.fluid.layers.learning_rate_scheduler as lr_scheduler

            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
                x = fluid.data(name='x', shape=[None, 13], dtype='float32')
                y = fluid.data(name='y', shape=[None, 1], dtype='float32')
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
                cost = fluid.layers.square_error_cost(input=y_predict, label=y)
                avg_cost = fluid.layers.mean(cost)

                # define beta decay variable
2422
                def get_decayed_betas(beta1_init, beta2_init, decay_steps, decay_rate, epsilon_init):
2423 2424 2425 2426 2427 2428 2429 2430 2431 2432 2433 2434 2435 2436 2437 2438
                    global_step = lr_scheduler._decay_step_counter()

                    beta1 = fluid.layers.create_global_var(
                        shape=[1],
                        value=float(beta1_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="beta1")
                    beta2 = fluid.layers.create_global_var(
                        shape=[1],
                        value=float(beta2_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="beta2")
2439 2440 2441 2442 2443 2444 2445
                    epsilon = fluid.layers.create_global_var(
                        shape=[1],
                        value=float(epsilon_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="epsilon")
2446 2447 2448 2449 2450 2451 2452

                    div_res = global_step / decay_steps
                    decayed_beta1 = beta1_init * (decay_rate**div_res)
                    decayed_beta2 = beta2_init * (decay_rate**div_res)
                    fluid.layers.assign(decayed_beta1, beta1)
                    fluid.layers.assign(decayed_beta2, beta2)

2453
                    return beta1, beta2, epsilon
2454

2455
                beta1, beta2, epsilon = get_decayed_betas(0.9, 0.99, 1e5, 0.9, 1e-8)
2456 2457
                adam_optimizer = fluid.optimizer.AdamOptimizer(
                                                    learning_rate=0.01,
2458
                                                    beta1=beta1,
2459 2460
                                                    beta2=beta2,
                                                    epsilon=epsilon)
2461 2462 2463 2464 2465 2466 2467 2468 2469 2470
                adam_optimizer.minimize(avg_cost)

                fetch_list = [avg_cost]
                train_reader = paddle.batch(
                    paddle.dataset.uci_housing.train(), batch_size=1)
                feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
                exe = fluid.Executor(place)
                exe.run(fluid.default_startup_program())
                for data in train_reader():
                    exe.run(main, feed=feeder.feed(data), fetch_list=fetch_list)
2471 2472 2473
    """
    _moment1_acc_str = "moment1"
    _moment2_acc_str = "moment2"
Q
qiaolongfei 已提交
2474 2475
    _beta1_pow_acc_str = "beta1_pow_acc"
    _beta2_pow_acc_str = "beta2_pow_acc"
2476 2477 2478 2479 2480

    def __init__(self,
                 learning_rate=0.001,
                 beta1=0.9,
                 beta2=0.999,
2481
                 epsilon=1e-8,
2482
                 parameter_list=None,
X
Xin Pan 已提交
2483
                 regularization=None,
2484
                 grad_clip=None,
Q
Qiao Longfei 已提交
2485
                 name=None,
2486
                 lazy_mode=False,
2487 2488 2489
                 use_global_beta_pow=False,
                 flatten_param_grads=False,
                 align_size=-1):
2490 2491 2492 2493
        assert learning_rate is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
2494 2495 2496 2497 2498 2499 2500 2501
        super(AdamOptimizer,
              self).__init__(learning_rate=learning_rate,
                             parameter_list=parameter_list,
                             regularization=regularization,
                             grad_clip=grad_clip,
                             flatten_param_grads=flatten_param_grads,
                             align_size=align_size,
                             name=name)
2502 2503 2504 2505
        self.type = "adam"
        self._beta1 = beta1
        self._beta2 = beta2
        self._epsilon = epsilon
Q
Qiao Longfei 已提交
2506
        self._lazy_mode = lazy_mode
2507
        self._use_global_beta_pow = use_global_beta_pow
2508 2509 2510 2511 2512 2513

    def _create_accumulators(self, block, parameters):
        assert isinstance(block, framework.Block)

        # Create accumulator tensors for first and second moments
        for p in parameters:
Q
Qiao Longfei 已提交
2514 2515
            self._add_accumulator(self._moment1_acc_str, p)
            self._add_accumulator(self._moment2_acc_str, p)
2516 2517 2518 2519 2520 2521 2522 2523 2524 2525 2526 2527 2528 2529 2530 2531 2532
            if not self._use_global_beta_pow:
                self._add_accumulator(
                    name=self._beta1_pow_acc_str,
                    param=p,
                    fill_value=0.9 if isinstance(self._beta1, Variable) \
                            else self._beta1,
                    shape=[1],
                    type=core.VarDesc.VarType.LOD_TENSOR, device='cpu')
                self._add_accumulator(
                    name=self._beta2_pow_acc_str,
                    param=p,
                    fill_value=0.999 if isinstance(self._beta2, Variable) \
                            else self._beta2,
                    shape=[1],
                    type=core.VarDesc.VarType.LOD_TENSOR, device='cpu')
        if self._use_global_beta_pow:
            self._add_global_accumulator(
Q
qiaolongfei 已提交
2533
                name=self._beta1_pow_acc_str,
2534 2535
                fill_value=0.9 if isinstance(self._beta1, Variable) \
                        else self._beta1,
2536
                shape=[1],
2537
                type=core.VarDesc.VarType.LOD_TENSOR, device='cpu')
2538
            self._add_global_accumulator(
Q
qiaolongfei 已提交
2539
                name=self._beta2_pow_acc_str,
2540 2541
                fill_value=0.999 if isinstance(self._beta2, Variable) \
                        else self._beta2,
2542
                shape=[1],
2543
                type=core.VarDesc.VarType.LOD_TENSOR, device='cpu')
2544 2545 2546 2547 2548 2549 2550 2551

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)

        moment1 = self._get_accumulator(self._moment1_acc_str,
                                        param_and_grad[0])
        moment2 = self._get_accumulator(self._moment2_acc_str,
                                        param_and_grad[0])
2552 2553 2554 2555 2556 2557 2558 2559 2560 2561
        if self._use_global_beta_pow:
            beta1_pow_acc = self._get_global_accumulator(
                self._beta1_pow_acc_str)
            beta2_pow_acc = self._get_global_accumulator(
                self._beta2_pow_acc_str)
        else:
            beta1_pow_acc = self._get_accumulator(self._beta1_pow_acc_str,
                                                  param_and_grad[0])
            beta2_pow_acc = self._get_accumulator(self._beta2_pow_acc_str,
                                                  param_and_grad[0])
2562
        lr = self._create_param_lr(param_and_grad)
2563
        # create the adam optimize op
2564

J
Jiabin Yang 已提交
2565
        if framework._non_static_mode():
2566 2567 2568 2569
            _beta1 = self._beta1 if not isinstance(
                self._beta1, Variable) else self._beta1.numpy().item(0)
            _beta2 = self._beta2 if not isinstance(
                self._beta2, Variable) else self._beta2.numpy().item(0)
2570 2571
            master_weight = None
            _, _, _, _, _, _ = _C_ops.adam(
2572
                param_and_grad[0], param_and_grad[1], lr, moment1, moment2,
2573 2574 2575 2576 2577
                beta1_pow_acc, beta2_pow_acc, master_weight, param_and_grad[0],
                moment1, moment2, beta1_pow_acc, beta2_pow_acc, master_weight,
                'epsilon', self._epsilon, 'lazy_mode', self._lazy_mode,
                'min_row_size_to_use_multithread', 1000, 'beta1', _beta1,
                'beta2', _beta2, 'use_global_beta_pow',
2578
                self._use_global_beta_pow)
2579 2580 2581

            return None

2582
        inputs = {
2583 2584
            "Param": [param_and_grad[0]],
            "Grad": [param_and_grad[1]],
2585
            "LearningRate": [lr],
2586 2587 2588 2589
            "Moment1": [moment1],
            "Moment2": [moment2],
            "Beta1Pow": [beta1_pow_acc],
            "Beta2Pow": [beta2_pow_acc]
2590
        }
2591 2592 2593 2594 2595 2596 2597

        # Pass found_inf to adam, to skip update for not only param, but also momentum and beta_pow
        found_inf = self._get_auxiliary_var('found_inf')

        if found_inf:
            inputs['SkipUpdate'] = found_inf

2598
        outputs = {
2599 2600 2601 2602 2603
            "ParamOut": [param_and_grad[0]],
            "Moment1Out": [moment1],
            "Moment2Out": [moment2],
            "Beta1PowOut": [beta1_pow_acc],
            "Beta2PowOut": [beta2_pow_acc],
2604 2605 2606
        }
        attrs = {
            "lazy_mode": self._lazy_mode,
2607 2608
            "min_row_size_to_use_multithread": 1000,
            'use_global_beta_pow': self._use_global_beta_pow
2609 2610 2611 2612 2613 2614 2615 2616 2617 2618
        }

        if isinstance(self._beta1, Variable):
            inputs['Beta1Tensor'] = self._beta1
        else:
            attrs['beta1'] = self._beta1
        if isinstance(self._beta2, Variable):
            inputs['Beta2Tensor'] = self._beta2
        else:
            attrs['beta2'] = self._beta2
2619 2620 2621 2622
        if isinstance(self._epsilon, Variable):
            inputs['EpsilonTensor'] = self._epsilon
        else:
            attrs['epsilon'] = self._epsilon
2623

2624 2625 2626 2627 2628
        adam_op = block.append_op(type=self.type,
                                  inputs=inputs,
                                  outputs=outputs,
                                  attrs=attrs,
                                  stop_gradient=True)
2629 2630 2631

        return adam_op

2632 2633 2634 2635 2636 2637 2638 2639 2640 2641 2642 2643
    def _finish_update(self, block, parameters_and_grads):
        r"""Update beta1_pow and beta2_pow accumulator
        """
        assert isinstance(block, framework.Block)
        if self._use_global_beta_pow:
            beta1_pow_acc = self._get_global_accumulator(
                self._beta1_pow_acc_str)
            beta2_pow_acc = self._get_global_accumulator(
                self._beta2_pow_acc_str)

            with block.program._optimized_guard([]):
                inputs = {"X": beta1_pow_acc}
2644
                outputs = {"Out": beta1_pow_acc}
2645 2646
                attrs = {}
                if isinstance(self._beta1, Variable):
2647 2648
                    inputs["Y"] = self._beta1
                    # use elementwise_mul for better performance
2649 2650 2651 2652 2653
                    block.append_op(type="elementwise_mul",
                                    inputs=inputs,
                                    outputs=outputs,
                                    attrs=attrs,
                                    stop_gradient=True)
2654 2655
                else:
                    attrs['scale'] = self._beta1
2656 2657 2658 2659 2660
                    block.append_op(type="scale",
                                    inputs=inputs,
                                    outputs=outputs,
                                    attrs=attrs,
                                    stop_gradient=True)
2661 2662

                inputs = {"X": beta2_pow_acc}
2663
                outputs = {"Out": beta2_pow_acc}
2664 2665
                attrs = {}
                if isinstance(self._beta2, Variable):
2666 2667
                    inputs["Y"] = self._beta2
                    # use elementwise_mul for better performance
2668 2669 2670 2671 2672
                    block.append_op(type="elementwise_mul",
                                    inputs=inputs,
                                    outputs=outputs,
                                    attrs=attrs,
                                    stop_gradient=True)
2673 2674
                else:
                    attrs['scale'] = self._beta2
2675 2676 2677 2678 2679
                    block.append_op(type="scale",
                                    inputs=inputs,
                                    outputs=outputs,
                                    attrs=attrs,
                                    stop_gradient=True)
2680

2681 2682

class AdamaxOptimizer(Optimizer):
2683
    r"""
2684 2685 2686 2687
    The Adamax optimizer is implemented based on the Adamax Optimization 
    in Section 7 of `Adam paper <https://arxiv.org/abs/1412.6980>`_.
    The Adamax algorithm is a variant of the Adam algorithm based on the infinite norm,
    which makes the learning rate update algorithm more stable and simple.
Q
qiaolongfei 已提交
2688

2689
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2690 2691 2692 2693 2694 2695 2696 2697 2698 2699 2700 2701 2702

    .. math::

        t & = t + 1

        moment\_out & = {\\beta}_1 * moment + (1 - {\\beta}_1) * grad

        inf\_norm\_out & = max({\\beta}_2 * inf\_norm + \epsilon, |grad|)

        learning\_rate & = \\frac{learning\_rate}{1 - {\\beta}_1^t}

        param\_out & = param - learning\_rate * \\frac{moment\_out}{inf\_norm\_out}

2703
    Related paper: `Adam: A Method for Stochastic Optimization <https://arxiv.org/abs/1412.6980>`_
Q
qiaolongfei 已提交
2704

2705 2706 2707 2708 2709 2710 2711 2712 2713 2714 2715 2716
    The original paper does not have an ``epsilon`` attribute,
    it is added here for numerical stability to prevent the division by 0 error.

    Args:
        learning_rate (float|Variable, optional): The learning rate used to update ``Parameter``.
            It can be a float value or a ``Variable`` with a float type. The default value is 0.001.
        beta1 (float, optional): The exponential decay rate for the 1st moment estimates.
            The default value is 0.9.
        beta2 (float, optional): The exponential decay rate for the 2nd moment estimates.
            The default value is 0.999.
        epsilon (float, optional): A small float value for numerical stability.
            The default value is 1e-08.
H
hong 已提交
2717
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2718 2719
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2720 2721 2722 2723 2724
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
2725 2726 2727 2728
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2729 2730 2731 2732 2733 2734
        name (str, optional): Normally there is no need for user to set this property.
            For more information, please refer to :ref:`api_guide_Name`.
            The default value is None.

    **Notes**:
        **Currently, AdamaxOptimizer doesn't support sparse parameter optimization.**
Q
qiaolongfei 已提交
2735

2736 2737 2738 2739 2740 2741 2742 2743 2744 2745 2746 2747 2748
    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          import numpy

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

          train_program = fluid.Program()
          startup_program = fluid.Program()
          with fluid.program_guard(train_program, startup_program):
2749
              data = fluid.data(name='X', shape=[None, 1], dtype='float32')
2750 2751
              hidden = fluid.layers.fc(input=data, size=10)
              loss = fluid.layers.mean(hidden)
2752
              adam = fluid.optimizer.AdamaxOptimizer(learning_rate=0.2)
2753 2754 2755 2756 2757 2758 2759 2760 2761
              adam.minimize(loss)

          # Run the startup program once and only once.
          exe.run(startup_program)

          x = numpy.random.random(size=(10, 1)).astype('float32')
          outs = exe.run(program=train_program,
                        feed={'X': x},
                         fetch_list=[loss.name])
2762 2763 2764
    """
    _moment_acc_str = "moment"
    _inf_norm_acc_str = "inf_norm"
Q
qiaolongfei 已提交
2765
    _beta1_pow_acc_str = "beta1_pow_acc"
2766 2767 2768 2769 2770

    def __init__(self,
                 learning_rate=0.001,
                 beta1=0.9,
                 beta2=0.999,
2771
                 epsilon=1e-8,
2772
                 parameter_list=None,
X
Xin Pan 已提交
2773
                 regularization=None,
2774
                 grad_clip=None,
X
Xin Pan 已提交
2775
                 name=None):
2776 2777 2778 2779
        assert learning_rate is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
2780 2781 2782 2783 2784
        super(AdamaxOptimizer, self).__init__(learning_rate=learning_rate,
                                              parameter_list=parameter_list,
                                              regularization=regularization,
                                              grad_clip=grad_clip,
                                              name=name)
2785 2786 2787 2788 2789 2790 2791 2792
        self.type = "adamax"
        self._beta1 = beta1
        self._beta2 = beta2
        self._epsilon = epsilon

    def _create_accumulators(self, block, parameters):
        # Create accumulator tensors for first moment and infinity norm
        for p in parameters:
Q
Qiao Longfei 已提交
2793 2794
            self._add_accumulator(self._moment_acc_str, p)
            self._add_accumulator(self._inf_norm_acc_str, p)
2795 2796 2797 2798
            self._add_accumulator(name=self._beta1_pow_acc_str,
                                  param=p,
                                  fill_value=self._beta1,
                                  shape=[1])
2799 2800 2801 2802 2803 2804 2805

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)

        moment = self._get_accumulator(self._moment_acc_str, param_and_grad[0])
        inf_norm = self._get_accumulator(self._inf_norm_acc_str,
                                         param_and_grad[0])
Q
qiaolongfei 已提交
2806 2807
        beta1_pow_acc = self._get_accumulator(self._beta1_pow_acc_str,
                                              param_and_grad[0])
J
Jiabin Yang 已提交
2808
        if framework._non_static_mode():
2809 2810 2811 2812 2813 2814 2815 2816 2817 2818 2819 2820 2821 2822 2823 2824 2825 2826 2827 2828 2829 2830 2831 2832 2833 2834 2835 2836
            _C_ops.adamax(param_and_grad[0], param_and_grad[1],
                          self._create_param_lr(param_and_grad), moment,
                          inf_norm, beta1_pow_acc, param_and_grad[0], moment,
                          inf_norm, "beta1", self._beta1, "beta2", self._beta2,
                          "epsilon", self._epsilon)
        else:
            # create the adamax optimize op
            adamax_op = block.append_op(
                type=self.type,
                inputs={
                    "Param": param_and_grad[0],
                    "Grad": param_and_grad[1],
                    "LearningRate": self._create_param_lr(param_and_grad),
                    "Moment": moment,
                    "InfNorm": inf_norm,
                    "Beta1Pow": beta1_pow_acc
                },
                outputs={
                    "ParamOut": param_and_grad[0],
                    "MomentOut": moment,
                    "InfNormOut": inf_norm
                },
                attrs={
                    "beta1": self._beta1,
                    "beta2": self._beta2,
                    "epsilon": self._epsilon
                },
                stop_gradient=True)
2837

2838
            return adamax_op
2839

2840
    def _finish_update(self, block, parameters_and_grads):
2841 2842 2843
        """Update Beta1 Power accumulator
        """
        assert isinstance(block, framework.Block)
2844
        for param, grad in parameters_and_grads:
C
chengduo 已提交
2845
            if grad is None or param.trainable is False:
2846
                continue
2847 2848
            with param.block.program._optimized_guard([param, grad
                                                       ]), name_scope('adamx'):
2849 2850
                beta1_pow_acc = self._get_accumulator(self._beta1_pow_acc_str,
                                                      param)
J
Jiabin Yang 已提交
2851
                if framework._non_static_mode():
2852 2853 2854 2855 2856
                    if framework.in_dygraph_mode():
                        tmp = _C_ops.final_state_scale(beta1_pow_acc,
                                                       self._beta1, 0.0, True)
                    else:
                        tmp = _C_ops.scale(beta1_pow_acc, "scale", self._beta1)
2857 2858
                    beta1_pow_acc.copy_(tmp, False)
                else:
2859 2860 2861 2862 2863
                    block.append_op(type="scale",
                                    inputs={"X": beta1_pow_acc},
                                    outputs={"Out": beta1_pow_acc},
                                    attrs={"scale": self._beta1},
                                    stop_gradient=True)
2864 2865


2866
class DpsgdOptimizer(Optimizer):
2867
    r"""
2868 2869 2870 2871 2872 2873 2874 2875 2876 2877 2878 2879 2880 2881 2882 2883 2884 2885 2886 2887 2888 2889 2890 2891 2892 2893 2894 2895 2896 2897 2898 2899 2900 2901 2902 2903
    We implement the Dpsgd optimizer according to CCS16 paper -
    Deep Learning with Differential Privacy.

    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          import numpy

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

          train_program = fluid.Program()
          startup_program = fluid.Program()
          with fluid.program_guard(train_program, startup_program):
              data = fluid.layers.data(name='X', shape=[1], dtype='float32')
              hidden = fluid.layers.fc(input=data, size=10)
              loss = fluid.layers.mean(hidden)
              optimizer = fluid.optimizer.Dpsgd(learning_rate=0.01, clip=10.0, batch_size=16.0, sigma=1.0)
              optimizer.minimize(loss)

          # Run the startup program once and only once.
          exe.run(startup_program)

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

    Args:
        learning_rate (float|Variable): the learning rate used to update parameters. \
        Can be a float value or a Variable with one float value as data element.
        clip (float): clipping threshold
        batch_size (float): batch size.
        sigma (float): for gaussian noise.
H
hong 已提交
2904
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2905 2906
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2907 2908 2909 2910 2911 2912 2913 2914
    Notes:
       Currently, DpsgdOptimizer doesn't support sparse parameter optimization.
    """

    def __init__(self,
                 learning_rate=0.001,
                 clip=0.9,
                 batch_size=0.999,
2915 2916
                 sigma=1e-8,
                 parameter_list=None):
2917 2918 2919 2920
        assert learning_rate is not None
        assert clip is not None
        assert batch_size is not None
        assert sigma is not None
2921 2922
        super(DpsgdOptimizer, self).__init__(learning_rate=learning_rate,
                                             parameter_list=parameter_list)
2923 2924 2925 2926
        self.type = "dpsgd"
        self._clip = clip
        self._batch_size = batch_size
        self._sigma = sigma
Z
zhongpu 已提交
2927 2928 2929 2930 2931 2932 2933
        '''
        Note(wangzhongpu):
        This property is only used for debugging, do not need to set it!
        Dpsgd operator use time(NULL) as random seed to generate random number.
        However, during debugging, we need determinated result, so we will set self._seed to a fixed number.
        '''
        self._seed = None
2934 2935 2936 2937 2938

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)

        # create the dpsgd optimize op
Z
zhongpu 已提交
2939 2940 2941
        if self._seed == None:
            self._seed = 0

J
Jiabin Yang 已提交
2942
        if framework._non_static_mode():
2943 2944 2945 2946 2947 2948
            _C_ops.dpsgd(param_and_grad[0], param_and_grad[1],
                         self._create_param_lr(param_and_grad),
                         param_and_grad[0], "clip", self._clip, "batch_size",
                         self._batch_size, "sigma", self._sigma, "seed",
                         self._seed)
        else:
2949 2950 2951 2952 2953 2954 2955 2956 2957 2958 2959 2960 2961 2962 2963 2964 2965
            dpsgd_op = block.append_op(type=self.type,
                                       inputs={
                                           "Param":
                                           param_and_grad[0],
                                           "Grad":
                                           param_and_grad[1],
                                           "LearningRate":
                                           self._create_param_lr(param_and_grad)
                                       },
                                       outputs={"ParamOut": param_and_grad[0]},
                                       attrs={
                                           "clip": self._clip,
                                           "batch_size": self._batch_size,
                                           "sigma": self._sigma,
                                           "seed": self._seed
                                       },
                                       stop_gradient=True)
2966

2967
            return dpsgd_op
2968 2969


2970
class DecayedAdagradOptimizer(Optimizer):
2971
    r"""
2972 2973 2974
    The Decayed Adagrad optimizer can be seen as an Adagrad algorithm that introduces
    the decay rate to solve the problem of a sharp drop in the learning rate
    during model training when using the AdagradOptimizer.
2975

2976
    The parameter ``param_out`` update rule with gradient ``grad``:
2977 2978 2979 2980 2981 2982 2983

    .. math::

        moment\_out & = decay * moment + (1 - decay) * grad * grad

        param\_out & = param - \\frac{learning\_rate * grad}{\sqrt{moment\_out} + \epsilon}

2984 2985 2986 2987
    Related paper: `Adaptive Subgradient Methods for Online Learning and Stochastic
    Optimization <http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf>`_.

    The original paper does not have an ``epsilon`` attribute. It is added here for numerical
2988 2989 2990
    stability to avoid the division by zero error.

    Args:
2991 2992 2993 2994 2995
        learning_rate (float|Variable): The learning rate used to update ``Parameter``.
            It can be a float value or a ``Variable`` with a float type.
        decay (float, optional): The decay rate. The default value is 0.95.
        epsilon (float, optional): A small float value for numerical stability.
            The default value is 1e-06.
H
hong 已提交
2996
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2997 2998
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2999 3000 3001 3002 3003
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
3004 3005 3006 3007
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3008 3009 3010 3011 3012 3013
        name (str, optional): Normally there is no need for user to set this property.
            For more information, please refer to :ref:`api_guide_Name`.
            The default value is None.

    **Notes**:
        **Currently, DecayedAdagradOptimizer doesn't support sparse parameter optimization.**
3014 3015 3016 3017

    Examples:
        .. code-block:: python

3018 3019
            import paddle.fluid as fluid

3020 3021 3022 3023
            x = fluid.data( name='x', shape=[None, 10], dtype='float32' )
            trans = fluid.layers.fc( x, 100 )
            cost = fluid.layers.reduce_mean( trans )
            optimizer = fluid.optimizer.DecayedAdagradOptimizer(learning_rate=0.2)
3024
            optimizer.minimize(cost)
3025 3026 3027
    """
    _moment_acc_str = "moment"

X
Xin Pan 已提交
3028 3029 3030 3031
    def __init__(self,
                 learning_rate,
                 decay=0.95,
                 epsilon=1.0e-6,
3032
                 parameter_list=None,
X
Xin Pan 已提交
3033
                 regularization=None,
3034
                 grad_clip=None,
X
Xin Pan 已提交
3035
                 name=None):
3036 3037 3038 3039
        assert learning_rate is not None
        assert decay is not None
        assert epsilon is not None

3040 3041 3042 3043 3044 3045
        super(DecayedAdagradOptimizer,
              self).__init__(learning_rate=learning_rate,
                             parameter_list=parameter_list,
                             regularization=regularization,
                             grad_clip=grad_clip,
                             name=name)
3046 3047 3048 3049 3050 3051 3052 3053 3054 3055 3056 3057 3058 3059 3060 3061
        self.type = "decayed_adagrad"
        self._decay = decay
        self._epsilon = epsilon

    def _create_accumulators(self, block, parameters):
        assert isinstance(block, framework.Block)

        for p in parameters:
            self._add_accumulator(self._moment_acc_str, p)

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)

        moment_acc = self._get_accumulator(self._moment_acc_str,
                                           param_and_grad[0])

J
Jiabin Yang 已提交
3062
        if framework._non_static_mode():
3063 3064 3065 3066 3067
            _C_ops.decayed_adagrad(param_and_grad[0], param_and_grad[1],
                                   moment_acc,
                                   self._create_param_lr(param_and_grad),
                                   param_and_grad[0], moment_acc, "epsilon",
                                   self._epsilon, "decay", self._decay)
3068 3069 3070 3071 3072 3073 3074 3075 3076 3077 3078 3079 3080 3081
        else:
            # Create the decayed adagrad optimizer op
            decayed_adagrad_op = block.append_op(
                type=self.type,
                inputs={
                    "Param": param_and_grad[0],
                    "Grad": param_and_grad[1],
                    "Moment": moment_acc,
                    "LearningRate": self._create_param_lr(param_and_grad)
                },
                outputs={
                    "ParamOut": param_and_grad[0],
                    "MomentOut": moment_acc
                },
3082 3083 3084 3085
                attrs={
                    "epsilon": self._epsilon,
                    "decay": self._decay
                },
3086
                stop_gradient=True)
3087

3088
            return decayed_adagrad_op
3089 3090


3091
class AdadeltaOptimizer(Optimizer):
3092
    r"""
Z
Zeng Jinle 已提交
3093
    **Notes: This API does not support sparse parameter optimization.**
Q
qiaolongfei 已提交
3094

Z
Zeng Jinle 已提交
3095
    Adadelta Optimizer. Please refer to this for details:
Z
Zeng Jinle 已提交
3096 3097 3098
    `ADADELTA: AN ADAPTIVE LEARNING RATE METHOD <https://arxiv.org/abs/1212.5701>`_.

    The update is done as follows:
3099

Z
Zeng Jinle 已提交
3100 3101
    .. math::

Z
Zeng Jinle 已提交
3102
        E(g_t^2) &= \\rho * E(g_{t-1}^2) + (1-\\rho) * g^2
3103

Z
Zeng Jinle 已提交
3104
        learning\_rate &= \sqrt{ ( E(dx_{t-1}^2) + \\epsilon ) / ( E(g_t^2) + \\epsilon ) }
Z
Zeng Jinle 已提交
3105

Z
Zeng Jinle 已提交
3106
        E(dx_t^2) &= \\rho * E(dx_{t-1}^2) + (1-\\rho) * (-g*learning\_rate)^2
3107 3108

    Args:
Z
Zeng Jinle 已提交
3109 3110 3111
        learning_rate (float|Variable): global learning rate.
        epsilon (float): a small float number for numeric stability. Default 1.0e-6.
        rho (float): a floating point value indicating the decay rate. Default 0.95.
H
hong 已提交
3112
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3113 3114
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3115 3116 3117 3118 3119
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
3120 3121 3122 3123
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3124 3125 3126
        name (str, optional): The default value is None. Normally there is no need for user
                to set this property. For more information, please refer to
                :ref:`api_guide_Name` .
3127 3128 3129 3130

    Examples:
        .. code-block:: python

3131
            import paddle.fluid as fluid
Z
Zeng Jinle 已提交
3132

3133
            image = fluid.data(name='image', shape=[None, 28], dtype='float32')
Z
Zeng Jinle 已提交
3134 3135
            fc = fluid.layers.fc(image, size=10)
            cost = fluid.layers.reduce_mean(fc)
3136 3137
            optimizer = fluid.optimizer.Adadelta(
                learning_rate=0.0003, epsilon=1.0e-6, rho=0.95)
C
chengduo 已提交
3138

Z
Zeng Jinle 已提交
3139 3140 3141 3142
            # optimizer_ops is a list of optimizer operators to update parameters
            # params_grads is a list of (param, param_grad), where param is each
            # parameter and param_grad is the gradient variable of param.
            optimizer_ops, params_grads = optimizer.minimize(cost)
3143
    """
3144

3145 3146 3147
    _avg_squared_grad_acc_str = "_avg_squared_grad"
    _avg_squared_update_acc_str = "_avg_squared_update"

X
Xin Pan 已提交
3148 3149 3150 3151
    def __init__(self,
                 learning_rate,
                 epsilon=1.0e-6,
                 rho=0.95,
3152
                 parameter_list=None,
X
Xin Pan 已提交
3153
                 regularization=None,
3154
                 grad_clip=None,
X
Xin Pan 已提交
3155
                 name=None):
3156 3157 3158 3159 3160 3161
        if learning_rate is None:
            raise ValueError("learning_rate is not set.")
        if epsilon is None:
            raise ValueError("epsilon is not set.")
        if rho is None:
            raise ValueError("rho is not set.")
3162 3163 3164 3165 3166
        super(AdadeltaOptimizer, self).__init__(learning_rate=learning_rate,
                                                parameter_list=parameter_list,
                                                regularization=regularization,
                                                grad_clip=grad_clip,
                                                name=name)
3167 3168 3169 3170 3171
        self.type = "adadelta"
        self._epsilon = epsilon
        self._rho = rho

    def _create_accumulators(self, block, parameters):
3172 3173
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")
3174 3175 3176 3177 3178 3179

        for p in parameters:
            self._add_accumulator(self._avg_squared_grad_acc_str, p)
            self._add_accumulator(self._avg_squared_update_acc_str, p)

    def _append_optimize_op(self, block, param_and_grad):
3180 3181
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")
3182 3183 3184 3185 3186 3187

        avg_squared_grad_acc = self._get_accumulator(
            self._avg_squared_grad_acc_str, param_and_grad[0])
        avg_squared_update_acc = self._get_accumulator(
            self._avg_squared_update_acc_str, param_and_grad[0])

J
Jiabin Yang 已提交
3188
        if framework._non_static_mode():
3189 3190 3191 3192 3193 3194 3195
            _C_ops.adadelta(param_and_grad[0], param_and_grad[1],
                            avg_squared_grad_acc, avg_squared_update_acc,
                            param_and_grad[0], avg_squared_grad_acc,
                            avg_squared_update_acc, "epsilon", self._epsilon,
                            "rho", self._rho)
        else:
            # Create the adadelta optimizer op
3196 3197 3198 3199 3200 3201 3202 3203 3204 3205 3206 3207 3208 3209 3210 3211 3212 3213 3214 3215 3216 3217 3218 3219
            adadelta_op = block.append_op(type=self.type,
                                          inputs={
                                              "Param":
                                              param_and_grad[0],
                                              "Grad":
                                              param_and_grad[1],
                                              "AvgSquaredGrad":
                                              avg_squared_grad_acc,
                                              "AvgSquaredUpdate":
                                              avg_squared_update_acc
                                          },
                                          outputs={
                                              "ParamOut":
                                              param_and_grad[0],
                                              "AvgSquaredGradOut":
                                              avg_squared_grad_acc,
                                              "AvgSquaredUpdateOut":
                                              avg_squared_update_acc
                                          },
                                          attrs={
                                              "epsilon": self._epsilon,
                                              "rho": self._rho
                                          },
                                          stop_gradient=True)
3220

3221
            return adadelta_op
3222 3223


Q
qingqing01 已提交
3224
class RMSPropOptimizer(Optimizer):
3225
    r"""
Q
qingqing01 已提交
3226 3227 3228 3229 3230 3231 3232 3233
    Root Mean Squared Propagation (RMSProp) is an unpublished, adaptive learning
    rate method. The original slides proposed RMSProp: Slide 29 of
    http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf .

    The original equation is as follows:

    ..  math::

Q
qiaolongfei 已提交
3234
        r(w, t) & = \\rho r(w, t-1) + (1 - \\rho)(\\nabla Q_{i}(w))^2
Q
qingqing01 已提交
3235 3236 3237 3238

        w & = w - \\frac{\\eta} {\\sqrt{r(w,t) + \\epsilon}} \\nabla Q_{i}(w)

    The first equation calculates moving average of the squared gradient for
Q
qiaolongfei 已提交
3239
    each weight. Then dividing the gradient by :math:`sqrt{v(w,t)}`.
Q
qingqing01 已提交
3240 3241 3242 3243 3244 3245

    In some cases, adding a momentum term :math: `\\beta` is beneficial.
    In our implementation, Nesterov momentum is used:

    ..  math::

Q
qiaolongfei 已提交
3246
        r(w, t) & = \\rho r(w, t-1) + (1 - \\rho)(\\nabla Q_{i}(w))^2
Q
qingqing01 已提交
3247

3248 3249 3250 3251 3252 3253 3254 3255 3256 3257 3258 3259 3260 3261
        v(w, t) & = \\beta v(w, t-1) + \\frac{\\eta} {\\sqrt{r(w,t) +
            \\epsilon}} \\nabla Q_{i}(w)

        w & = w - v(w, t)

    if centered is True:

    ..  math::

        r(w, t) & = \\rho r(w, t-1) + (1 - \\rho)(\\nabla Q_{i}(w))^2

        g(w, t) & = \\rho g(w, t-1) + (1 - \\rho)\\nabla Q_{i}(w)

        v(w, t) & = \\beta v(w, t-1) + \\frac{\\eta} {\\sqrt{r(w,t) - (g(w, t))^2 +
Q
qingqing01 已提交
3262 3263 3264 3265
            \\epsilon}} \\nabla Q_{i}(w)

        w & = w - v(w, t)

Q
qiaolongfei 已提交
3266
    where, :math:`\\rho` is a hyperparameter and typical values are 0.9, 0.95
Q
qingqing01 已提交
3267 3268 3269 3270 3271
    and so on. :math: `beta` is the momentum term. :math: `\\epsilon` is a
    smoothing term to avoid division by zero, usually set somewhere in range
    from 1e-4 to 1e-8.


3272 3273 3274
    Parameters:
        learning_rate(float): Global learning rate.
        rho(float): rho is :math: `\\rho` in equation, default is 0.95.
Q
qingqing01 已提交
3275
        epsilon(float): :math: `\\epsilon` in equation is smoothing term to
3276
            avoid division by zero, default is 1e-6.
Q
qiaolongfei 已提交
3277
        momentum(float): :math:`\\beta` in equation is the momentum term,
3278
            default is 0.0.
3279 3280 3281 3282
        centered(bool): If True, gradients are normalized by the estimated variance of
            the gradient; if False, by the uncentered second moment. Setting this to
            True may help with training, but is slightly more expensive in terms of
            computation and memory. Defaults to False.
H
hong 已提交
3283
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3284 3285
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3286 3287 3288 3289 3290
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
3291 3292 3293 3294
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3295 3296
        name (str, optional): This parameter is used by developers to print debugging information. \
            For details, please refer to :ref:`api_guide_Name`. Default is None.
Q
qingqing01 已提交
3297 3298 3299 3300 3301 3302 3303

    Raises:
        ValueError: If learning_rate, rho, epsilon, momentum are None.

    Examples:
          .. code-block:: python

3304 3305 3306 3307 3308 3309 3310 3311 3312 3313 3314 3315 3316 3317 3318 3319 3320 3321 3322 3323 3324 3325 3326 3327 3328
            import paddle
            import paddle.fluid as fluid
            import numpy as np

            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
                x = fluid.layers.data(name='x', shape=[13], dtype='float32')
                y = fluid.layers.data(name='y', shape=[1], dtype='float32')
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
                cost = fluid.layers.square_error_cost(input=y_predict, label=y)
                avg_cost = fluid.layers.mean(cost)

                rms_optimizer = fluid.optimizer.RMSProp(learning_rate=0.1)
                rms_optimizer.minimize(avg_cost)

                fetch_list = [avg_cost]
                train_reader = paddle.batch(
                    paddle.dataset.uci_housing.train(), batch_size=1)
                feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
                exe = fluid.Executor(place)
                exe.run(fluid.default_startup_program())
                for data in train_reader():
                    exe.run(main, feed=feeder.feed(data), fetch_list=fetch_list)

Q
qingqing01 已提交
3329 3330 3331 3332
    """

    _momentum_acc_str = "momentum"
    _mean_square_acc_str = "mean_square"
3333
    _mean_grad_acc_str = "mean_grad"
Q
qingqing01 已提交
3334 3335 3336 3337 3338 3339

    def __init__(self,
                 learning_rate,
                 rho=0.95,
                 epsilon=1.0e-6,
                 momentum=0.0,
3340
                 centered=False,
3341
                 parameter_list=None,
X
Xin Pan 已提交
3342
                 regularization=None,
3343
                 grad_clip=None,
X
Xin Pan 已提交
3344
                 name=None):
3345 3346 3347 3348 3349
        super(RMSPropOptimizer, self).__init__(learning_rate=learning_rate,
                                               parameter_list=parameter_list,
                                               regularization=regularization,
                                               grad_clip=grad_clip,
                                               name=name)
Q
qingqing01 已提交
3350 3351 3352 3353 3354 3355 3356 3357 3358 3359 3360 3361 3362
        if learning_rate is None:
            raise ValueError("learning_rate is not set.")
        if rho is None:
            raise ValueError("rho is not set.")
        if epsilon is None:
            raise ValueError("epsilon is not set.")
        if momentum is None:
            raise ValueError("momentum is not set.")

        self.type = "rmsprop"
        self._rho = rho
        self._epsilon = epsilon
        self._momentum = momentum
3363
        self._centered = centered
Q
qingqing01 已提交
3364 3365 3366 3367 3368 3369 3370 3371

    def _create_accumulators(self, block, parameters):
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")

        for p in parameters:
            self._add_accumulator(self._momentum_acc_str, p)
            self._add_accumulator(self._mean_square_acc_str, p)
3372
            self._add_accumulator(self._mean_grad_acc_str, p)
Q
qingqing01 已提交
3373 3374 3375 3376 3377 3378 3379 3380 3381

    def _append_optimize_op(self, block, param_and_grad):
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")

        momentum_acc = self._get_accumulator(self._momentum_acc_str,
                                             param_and_grad[0])
        mean_square_acc = self._get_accumulator(self._mean_square_acc_str,
                                                param_and_grad[0])
3382 3383
        mean_grad_acc = self._get_accumulator(self._mean_grad_acc_str,
                                              param_and_grad[0])
C
caozhou 已提交
3384 3385 3386 3387 3388 3389 3390 3391
        if in_dygraph_mode():
            _C_ops.final_state_rmsprop_(param_and_grad[0], mean_square_acc,
                                        param_and_grad[1], momentum_acc,
                                        self._create_param_lr(param_and_grad),
                                        mean_grad_acc, self._epsilon, self._rho,
                                        self._momentum, self._centered)
            return None
        elif _in_legacy_dygraph():
3392 3393 3394 3395 3396 3397 3398
            _C_ops.rmsprop(param_and_grad[0], mean_square_acc,
                           self._create_param_lr(param_and_grad),
                           param_and_grad[1], momentum_acc, param_and_grad[0],
                           momentum_acc, mean_square_acc, mean_grad_acc,
                           "epsilon", self._epsilon, "decay", self._rho,
                           "momentum", self._momentum, "centered",
                           self._centered)
C
caozhou 已提交
3399
            return None
3400 3401 3402 3403 3404 3405 3406 3407 3408 3409 3410 3411 3412 3413 3414 3415 3416 3417 3418 3419 3420 3421 3422 3423
        else:
            rmsprop_op = block.append_op(
                type=self.type,
                inputs={
                    "Param": param_and_grad[0],
                    "Grad": param_and_grad[1],
                    "Moment": momentum_acc,
                    "MeanSquare": mean_square_acc,
                    "MeanGrad": mean_grad_acc,
                    "LearningRate": self._create_param_lr(param_and_grad),
                },
                outputs={
                    "ParamOut": param_and_grad[0],
                    "MomentOut": momentum_acc,
                    "MeanSquareOut": mean_square_acc,
                    "MeanGradOut": mean_grad_acc
                },
                attrs={
                    "epsilon": self._epsilon,
                    "decay": self._rho,
                    "momentum": self._momentum,
                    "centered": self._centered
                },
                stop_gradient=True)
Q
qingqing01 已提交
3424

3425
            return rmsprop_op
Q
qingqing01 已提交
3426 3427


Q
qiaolongfei 已提交
3428
class FtrlOptimizer(Optimizer):
3429
    r"""
Q
qiaolongfei 已提交
3430 3431 3432 3433 3434 3435 3436 3437 3438 3439 3440 3441 3442 3443 3444 3445 3446 3447 3448 3449 3450 3451 3452 3453 3454 3455 3456 3457 3458 3459 3460 3461 3462 3463 3464 3465 3466 3467
    FTRL (Follow The Regularized Leader) Optimizer.

    The paper that proposed Follow The Regularized Leader (FTRL):
    (https://www.eecs.tufts.edu/~dsculley/papers/ad-click-prediction.pdf)

    ..  math::

        &new\_accum = squared\_accum + grad^2

        &if (lr\_power == -0.5):

        &\quad  linear\_accum += grad - \\frac{\\sqrt{new\_accum} - \\sqrt{squared\_accum}}{learning\_rate * param}

        &else:

        &\quad   linear\_accum += grad - \\frac{new\_accum^{-lr\_power} - accum^{-lr\_power}}{learning\_rate * param}


        &x = l1 * sign(linear\_accum) - linear\_accum

        &if (lr\_power == -0.5):

        &\quad   y = \\frac{\\sqrt{new\_accum}}{learning\_rate} + (2 * l2)

        &\quad   pre\_shrink = \\frac{x}{y}

        &\quad   param = (abs(linear\_accum) > l1).select(pre\_shrink, 0.0)

        &else:

        &\quad   y = \\frac{new\_accum^{-lr\_power}}{learning\_rate} + (2 * l2)

        &\quad   pre\_shrink = \\frac{x}{y}

        &\quad   param = (abs(linear\_accum) > l1).select(pre\_shrink, 0.0)

        &squared\_accum += grad^2

3468 3469 3470 3471 3472
    Parameters:
        learning_rate (float|Variable): Global learning rate.
        l1 (float): L1 regularization strength, default is 0.0.
        l2 (float): L2 regularization strength, default is 0.0.
        lr_power (float): Learning Rate Power, default is -0.5.
H
hong 已提交
3473
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3474 3475
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3476 3477 3478 3479 3480
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
3481 3482 3483 3484
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
            ( :ref:`api_fluid_clip_GradientClipByGlobalNorm` , :ref:`api_fluid_clip_GradientClipByNorm` , 
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3485 3486
        name (str, optional): This parameter is used by developers to print debugging information. \
            For details, please refer to :ref:`api_guide_Name`. Default is None.
Q
qiaolongfei 已提交
3487 3488 3489 3490 3491 3492 3493

    Raises:
        ValueError: If learning_rate, rho, epsilon, momentum are None.

    Examples:
          .. code-block:: python

3494 3495 3496 3497 3498 3499 3500 3501 3502 3503 3504 3505 3506 3507 3508 3509 3510 3511 3512 3513 3514 3515 3516 3517
            import paddle
            import paddle.fluid as fluid
            import numpy as np

            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
                x = fluid.layers.data(name='x', shape=[13], dtype='float32')
                y = fluid.layers.data(name='y', shape=[1], dtype='float32')
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
                cost = fluid.layers.square_error_cost(input=y_predict, label=y)
                avg_cost = fluid.layers.mean(cost)

                ftrl_optimizer = fluid.optimizer.Ftrl(learning_rate=0.1)
                ftrl_optimizer.minimize(avg_cost)

                fetch_list = [avg_cost]
                train_reader = paddle.batch(
                    paddle.dataset.uci_housing.train(), batch_size=1)
                feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
                exe = fluid.Executor(place)
                exe.run(fluid.default_startup_program())
                for data in train_reader():
                    exe.run(main, feed=feeder.feed(data), fetch_list=fetch_list)
C
chengduo 已提交
3518

3519
    NOTE:
C
chengduo 已提交
3520
       Currently, FtrlOptimizer doesn't support sparse parameter optimization.
Q
qiaolongfei 已提交
3521 3522 3523 3524 3525
    """

    _squared_acc_str = "squared"
    _linear_acc_str = "linear"

X
Xin Pan 已提交
3526 3527 3528 3529 3530
    def __init__(self,
                 learning_rate,
                 l1=0.0,
                 l2=0.0,
                 lr_power=-0.5,
3531
                 parameter_list=None,
X
Xin Pan 已提交
3532
                 regularization=None,
3533
                 grad_clip=None,
X
Xin Pan 已提交
3534
                 name=None):
3535 3536 3537 3538 3539
        super(FtrlOptimizer, self).__init__(learning_rate=learning_rate,
                                            parameter_list=parameter_list,
                                            regularization=regularization,
                                            grad_clip=grad_clip,
                                            name=name)
Q
qiaolongfei 已提交
3540 3541 3542 3543 3544 3545 3546 3547 3548 3549 3550 3551 3552 3553 3554 3555 3556 3557 3558 3559 3560 3561 3562 3563
        if learning_rate is None:
            raise ValueError("learning_rate is not set.")

        self.type = "ftrl"
        self._l1 = l1
        self._l2 = l2
        self._lr_power = lr_power

    def _create_accumulators(self, block, parameters):
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")

        for p in parameters:
            self._add_accumulator(self._squared_acc_str, p)
            self._add_accumulator(self._linear_acc_str, p)

    def _append_optimize_op(self, block, param_and_grad):
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")

        squared_acc = self._get_accumulator(self._squared_acc_str,
                                            param_and_grad[0])
        linear_acc = self._get_accumulator(self._linear_acc_str,
                                           param_and_grad[0])
J
Jiabin Yang 已提交
3564
        if framework._non_static_mode():
3565 3566
            _C_ops.ftrl(param_and_grad[0], squared_acc,
                        linear_acc, param_and_grad[1],
3567 3568 3569 3570 3571
                        self._create_param_lr(param_and_grad),
                        param_and_grad[0], squared_acc, linear_acc, "l1",
                        self._l1, "l2", self._l2, "lr_power", self._lr_power)

        else:
3572 3573 3574 3575 3576 3577 3578 3579 3580 3581 3582 3583 3584 3585 3586 3587 3588 3589 3590 3591 3592 3593 3594 3595
            ftrl_op = block.append_op(type=self.type,
                                      inputs={
                                          "Param":
                                          param_and_grad[0],
                                          "Grad":
                                          param_and_grad[1],
                                          "SquaredAccumulator":
                                          squared_acc,
                                          "LinearAccumulator":
                                          linear_acc,
                                          "LearningRate":
                                          self._create_param_lr(param_and_grad),
                                      },
                                      outputs={
                                          "ParamOut": param_and_grad[0],
                                          "SquaredAccumOut": squared_acc,
                                          "LinearAccumOut": linear_acc
                                      },
                                      attrs={
                                          "l1": self._l1,
                                          "l2": self._l2,
                                          "lr_power": self._lr_power
                                      },
                                      stop_gradient=True)
Q
qiaolongfei 已提交
3596

3597
            return ftrl_op
Q
qiaolongfei 已提交
3598 3599


Y
Yibing Liu 已提交
3600
class LambOptimizer(AdamOptimizer):
3601
    r"""
Y
Yibing Liu 已提交
3602 3603 3604 3605
    LAMB (Layer-wise Adaptive Moments optimizer for Batching training) Optimizer.

    LAMB Optimizer is designed to scale up the batch size of training without losing 
    accuracy, which supports adaptive element-wise updating and accurate layer-wise 
Y
Yibing Liu 已提交
3606 3607
    correction. For more information, please refer to `Large Batch Optimization for 
    Deep Learning: Training BERT in 76 minutes <https://arxiv.org/abs/1904.00962>`_ .
Y
Yibing Liu 已提交
3608 3609 3610 3611 3612

    The updating of parameters follows:

    ..  math::

Y
Yibing Liu 已提交
3613
        m_t &= \\beta_1 m_{t - 1}+ (1 - \\beta_1)g_t 
Y
Yibing Liu 已提交
3614

Y
Yibing Liu 已提交
3615
        v_t &= \\beta_2 v_{t - 1}  + (1 - \\beta_2)g_t^2
Y
Yibing Liu 已提交
3616

3617 3618 3619 3620
        m_t &= \\frac{m_t}{\\beta_1^t}

        v_t &= \\frac{v_t}{\\beta_2^t}

Y
Yibing Liu 已提交
3621
        r_t &= \\frac{m_t}{\\sqrt{v_t}+\\epsilon}
Y
Yibing Liu 已提交
3622

Y
Yibing Liu 已提交
3623
        w_t &= w_{t-1} -\\eta_t \\frac{\\left \| w_{t-1}\\right \|}{\\left \| r_t + \\lambda w_{t-1}\\right \|} (r_t + \\lambda w_{t-1})
Y
Yibing Liu 已提交
3624 3625 3626 3627 3628 3629


    where :math:`m` is the 1st moment, and :math:`v` the 2nd moment, :math:`\\eta` the 
    learning rate, :math:`\\lambda` the LAMB weight decay rate.

    Args:
Y
Yibing Liu 已提交
3630 3631 3632 3633 3634 3635 3636 3637
        learning_rate (float|Variable, optional): the learning rate used to update parameters. \
            Can be a float value or a Variable with data type float32. Default 0.001.
        lamb_weight_decay (float, optional): The LAMB weight decay rate. Default 0.01.
        beta1 (float, optional): The exponential decay rate for the 1st moment estimates.
            Default 0.9.
        beta2 (float, optional): The exponential decay rate for the 2nd moment estimates.
            Default 0.999.
        epsilon (float, optional): A small float value for numerical stability. Default 1e-6.
H
hong 已提交
3638
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3639 3640
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3641 3642 3643 3644 3645
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
3646 3647
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of 
            some derived class of ``GradientClipBase`` . There are three cliping strategies 
3648 3649 3650
            ( :ref:`api_paddle_fluid_clip_ClipGradByGlobalNorm` , :ref:`api_paddle_fluid_clip_ClipGradByNorm` ,
            :ref:`api_paddle_fluid_clip_ClipGradByValue` ). If you want better convergence, it is recommended
            to use :ref:`api_paddle_fluid_clip_ClipGradByGlobalNorm` . Default None, meaning there is no gradient clipping.
Y
Yibing Liu 已提交
3651 3652 3653 3654 3655
        exclude_from_weight_decay_fn (function|None): Exclude a parameter from weight 
            decay when **exclude_from_weight_decay_fn(parameter)** returns true. 
            Default None.
        name(str|None): For detailed information, please refer to 
            :ref:`api_guide_Name` . Usually name is no need to set and None by default.
Y
Yibing Liu 已提交
3656 3657 3658 3659 3660 3661

    Examples:
        .. code-block:: python
            
            import paddle.fluid as fluid 

Y
Yibing Liu 已提交
3662
            data = fluid.data(name='x', shape=[-1, 5], dtype='float32')
Y
Yibing Liu 已提交
3663 3664 3665
            hidden = fluid.layers.fc(input=data, size=10)
            cost = fluid.layers.mean(hidden)

Y
Yibing Liu 已提交
3666 3667 3668 3669 3670
            def exclude_fn(param):
                return param.name.endswith('.b_0')

            optimizer = fluid.optimizer.Lamb(learning_rate=0.002,
                                             exclude_from_weight_decay_fn=exclude_fn)
Y
Yibing Liu 已提交
3671 3672 3673 3674 3675 3676 3677 3678 3679 3680 3681 3682 3683
            optimizer.minimize(cost)
    """
    _moment1_acc_str = "moment1"
    _moment2_acc_str = "moment2"
    _beta1_pow_acc_str = "beta1_pow_acc"
    _beta2_pow_acc_str = "beta2_pow_acc"

    def __init__(self,
                 learning_rate=0.001,
                 lamb_weight_decay=0.01,
                 beta1=0.9,
                 beta2=0.999,
                 epsilon=1e-6,
3684
                 parameter_list=None,
Y
Yibing Liu 已提交
3685
                 regularization=None,
3686
                 grad_clip=None,
Y
Yibing Liu 已提交
3687
                 exclude_from_weight_decay_fn=None,
Y
Yibing Liu 已提交
3688 3689 3690 3691 3692 3693
                 name=None):
        assert learning_rate is not None
        assert lamb_weight_decay is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
3694 3695 3696 3697 3698 3699 3700 3701
        super(LambOptimizer, self).__init__(learning_rate=learning_rate,
                                            parameter_list=parameter_list,
                                            regularization=regularization,
                                            grad_clip=grad_clip,
                                            beta1=beta1,
                                            beta2=beta2,
                                            epsilon=epsilon,
                                            name=name)
Y
Yibing Liu 已提交
3702 3703
        self.type = "lamb"
        self._weight_decay = lamb_weight_decay
Y
Yibing Liu 已提交
3704
        self._exclude_from_weight_decay_fn = exclude_from_weight_decay_fn
Y
Yibing Liu 已提交
3705 3706 3707

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)
3708
        block.program._use_lamb = True
Y
Yibing Liu 已提交
3709 3710 3711 3712 3713 3714 3715 3716 3717 3718

        moment1 = self._get_accumulator(self._moment1_acc_str,
                                        param_and_grad[0])
        moment2 = self._get_accumulator(self._moment2_acc_str,
                                        param_and_grad[0])
        beta1_pow_acc = self._get_accumulator(self._beta1_pow_acc_str,
                                              param_and_grad[0])
        beta2_pow_acc = self._get_accumulator(self._beta2_pow_acc_str,
                                              param_and_grad[0])

Y
Yibing Liu 已提交
3719 3720 3721 3722 3723
        if self._exclude_from_weight_decay_fn is not None \
            and self._exclude_from_weight_decay_fn(param_and_grad[0]):
            weight_decay = 0.0
        else:
            weight_decay = self._weight_decay
3724
        lr = self._create_param_lr(param_and_grad)
3725
        master_weight = None
J
Jiabin Yang 已提交
3726
        if framework._non_static_mode():
3727 3728 3729 3730 3731 3732
            _C_ops.lamb(param_and_grad[0], param_and_grad[1], lr, moment1,
                        moment2, beta1_pow_acc, beta2_pow_acc, master_weight,
                        param_and_grad[0], moment1, moment2, beta1_pow_acc,
                        beta2_pow_acc, master_weight, 'beta1', self._beta1,
                        'beta2', self._beta2, 'epsilon', self._epsilon,
                        'weight_decay', weight_decay)
3733
            return None
Y
Yibing Liu 已提交
3734

Y
Yibing Liu 已提交
3735
        # create the lamb optimize op
3736 3737 3738 3739 3740 3741 3742 3743 3744 3745 3746 3747 3748 3749 3750 3751 3752 3753 3754 3755 3756 3757 3758 3759
        lamb_op = block.append_op(type=self.type,
                                  inputs={
                                      "Param": param_and_grad[0],
                                      "Grad": param_and_grad[1],
                                      "LearningRate": lr,
                                      "Moment1": moment1,
                                      "Moment2": moment2,
                                      "Beta1Pow": beta1_pow_acc,
                                      "Beta2Pow": beta2_pow_acc
                                  },
                                  outputs={
                                      "ParamOut": param_and_grad[0],
                                      "Moment1Out": moment1,
                                      "Moment2Out": moment2,
                                      "Beta1PowOut": beta1_pow_acc,
                                      "Beta2PowOut": beta2_pow_acc
                                  },
                                  attrs={
                                      "beta1": self._beta1,
                                      "beta2": self._beta2,
                                      "epsilon": self._epsilon,
                                      "weight_decay": weight_decay
                                  },
                                  stop_gradient=True)
Y
Yibing Liu 已提交
3760 3761 3762 3763

        return lamb_op


3764 3765 3766 3767 3768 3769 3770 3771 3772 3773 3774 3775 3776
# We short the class name, since users will use the optimizer with the package
# name. The sample code:
#
# import paddle.fluid as fluid
#
# sgd = fluid.optimizer.SGD(...)
#
# It is no need to add an `Optimizer` as the class suffix
SGD = SGDOptimizer
Momentum = MomentumOptimizer
Adagrad = AdagradOptimizer
Adam = AdamOptimizer
Adamax = AdamaxOptimizer
3777
Dpsgd = DpsgdOptimizer
3778
DecayedAdagrad = DecayedAdagradOptimizer
3779
Adadelta = AdadeltaOptimizer
Q
qingqing01 已提交
3780
RMSProp = RMSPropOptimizer
Q
qiaolongfei 已提交
3781
Ftrl = FtrlOptimizer
3782
LarsMomentum = LarsMomentumOptimizer
Y
Yibing Liu 已提交
3783
Lamb = LambOptimizer
3784 3785 3786


class ModelAverage(Optimizer):
3787
    r"""
3788
	:api_attr: Static Graph
S
swtkiwi 已提交
3789

3790 3791 3792 3793 3794 3795 3796 3797 3798 3799 3800 3801 3802 3803 3804 3805 3806 3807
    The ModelAverage optimizer accumulates specific continuous historical parameters
    during training. The accumulated historical range can be controlled by the passed
    ``average_window_rate`` argument. The averaged ``Parameter`` are used in the prediction,
    which usually can improve the accuracy of the prediction.

    Accumulate the average of the ``Parameter`` in the sliding window, the result will be saved
    in a temporary variable, can be applied to the current model's ``Parameter`` by calling
    the ``apply()`` method, and the current model ``Parameter`` can be restored by calling
    the ``restore()`` method.

    The window size for calculating the average is determined by ``average_window_rate``,
    ``min_average_window``, ``max_average_window`` and the current ``Parameter`` update times (num_updates).

    When the cumulative times (num_accumulates) is greater than the specific window
    threshold (average_window), the accumulated ``Parameter`` temporary variable is set to 0.0.
    The following example will help to understand the role of these arguments:

    ::
3808

3809 3810 3811 3812 3813 3814 3815 3816 3817
        if num_accumulates >= min_average_window and num_accumulates >= min(max_average_window, num_updates * average_window_rate):
            num_accumulates = 0

    In the above conditional judgment statement, ``num_accumulates`` indicates the current
    accumulated number, which can be abstractly understood as the length of the cumulative window.
    The length of the window must be at least the length set by the ``min_average_window`` argument,
    and cannot exceed the length specified by the ``max_average_window`` argument or
    ``num_updates * average_window_rate``, where ``num_updates`` indicates the current ``Parameter``
    update times, ``average_window_rate`` is a coefficient that calculates the length of the window.
3818 3819

    Args:
3820 3821 3822
        average_window_rate (float): The calculate ratio of the window length relative to ``Parameter`` update times.
        min_average_window (int, optional): the minimum size of average window length. The default value is 10000.
        max_average_window (int, optional): The maximum size of average window length. The default value is 10000.
3823 3824 3825 3826 3827
        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
             :ref:`api_fluid_regularizer_L1Decay` , :ref:`api_fluid_regularizer_L2Decay` . If a parameter has set \
            regularizer using :ref:`api_fluid_ParamAttr` already, the regularization setting here in optimizer will be \
            ignored for this parameter. Otherwise, the regularization setting here in optimizer will take effect.  \
            Default None, meaning there is no regularization.
3828 3829 3830
        name (str, optional): Normally there is no need for user to set this property.
            For more information, please refer to :ref:`api_guide_Name`.
            The default value is None.
3831

3832
    Examples:
Q
qiaolongfei 已提交
3833 3834 3835

      .. code-block:: python

3836 3837 3838 3839 3840 3841
        import paddle.fluid as fluid
        import numpy

        # First create the Executor.
        place = fluid.CPUPlace()  # fluid.CUDAPlace(0)
        exe = fluid.Executor(place)
3842

3843 3844 3845 3846
        train_program = fluid.Program()
        startup_program = fluid.Program()
        with fluid.program_guard(train_program, startup_program):
            # build net
3847
            data = fluid.data(name='X', shape=[None, 1], dtype='float32')
3848 3849 3850 3851 3852 3853 3854 3855
            hidden = fluid.layers.fc(input=data, size=10)
            loss = fluid.layers.mean(hidden)
            optimizer = fluid.optimizer.Momentum(learning_rate=0.2, momentum=0.1)
            optimizer.minimize(loss)

            # build ModelAverage optimizer
            model_average = fluid.optimizer.ModelAverage(0.15,
                                                         min_average_window=10000,
3856
                                                         max_average_window=12500)
3857 3858

            exe.run(startup_program)
3859 3860 3861 3862 3863
            for i in range(12500):
                x = numpy.random.random(size=(10, 1)).astype('float32')
                outs = exe.run(program=train_program,
                               feed={'X': x},
                               fetch_list=[loss.name])
3864 3865

            # apply ModelAverage
3866
            with model_average.apply(exe):
3867 3868 3869 3870
                x = numpy.random.random(size=(10, 1)).astype('float32')
                exe.run(program=train_program,
                        feed={'X': x},
                        fetch_list=[loss.name])
3871 3872 3873
    """

    def __init__(self,
W
wanghaoshuang 已提交
3874
                 average_window_rate,
3875 3876
                 min_average_window=10000,
                 max_average_window=10000,
X
Xin Pan 已提交
3877 3878
                 regularization=None,
                 name=None):
J
Jiabin Yang 已提交
3879
        if framework._non_static_mode():
Z
zhongpu 已提交
3880
            raise Exception("In dygraph, don't support ModelAverage.")
3881 3882 3883
        super(ModelAverage, self).__init__(0.0,
                                           regularization=regularization,
                                           name=name)
3884 3885 3886
        self.average_window = average_window_rate
        self.min_average_window = min_average_window
        self.max_average_window = max_average_window
3887

3888
        self.params_grads = []
3889 3890
        for param in framework.default_main_program().global_block(
        ).all_parameters():
3891
            if param.do_model_average != False:
3892
                grad = param.block.create_var(
3893 3894
                    name=unique_name.generate_with_ignorable_key(".".join(
                        [param.name, 'tmp'])),
3895 3896
                    dtype=param.dtype,
                    persistable=False,
W
wanghaoshuang 已提交
3897
                    stop_gradient=True)
3898
                self.params_grads.append((param, grad))
3899

3900
        for param, grad in self.params_grads:
3901 3902
            if grad is None:
                continue
X
Xin Pan 已提交
3903 3904
            with param.block.program._optimized_guard(
                [param, grad]), name_scope('move_average'):
3905
                self._append_average_accumulate_op(param)
3906

3907 3908 3909 3910
        self.apply_program = Program()
        block = self.apply_program.global_block()
        with program_guard(main_program=self.apply_program):
            for param_grad in self.params_grads:
3911
                self._add_average_apply_op(block, param_grad)
3912 3913 3914 3915 3916

        self.restore_program = Program()
        block = self.restore_program.global_block()
        with program_guard(main_program=self.restore_program):
            for param_grad in self.params_grads:
3917
                self._add_average_restore_op(block, param_grad)
3918

3919
    def _add_average_apply_op(self, block, param_grad):
L
Luo Tao 已提交
3920 3921 3922 3923 3924 3925
        param = block._clone_variable(param_grad[0])
        grad = block._clone_variable(param_grad[1])
        sum_1 = block._clone_variable(self._get_accumulator('sum_1', param))
        sum_2 = block._clone_variable(self._get_accumulator('sum_2', param))
        sum_3 = block._clone_variable(self._get_accumulator('sum_3', param))
        num_accumulates = block._clone_variable(
3926
            self._get_accumulator('num_accumulates', param))
L
Luo Tao 已提交
3927
        old_num_accumulates = block._clone_variable(
3928
            self._get_accumulator('old_num_accumulates', param))
L
Luo Tao 已提交
3929
        num_updates = block._clone_variable(
3930 3931 3932 3933 3934 3935
            self._get_accumulator('num_updates', param))
        # backup param value to grad
        layers.assign(input=param, output=grad)
        # param = (sum_1 + sum_2 + sum_3) / (num_accumulates + old_num_accumulates)
        tmp = layers.sum(x=[num_accumulates, old_num_accumulates])
        sum = layers.sum(x=[sum_1, sum_2, sum_3])
D
dzhwinter 已提交
3936 3937 3938 3939
        tmp = layers.cast(
            x=tmp, dtype='float32' if self._dtype == None else self._dtype)
        sum = layers.cast(
            x=sum, dtype='float32' if self._dtype == None else self._dtype)
S
sneaxiy 已提交
3940
        ops._elementwise_div(x=sum, y=tmp, out=param)
3941 3942

    def _add_average_restore_op(self, block, param_grad):
L
Luo Tao 已提交
3943 3944
        param = block._clone_variable(param_grad[0])
        grad = block._clone_variable(param_grad[1])
3945 3946 3947 3948 3949 3950 3951
        layers.assign(input=grad, output=param)

    def _append_average_accumulate_op(self, param):
        self.helper = LayerHelper("average_accumulate")
        sum_1 = self._add_accumulator('sum_1', param)
        sum_2 = self._add_accumulator('sum_2', param)
        sum_3 = self._add_accumulator('sum_3', param)
3952 3953 3954 3955 3956 3957 3958 3959 3960 3961 3962 3963 3964 3965 3966 3967 3968 3969 3970 3971 3972 3973 3974 3975 3976 3977 3978 3979 3980 3981 3982 3983 3984 3985 3986 3987 3988 3989
        num_accumulates = self._add_accumulator('num_accumulates',
                                                param,
                                                dtype='int64',
                                                shape=[1])
        old_num_accumulates = self._add_accumulator('old_num_accumulates',
                                                    param,
                                                    dtype='int64',
                                                    shape=[1])
        num_updates = self._add_accumulator('num_updates',
                                            param,
                                            dtype='int64',
                                            shape=[1])

        self.helper.append_op(type='average_accumulates',
                              inputs={
                                  "param": param,
                                  "in_sum_1": sum_1,
                                  "in_sum_2": sum_2,
                                  "in_sum_3": sum_3,
                                  "in_num_accumulates": num_accumulates,
                                  "in_old_num_accumulates": old_num_accumulates,
                                  "in_num_updates": num_updates
                              },
                              outputs={
                                  "out_sum_1": sum_1,
                                  "out_sum_2": sum_2,
                                  "out_sum_3": sum_3,
                                  "out_num_accumulates": num_accumulates,
                                  "out_old_num_accumulates":
                                  old_num_accumulates,
                                  "out_num_updates": num_updates,
                              },
                              attrs={
                                  "average_window": self.average_window,
                                  "min_average_window": self.min_average_window,
                                  "max_average_window": self.max_average_window,
                              },
                              stop_gradient=True)
3990

S
rename  
sneaxiy 已提交
3991
    @signature_safe_contextmanager
3992
    def apply(self, executor, need_restore=True):
3993 3994
        """
        Apply the average of the cumulative ``Parameter`` to the parameters of the current model.
3995 3996

        Args:
3997 3998 3999 4000 4001 4002 4003 4004 4005 4006 4007 4008 4009 4010 4011 4012 4013 4014 4015 4016 4017 4018 4019 4020 4021 4022 4023 4024 4025 4026 4027 4028 4029 4030 4031 4032 4033 4034 4035 4036 4037 4038 4039 4040
            executor(fluid.Executor): The current network executor.
            need_restore(bool): Restore flag variable, if set to True, the network will restore
                the parameters of the network to the default value, if set to False,
                it will not be restored. The default value is True.

        Examples:

          .. code-block:: python

            import paddle.fluid as fluid
            import numpy

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

            train_program = fluid.Program()
            startup_program = fluid.Program()
            with fluid.program_guard(train_program, startup_program):
                # build net
                data = fluid.data(name='X', shape=[None, 1], dtype='float32')
                hidden = fluid.layers.fc(input=data, size=10)
                loss = fluid.layers.mean(hidden)
                optimizer = fluid.optimizer.Momentum(learning_rate=0.2, momentum=0.1)
                optimizer.minimize(loss)

                # build ModelAverage optimizer
                model_average = fluid.optimizer.ModelAverage(0.15,
                                                            min_average_window=10000,
                                                            max_average_window=12500)

                exe.run(startup_program)
                for i in range(12500):
                    x = numpy.random.random(size=(10, 1)).astype('float32')
                    outs = exe.run(program=train_program,
                                feed={'X': x},
                                fetch_list=[loss.name])

                # apply ModelAverage
                with model_average.apply(exe):
                    x = numpy.random.random(size=(10, 1)).astype('float32')
                    exe.run(program=train_program,
                            feed={'X': x},
                            fetch_list=[loss.name])
4041
        """
4042 4043 4044 4045 4046 4047
        executor.run(self.apply_program)
        try:
            yield
        finally:
            if need_restore:
                self.restore(executor)
4048 4049

    def restore(self, executor):
4050 4051
        """
        Restore ``Parameter`` values of current model.
4052 4053
        
        Args:
4054 4055 4056 4057 4058 4059 4060 4061 4062 4063 4064 4065 4066 4067 4068 4069 4070 4071 4072 4073 4074 4075 4076 4077 4078 4079 4080 4081 4082 4083 4084 4085 4086 4087 4088 4089 4090 4091 4092 4093 4094 4095 4096 4097
            executor(fluid.Executor): The current network executor.

        Examples:

          .. code-block:: python

            import paddle.fluid as fluid
            import numpy

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

            train_program = fluid.Program()
            startup_program = fluid.Program()
            with fluid.program_guard(train_program, startup_program):
                # build net
                data = fluid.data(name='X', shape=[None, 1], dtype='float32')
                hidden = fluid.layers.fc(input=data, size=10)
                loss = fluid.layers.mean(hidden)
                optimizer = fluid.optimizer.Momentum(learning_rate=0.2, momentum=0.1)
                optimizer.minimize(loss)

                # build ModelAverage optimizer
                model_average = fluid.optimizer.ModelAverage(0.15,
                                                            min_average_window=10000,
                                                            max_average_window=12500)

                exe.run(startup_program)
                for i in range(12500):
                    x = numpy.random.random(size=(10, 1)).astype('float32')
                    outs = exe.run(program=train_program,
                                feed={'X': x},
                                fetch_list=[loss.name])

                # apply ModelAverage
                with model_average.apply(exe, False):
                    x = numpy.random.random(size=(10, 1)).astype('float32')
                    exe.run(program=train_program,
                            feed={'X': x},
                            fetch_list=[loss.name])

                # restore Parameters
                model_average.restore(exe)
4098
        """
4099
        executor.run(self.restore_program)
4100 4101 4102


class ExponentialMovingAverage(object):
4103
    r"""
4104
	:api_attr: Static Graph
S
swtkiwi 已提交
4105

4106 4107 4108 4109 4110 4111
    Compute the moving average of parameters with exponential decay.
    Given a parameter :math:`\\theta`, its exponential moving average (EMA)
    will be

    ..  math::

4112
        \\text{EMA}_0 & = 0
4113

4114 4115
	\\text{EMA}_t & = \\text{decay} * \\text{EMA}_{t-1} + (1 - \\text{decay}) * \\theta_t

Y
Yibing Liu 已提交
4116 4117 4118 4119
    The average results calculated by **update()** method will be saved in 
    temporary variables which are created and maintained by the object, and can 
    be applied to parameters of current model by calling **apply()** method. And 
    the **restore()** method is used to restore the parameters.
4120 4121 4122 4123 4124 4125 4126 4127 4128 4129 4130 4131 4132 4133 4134 4135 4136 4137 4138 4139 4140

    **Bias correction**. All EMAs are initialized to :math:`0` and hence they will be 
    zero biased, which can be corrected by divided by a factor 
    :math:`(1 - \\text{decay}^t)` , i.e., the actual EMAs applied to parameters 
    when calling **apply()** method would be 

    ..  math::
    
        \\widehat{\\text{EMA}}_t = \\frac{\\text{EMA}_t}{1 - \\text{decay}^t}

    **Decay rate scheduling**. A large decay rate very close to 1 would result 
    in that the averages move very slowly. And a better strategy is to set a 
    relative smaller decay rate in the very beginning. The argument **thres_steps**
    allows users to pass a Variable to schedule the decay rate, in this case, 
    the actual decay rate becomes
     
    ..  math::
    
        \\min(\\text{decay}, \\frac{1 + \\text{thres_steps}}{10 + \\text{thres_steps}})

    Usually **thres_steps** can be the global training steps.
4141 4142 4143


    Args:
4144 4145 4146
        decay (float, optional): The exponential decay rate, usually close to 1, such as 0.999, 0.9999, ... . Default 0.999.
        thres_steps (Variable|None, optional): If not `None`, schedule the decay rate. Default None.
        name (str|None, optional): For detailed information, please refer to :ref:`api_guide_Name`. Usually name is no need to set and None by default.
4147 4148 4149 4150


    Examples:

4151 4152 4153 4154 4155 4156 4157 4158 4159 4160 4161 4162 4163 4164 4165 4166 4167 4168 4169 4170 4171 4172 4173 4174 4175 4176 4177 4178 4179 4180 4181 4182 4183 4184 4185 4186 4187 4188 4189 4190 4191 4192 4193 4194 4195 4196
        .. code-block:: python

            import numpy
            import paddle
            import paddle.static as static
            from paddle.static import ExponentialMovingAverage

            paddle.enable_static()

            data = static.data(name='x', shape=[-1, 5], dtype='float32')
            hidden = static.nn.fc(x=data, size=10)
            cost = paddle.mean(hidden)

            test_program = static.default_main_program().clone(for_test=True)
            optimizer = paddle.optimizer.Adam(learning_rate=0.001)
            optimizer.minimize(cost)

            ema = ExponentialMovingAverage(0.999)
            ema.update()

            place = paddle.CPUPlace()
            exe = static.Executor(place)
            exe.run(static.default_startup_program())

            for pass_id in range(3):
                for batch_id in range(6):
                    data = numpy.random.random(size=(10, 5)).astype('float32')
                    exe.run(program=static.default_main_program(),
                    feed={'x': data}, 
                    fetch_list=[cost.name])

                # usage 1
                with ema.apply(exe):
                    data = numpy.random.random(size=(10, 5)).astype('float32')
                    exe.run(program=test_program,
                        feed={'x': data}, 
                        fetch_list=[hidden.name])

                # usage 2
                with ema.apply(exe, need_restore=False):
                    data = numpy.random.random(size=(10, 5)).astype('float32')
                    exe.run(program=test_program,
                        feed={'x': data}, 
                        fetch_list=[hidden.name])
                ema.restore(exe)

4197 4198
    """

4199
    def __init__(self, decay=0.999, thres_steps=None, name=None):
J
Jiabin Yang 已提交
4200
        if framework._non_static_mode():
Z
zhongpu 已提交
4201 4202
            raise Exception(
                "In dygraph, don't support ExponentialMovingAverage.")
4203
        self._decay = decay
4204
        self._thres_steps = thres_steps
4205
        self._name = name if name is not None else ''
4206 4207
        self._decay_var = self._get_ema_decay()

4208
        self._step_counter_name = "@EMA_STEP_COUNTER@"
Y
Yibing Liu 已提交
4209
        self._params_tmps = []
4210
        for param in default_main_program().global_block().all_parameters():
4211
            if param.do_model_average != False:
4212 4213 4214 4215 4216
                tmp = param.block.create_var(name=unique_name.generate(".".join(
                    [self._name + param.name, 'ema_tmp'])),
                                             dtype=param.dtype,
                                             persistable=False,
                                             stop_gradient=True)
Y
Yibing Liu 已提交
4217
                self._params_tmps.append((param, tmp))
4218

Y
Yibing Liu 已提交
4219 4220
        self._ema_vars = {}
        for param, tmp in self._params_tmps:
4221 4222
            with param.block.program._optimized_guard(
                [param, tmp]), name_scope('moving_average'):
Y
Yibing Liu 已提交
4223
                self._ema_vars[param.name] = self._create_ema_vars(param)
4224 4225 4226 4227

        self.apply_program = Program()
        block = self.apply_program.global_block()
        with program_guard(main_program=self.apply_program):
4228
            decay_pow, global_step = self._get_decay_pow(block)
Y
Yibing Liu 已提交
4229
            for param, tmp in self._params_tmps:
4230 4231
                param = block._clone_variable(param)
                tmp = block._clone_variable(tmp)
Y
Yibing Liu 已提交
4232
                ema = block._clone_variable(self._ema_vars[param.name])
4233
                layers.assign(input=param, output=tmp)
4234
                # bias correction
4235 4236
                with layers.control_flow.Switch() as switch:
                    with switch.case(global_step > 0):
4237 4238
                        layers.assign(output=param,
                                      input=ema / (1.0 - decay_pow))
4239 4240
                    with switch.default():
                        layers.assign(output=param, input=ema)
4241 4242 4243 4244

        self.restore_program = Program()
        block = self.restore_program.global_block()
        with program_guard(main_program=self.restore_program):
Y
Yibing Liu 已提交
4245
            for param, tmp in self._params_tmps:
4246 4247 4248 4249
                tmp = block._clone_variable(tmp)
                param = block._clone_variable(param)
                layers.assign(input=tmp, output=param)

4250 4251 4252 4253 4254 4255 4256 4257 4258 4259 4260 4261 4262 4263 4264 4265
    def _get_ema_decay(self):
        with default_main_program()._lr_schedule_guard():
            decay_var = layers.tensor.create_global_var(
                shape=[1],
                value=self._decay,
                dtype='float32',
                persistable=True,
                name="scheduled_ema_decay_rate")

            if self._thres_steps is not None:
                decay_t = (self._thres_steps + 1.0) / (self._thres_steps + 10.0)
                with layers.control_flow.Switch() as switch:
                    with switch.case(decay_t < self._decay):
                        layers.tensor.assign(decay_t, decay_var)
                    with switch.default():
                        layers.tensor.assign(
4266
                            np.array([self._decay], dtype=np.float32),
4267 4268 4269 4270
                            decay_var)
        return decay_var

    def _get_decay_pow(self, block):
4271 4272 4273 4274 4275
        global_step = layers.create_global_var(name=self._step_counter_name,
                                               shape=[1],
                                               value=0,
                                               dtype='int64',
                                               persistable=True)
4276
        global_step = layers.cast(global_step, "float32")
4277
        decay_var = block._clone_variable(self._decay_var)
4278 4279
        decay_pow_acc = layers.elementwise_pow(decay_var, global_step)
        return decay_pow_acc, global_step
4280

Y
Yibing Liu 已提交
4281
    def _create_ema_vars(self, param):
4282 4283 4284 4285 4286 4287 4288 4289 4290
        param_ema = layers.create_global_var(
            name=unique_name.generate(self._name + param.name + '_ema'),
            shape=param.shape,
            value=0.0,
            dtype=param.dtype,
            persistable=True)

        return param_ema

Y
Yibing Liu 已提交
4291 4292 4293 4294 4295
    def update(self):
        """ 
        Update Exponential Moving Average. Should only call this method in 
        train program.
        """
4296 4297
        global_step = layers.autoincreased_step_counter(
            counter_name=self._step_counter_name)
4298
        param_master_emas = []
Y
Yibing Liu 已提交
4299 4300 4301 4302
        for param, tmp in self._params_tmps:
            with param.block.program._optimized_guard(
                [param, tmp]), name_scope('moving_average'):
                param_ema = self._ema_vars[param.name]
4303
                if param.name + '.master' in self._ema_vars:
4304 4305 4306 4307 4308 4309 4310 4311 4312 4313 4314 4315 4316 4317 4318 4319 4320
                    master_ema = self._ema_vars[param.name + '.master']
                    param_master_emas.append([param_ema, master_ema])
                else:
                    ema_t = param_ema * self._decay_var + param * (
                        1 - self._decay_var)
                    layers.assign(input=ema_t, output=param_ema)

        # for fp16 params
        for param_ema, master_ema in param_master_emas:
            default_main_program().global_block().append_op(
                type="cast",
                inputs={"X": master_ema},
                outputs={"Out": param_ema},
                attrs={
                    "in_dtype": master_ema.dtype,
                    "out_dtype": param_ema.dtype
                })
Y
Yibing Liu 已提交
4321

4322 4323 4324 4325 4326 4327 4328
    @signature_safe_contextmanager
    def apply(self, executor, need_restore=True):
        """
        Apply moving average to parameters for evaluation.
        
        Args:
            executor (Executor): The Executor to execute applying.
Y
Yibing Liu 已提交
4329 4330
            need_restore (bool, optional): Whether to restore parameters after 
                applying. Default True.
4331 4332 4333 4334 4335 4336 4337 4338 4339 4340 4341 4342 4343 4344 4345
        """
        executor.run(self.apply_program)
        try:
            yield
        finally:
            if need_restore:
                self.restore(executor)

    def restore(self, executor):
        """Restore parameters.
        
        Args:
            executor (Executor): The Executor to execute restoring.
        """
        executor.run(self.restore_program)
H
hutuxian 已提交
4346 4347 4348


class PipelineOptimizer(object):
4349
    """
4350
	:api_attr: Static Graph
S
swtkiwi 已提交
4351

4352 4353 4354 4355
    Pipeline Optimizer: Make a program to run as pipeline, that is splitting a
    program into multiple sections (sub-programs) and each section run on a
    device to enable the training of large scale models and the use of
    heterogeneous devices. Meanwhile, all sections run in the stype of pipeline.
H
hutuxian 已提交
4356

4357
    Args:
4358 4359 4360 4361
        optimizer (Optimizer): The optimizer to use, such as SGD.
        num_microbatches (int): Number of microbatches. [Optional. Default:1].
        start_cpu_core_id (int): The first cpu core id to use. [Optional. Default:0].
    
4362 4363
    Examples:
        .. code-block:: python
H
hutuxian 已提交
4364

4365
            import paddle.fluid as fluid
H
hutuxian 已提交
4366 4367
            import paddle.fluid.layers as layers

4368 4369 4370 4371 4372 4373 4374 4375 4376 4377 4378 4379 4380 4381 4382 4383
            with fluid.device_guard("gpu:0"):
                x = fluid.layers.data(name='x', shape=[1], dtype='int64', lod_level=0)
                y = fluid.layers.data(name='y', shape=[1], dtype='int64', lod_level=0)
                data_loader = fluid.io.DataLoader.from_generator(
                    feed_list=[x, y],
                    capacity=64,
                    use_double_buffer=True,
                    iterable=False)

                emb_x = layers.embedding(input=x, param_attr=fluid.ParamAttr(name="embx"), size=[10,2], is_sparse=False)
                emb_y = layers.embedding(input=y, param_attr=fluid.ParamAttr(name="emby",learning_rate=0.9), size=[10,2], is_sparse=False)

            with fluid.device_guard("gpu:1"):
                concat = layers.concat([emb_x, emb_y], axis=1)
                fc = layers.fc(input=concat, name="fc", size=1, num_flatten_dims=1, bias_attr=False)
                loss = layers.reduce_mean(fc)
H
hutuxian 已提交
4384
            optimizer = fluid.optimizer.SGD(learning_rate=0.5)
4385
            optimizer = fluid.optimizer.PipelineOptimizer(optimizer)
H
hutuxian 已提交
4386
            optimizer.minimize(loss)
4387 4388 4389 4390 4391 4392 4393 4394 4395

            def train_reader():
                for _ in range(4):
                    x = np.random.random(size=[1]).astype('int64')
                    y = np.random.random(size=[1]).astype('int64')
                    yield x, y
            data_loader.set_sample_generator(train_reader, batch_size=1)

            place = fluid.CUDAPlace(0)
H
hutuxian 已提交
4396 4397
            exe = fluid.Executor(place)
            exe.run(fluid.default_startup_program())
4398 4399
            batch_size = 1
            data_loader.start()
H
hutuxian 已提交
4400
            exe.train_from_dataset(
4401
                    fluid.default_main_program())
4402
            data_loader.reset()
4403 4404
    """

4405
    def __init__(self, optimizer, num_microbatches=1, start_cpu_core_id=0):
4406 4407 4408 4409 4410
        self._device = 'cpu'
        if core.is_compiled_with_npu():
            self._device = "npu"
        elif core.is_compiled_with_cuda():
            self._device = "gpu"
J
Jiabin Yang 已提交
4411
        if framework._non_static_mode():
Z
zhongpu 已提交
4412
            raise Exception("In dygraph, don't support PipelineOptimizer.")
4413 4414 4415 4416
        valid_optimizers = (Optimizer, paddle.optimizer.Optimizer,
                            paddle.fluid.contrib.mixed_precision.decorator.
                            OptimizerWithMixedPrecision)
        if not isinstance(optimizer, valid_optimizers):
4417 4418
            raise ValueError("The 'optimizer' parameter for "
                             "PipelineOptimizer must be an instance of "
4419 4420
                             "{}, but the given type is {}.".format(
                                 valid_optimizers, type(optimizer)))
H
hutuxian 已提交
4421
        self._optimizer = optimizer
4422 4423 4424 4425 4426 4427

        # Get the original optimizer defined by users, such as SGD
        self._origin_optimizer = self._optimizer
        while hasattr(self._origin_optimizer, "inner_opt"):
            self._origin_optimizer = self._origin_optimizer.inner_opt

4428 4429 4430 4431
        assert num_microbatches >= 1, (
            "num_microbatches must be a positive value.")
        self._num_microbatches = num_microbatches
        assert start_cpu_core_id >= 0, (
4432
            "start_cpu_core_id must be a non-negative integer.")
H
hutuxian 已提交
4433
        self._start_cpu_core_id = start_cpu_core_id
4434 4435 4436 4437 4438 4439
        self._place_list = None
        op_maker = core.op_proto_and_checker_maker
        self._op_role = op_maker.OpRole
        self._op_role_key = op_maker.kOpRoleAttrName()
        self._op_role_var_key = op_maker.kOpRoleVarAttrName()
        self._op_device_key = op_maker.kOpDeviceAttrName()
4440
        self._param_device_map = None
4441 4442
        self._pipeline_pair = []
        self._pp_ring_map = dict()
4443 4444
        self.output_var_to_op = None
        self.input_var_to_op = None
4445 4446 4447 4448 4449 4450 4451 4452 4453 4454 4455 4456 4457 4458 4459

    # insert allreduce op to sync global information for global
    # gradient clip and amp
    def _insert_allreduce_op(self, op_idx, block):
        """
        Insert allreduce op to sync global information for global
        gradient clip and amp.
        """
        op = block.ops[op_idx]
        out_name = op.desc.output_arg_names()[0]
        out_var = block.var(out_name)
        offset = 0
        if op.type == "reduce_any":
            # cast the bool var to int32 to use allreduce_max op
            temp_var_name = unique_name.generate(out_name + "_cast_int32")
4460 4461 4462 4463 4464 4465 4466 4467 4468 4469 4470 4471
            temp_var = block.create_var(name=temp_var_name,
                                        shape=[1],
                                        dtype="int32")
            block._insert_op(op_idx + 1 + offset,
                             type='cast',
                             inputs={'X': out_var},
                             outputs={'Out': temp_var},
                             attrs={
                                 'in_dtype': out_var.dtype,
                                 'out_dtype': temp_var.dtype,
                                 self._op_role_key: self._op_role.Optimize
                             })
4472 4473 4474 4475 4476 4477 4478 4479
            offset += 1
        block._insert_op(
            op_idx + 1 + offset,
            type='c_allreduce_max'
            if op.type == "reduce_any" else 'c_allreduce_sum',
            inputs={'X': temp_var if op.type == "reduce_any" else out_var},
            outputs={'Out': temp_var if op.type == "reduce_any" else out_var},
            attrs={
4480
                'ring_id': self.global_ring_id,
4481 4482 4483 4484 4485
                self._op_role_key: self._op_role.Optimize,
                'use_calc_stream': True
            })
        offset += 1
        if op.type == "reduce_any":
4486 4487 4488 4489 4490 4491 4492 4493 4494
            block._insert_op(op_idx + 1 + offset,
                             type='cast',
                             inputs={'X': temp_var},
                             outputs={'Out': out_var},
                             attrs={
                                 'in_dtype': temp_var.dtype,
                                 'out_dtype': out_var.dtype,
                                 self._op_role_key: self._op_role.Optimize
                             })
4495
            offset += 1
4496
        return offset
H
hutuxian 已提交
4497

4498
    def _create_vars(self, block, ori_block):
4499
        # Create vars for block, copied from ori_block
H
hutuxian 已提交
4500
        used_var_set = set()
4501 4502 4503 4504 4505 4506 4507 4508 4509
        added_op_num = 0
        op_idx = 0
        op_size = block.desc.op_size()
        while op_idx < op_size + added_op_num:
            # Whether to insert allreduce_sum or allreduce_max op.
            # For amp and global gradient clip strategies, we should
            # get the global information, so allreduce op is needed.
            should_insert = False
            op = block.ops[op_idx]
4510
            # For op process vars on all devices, remove its input
4511 4512 4513 4514 4515 4516 4517 4518 4519 4520 4521 4522 4523 4524 4525
            # vars not in this block
            reserved_x = []
            if op.type == 'reduce_any' and self._is_optimize_op(op):
                should_insert = True
            elif op.type == 'concat' and self._is_optimize_op(op):
                for input_name in op.desc.input("X"):
                    if block._find_var_recursive(input_name):
                        reserved_x.append(input_name)
                op.desc.set_input('X', reserved_x)
            elif op.type == 'update_loss_scaling':
                for input_name in op.desc.input("X"):
                    if block._find_var_recursive(input_name):
                        reserved_x.append(input_name)
                op.desc.set_input('X', reserved_x)
                op.desc.set_output('Out', reserved_x)
4526 4527 4528 4529 4530 4531 4532 4533 4534 4535
            elif op.type == 'check_finite_and_unscale':
                for input_name in op.desc.input("X"):
                    if block._find_var_recursive(input_name):
                        reserved_x.append(input_name)
                op.desc.set_input('X', reserved_x)
                op.desc.set_output('Out', reserved_x)
                if len(reserved_x) == 0:
                    block._remove_op(op_idx)
                    op_size -= 1
                    continue
4536 4537 4538 4539 4540 4541 4542 4543
            elif op.type == 'sum' and self._is_gradient_clip_op(op):
                for input_name in op.desc.input("X"):
                    if block._find_var_recursive(input_name):
                        reserved_x.append(input_name)
                op.desc.set_input('X', reserved_x)
                should_insert = True

            vars = op.desc.input_arg_names() + op.desc.output_arg_names()
H
hutuxian 已提交
4544
            for var in vars:
4545 4546
                # a var whose name contains "blocking_queue"
                # only exists in startup program
4547
                if var in used_var_set or "_blocking_queue" in var:
H
hutuxian 已提交
4548 4549
                    continue
                used_var_set.add(var)
4550 4551
                if block._find_var_recursive(str(var)): continue
                source_var = ori_block._var_recursive(str(var))
4552
                if source_var.type == core.VarDesc.VarType.READER:
4553
                    dest_var = block.create_var(
4554 4555 4556
                        name=var,
                        type=core.VarDesc.VarType.READER,
                        persistable=source_var.persistable)
4557 4558 4559 4560 4561 4562 4563 4564 4565 4566 4567 4568
                elif isinstance(source_var, Parameter):
                    dest_var = block.create_parameter(
                        name=source_var.name,
                        shape=source_var.shape,
                        dtype=source_var.dtype,
                        type=source_var.type,
                        lod_level=source_var.lod_level,
                        stop_gradient=source_var.stop_gradient,
                        trainable=source_var.trainable,
                        optimize_attr=source_var.optimize_attr,
                        regularizer=source_var.regularizer,
                        error_clip=source_var.error_clip)
4569
                else:
4570
                    dest_var = block._clone_variable(source_var, False)
4571
                self._clone_var_attr(dest_var, source_var)
4572 4573 4574 4575 4576 4577 4578 4579
            # When use with sharding, allreduce_sum and allreduce_max
            # used for global gradient clip and amp will be added by sharding.
            op_idx += 1
            if self.use_sharding or not should_insert: continue
            inserted_ops = self._insert_allreduce_op(op_idx - 1, block)
            added_op_num += inserted_ops
            op_idx += inserted_ops
        block._sync_with_cpp()
H
hutuxian 已提交
4580

4581
    def _is_loss_grad_op(self, op):
4582 4583
        assert self._op_role_key in op.attr_names
        op_role = int(op.attr(self._op_role_key))
4584 4585 4586
        return op_role & int(self._op_role.Backward) and op_role & int(
            self._op_role.Loss)

4587
    def _is_forward_op(self, op):
4588 4589
        return self._op_role_key in op.attr_names and (int(
            op.attr(self._op_role_key)) == int(self._op_role.Forward))
4590

4591
    def _is_backward_op(self, op):
4592 4593 4594 4595 4596 4597
        return self._op_role_key in op.attr_names and (
            int(op.attr(self._op_role_key)) & int(self._op_role.Backward))

    def _is_loss_op(self, op):
        assert self._op_role_key in op.attr_names
        return int(op.attr(self._op_role_key)) == int(self._op_role.Loss)
4598 4599

    def _is_optimize_op(self, op):
4600 4601
        return self._op_role_key in op.attr_names and (
            int(op.attr(self._op_role_key)) & int(self._op_role.Optimize))
4602 4603 4604 4605 4606

    def _is_update_op(self, op):
        return 'Param' in op.input_names and 'Grad' in op.input_names and (
            "LearningRate" in op.input_names)

4607
    def _split_program(self, main_program, devices):
H
hutuxian 已提交
4608
        """
4609
        Split a program into sections according to devices that ops run on.
4610
        The op whose op_device attr is "gpu:all" is copied to all sections.
4611 4612 4613

        Args:
            main_program (Program): the main program
4614
            devices: all used devices
H
hutuxian 已提交
4615
        """
4616
        # Map from device to its corresponding section program info
4617
        device_program_map = defaultdict(Program)
4618

4619
        block = main_program.block(0)
4620 4621
        for op in block.ops:
            device = op.attr(self._op_device_key)
4622
            # Copy ops whose op_device set to "gpu:all" to all sections.
4623
            if device == f"{self._device}:all":
4624
                for device in devices:
4625 4626
                    program = device_program_map[device]
                    op_desc = op.desc
4627
                    ap_op = program.global_block().desc.append_op()
4628
                    ap_op.copy_from(op_desc)
4629
                    ap_op._set_attr(self._op_device_key, "")
4630 4631 4632
            else:
                program = device_program_map[device]
                op_desc = op.desc
4633
                ap_op = program.global_block().desc.append_op()
4634
                ap_op.copy_from(op_desc)
4635
                ap_op._set_attr(self._op_device_key, "")
4636

4637
        program_list = []
4638
        for key in devices:
4639
            program = device_program_map[key]
4640 4641
            program._sync_with_cpp()
            program_list.append(program)
H
hutuxian 已提交
4642

4643
        return program_list
H
hutuxian 已提交
4644

4645 4646 4647 4648 4649 4650 4651
    def _get_op_device_for_startup_program(self, var_name):
        """
        For adam optimizer, it will add accumulators and initialize them
        with fill_constant, and force the op device to cpu. Hence, we should
        get the real op_device attribute of the fill_constant as the device
        where the corresponding parameters on.
        """
4652 4653 4654
        assert "beta1_pow_acc" in var_name or "beta2_pow_acc" in var_name, \
            'For accumulators for Adam, the name must contain beta1_pow_acc ' \
            'or beta2_pow_acc.'
4655 4656 4657 4658
        param_name = var_name[0:var_name.index('_beta')]
        device = self._param_device_map[param_name]
        return device

4659 4660
    def _split_startup_program(self, startup_program, device_id):
        block = startup_program.global_block()
4661 4662 4663
        new_startup_program = Program()
        for op in block.ops:
            device = op.attr(self._op_device_key)
4664 4665
            if device == "cpu":
                assert op.type == "fill_constant", (
4666 4667
                    "For ops in startup program with the op_device attribute "
                    "of cpu, they must be of type fill_constant.")
4668 4669 4670
                output_var = op.output_arg_names[0]
                device = self._get_op_device_for_startup_program(output_var)

4671
            if device:
4672
                device_index = int(device.split(':')[1])
4673
            else:
4674 4675
                # LR related ops
                device = None
4676
            if device and device_index != device_id: continue
4677
            op_desc = op.desc
4678
            ap_op = new_startup_program.global_block().desc.append_op()
4679 4680 4681
            ap_op.copy_from(op_desc)
            ap_op._set_attr(self._op_device_key, "")
        new_startup_program._sync_with_cpp()
4682
        self._create_vars(new_startup_program.global_block(), block)
4683 4684
        return new_startup_program

4685
    def _find_post_op(self, index, var_name):
H
hutuxian 已提交
4686
        """
4687
        Find the post op that has variable named var_name as input.
H
hutuxian 已提交
4688
        """
4689 4690 4691 4692 4693 4694
        # bugfix for uniform hybrid parallelism
        if '.cast_fp32' in var_name:
            var_name = var_name.replace('.cast_fp32', '')
        if '.cast_fp16' in var_name:
            var_name = var_name.replace('.cast_fp16', '')

4695 4696 4697 4698 4699 4700 4701 4702
        post_ops = self.input_var_to_op[var_name]
        if post_ops == None: return None
        result_op = None
        for post_op, post_idx in reversed(post_ops):
            if post_idx > index:
                result_op = post_op
                break
        return result_op
4703

4704
    def _find_prev_op(self, index, var_name):
H
hutuxian 已提交
4705
        """
4706 4707
        Find the previous op of op with index that outputs
        variable named var_name.
H
hutuxian 已提交
4708
        """
4709 4710 4711 4712 4713 4714
        prev_ops = self.output_var_to_op[var_name]
        if prev_ops == None: return None
        result_op = None
        for prev_op, prev_idx in reversed(prev_ops):
            if prev_idx < index:
                result_op = prev_op
4715
                break
4716
        return result_op
4717 4718

    def _rename_arg(self, op, old_name, new_name):
4719 4720
        op._rename_input(old_name, new_name)
        op._rename_output(old_name, new_name)
4721

4722
    def _create_var(self, block, ref_var, name, dtype=None):
4723 4724 4725 4726 4727 4728 4729 4730
        """
        Create a new var for block, which has the same type,
        shape and dtype as ref_var, then rename it with the
        name `name`.
        """
        new_var = block.create_var(
            name=name,
            shape=ref_var.shape,
4731
            dtype=ref_var.dtype if dtype is None else dtype,
4732 4733
            type=ref_var.type,
            lod_level=ref_var.lod_level,
4734 4735
            persistable=ref_var.persistable,
            is_data=ref_var.is_data,
4736
            need_check_feed=ref_var.desc.need_check_feed())
4737
        self._clone_var_attr(new_var, ref_var)
4738 4739
        return new_var

4740 4741 4742 4743 4744
    def _clone_var_attr(self, dest, src):
        dest.stop_gradient = src.stop_gradient
        if hasattr(src, 'is_distributed'):
            dest.is_distributed = src.is_distributed

4745 4746 4747 4748 4749 4750
    def _strip_grad_suffix(self, name):
        """
        Strip the grad suffix from the given variable name
        """
        pos = name.find(core.grad_var_suffix())
        return name[:pos] if pos != -1 else name
H
hutuxian 已提交
4751

4752 4753 4754 4755 4756 4757
    def _append_grad_suffix(self, name):
        """
        Append grad suffix to the given variable name
        """
        return name + core.grad_var_suffix()

4758
    def _get_op_device_attr(self, op):
H
hutuxian 已提交
4759
        """
4760
        Get the op_device attribute of a op.
H
hutuxian 已提交
4761
        """
4762 4763 4764
        device = op.attr(self._op_device_key) \
            if op.has_attr(self._op_device_key) else None
        if device:
B
Baibaifan 已提交
4765
            assert device[0:3] == 'gpu' or device[0:3] == 'npu', "Now, only gpu and npu devices are " \
4766 4767 4768 4769 4770 4771 4772 4773 4774 4775 4776 4777 4778 4779
                "supported in pipeline parallemism."
        return device

    def _add_op_device_attr_for_op(self, op, idx, block):
        """
        Add op_device attrribute for ops that have not that attribute set.
        We use "gpu:all" to represent the op should be put on all
        sub-programs, such as lr-related ops. Note that: "gpu:all"
        is only used by pipeline as an indicator.
        """
        lrsched_role = int(self._op_role.LRSched)
        if op.attr(self._op_role_key) == lrsched_role:
            # For LRSched ops, we should put them on all sub-programs to
            # make sure each sub-program update the lr correctly
4780
            op._set_attr(self._op_device_key, f"{self._device}:all")
4781 4782 4783 4784 4785 4786 4787 4788 4789 4790 4791 4792 4793 4794 4795
        # bugfix in hybrid parallelism
        elif op.type == "sum" and self._is_backward_op(op):
            # For sum ops that compute the sum of @RENAMED@ vars
            for name in op.desc.input_arg_names():
                assert '@RENAME@' in name, \
                    "The op must be sum used to accumulate renamed vars."
            assert len(op.desc.output_arg_names()) == 1
            out_name = op.desc.output_arg_names()[0]
            post_op = self._find_post_op(idx, out_name)
            assert post_op.has_attr(
                'op_device'), "{} has no op_device attr for var {}".format(
                    post_op.type, out_name)
            device = post_op.attr(self._op_device_key)
            assert device, "The post op must have op_device set."
            op._set_attr(self._op_device_key, device)
4796 4797
        elif (op.type == "cast"
              or op.type == "scale") and self._is_backward_op(op):
4798
            prev_op = self._find_prev_op(idx, op.desc.input("X")[0])
4799 4800
            op._set_attr(self._op_device_key, prev_op.attr(self._op_device_key))
        elif op.type == "memcpy" and not self._is_optimize_op(op):
4801
            # for checkpoint offloading
4802 4803 4804 4805 4806
            assert len(op.input_arg_names) == 1 and len(
                op.output_arg_names) == 1
            input_name = op.input_arg_names[0]
            output_name = op.output_arg_names[0]
            if '@Fetch' in output_name:
4807
                post_op = self._find_post_op(idx, output_name)
4808 4809 4810
                op._set_attr(self._op_device_key,
                             post_op.attr(self._op_device_key))
            else:
4811
                prev_op = self._find_prev_op(idx, op.desc.input("X")[0])
4812 4813 4814 4815 4816
                op._set_attr(self._op_device_key,
                             prev_op.attr(self._op_device_key))
        elif self._is_loss_op(op):
            # For loss * loss_scaling op added by AMP
            offset = 1
4817 4818
            while (not block.ops[idx + offset].has_attr(self._op_device_key)
                   or not block.ops[idx + offset].attr(self._op_device_key)):
4819 4820 4821 4822 4823 4824 4825 4826 4827
                offset += 1
            device = block.ops[idx + offset].attr(self._op_device_key)
            assert device, "Please put you program within device_guard scope."
            for i in range(offset):
                block.ops[idx + i]._set_attr(self._op_device_key, device)
        elif self._is_optimize_op(op) and op.type == "cast":
            # For fp16-->fp32 cast added by AMP
            grad_name = op.output('Out')
            assert len(grad_name) == 1
4828
            param_name = self._strip_grad_suffix(grad_name[0])
4829 4830 4831 4832 4833 4834 4835 4836 4837 4838 4839 4840
            device = self._param_device_map[param_name]
            op._set_attr(self._op_device_key, device)
        elif self._is_gradient_clip_op(op) or self._is_regularization_op(op):
            # For gradient clip and regularization ops, we set their op_device
            # attribute to the device where their corresponding parameters on.
            assert self._op_role_var_key in op.attr_names, "gradient_clip " \
                "and regularization ops must have op_role_var attribute."
            op_role_var = op.attr(self._op_role_var_key)
            assert len(op_role_var) == 2, "op_role_var for gradient_clip " \
                "regularization ops must have two elements."
            param_name = op_role_var[0]
            device = self._param_device_map[param_name]
4841
            # For sum op added by global gradient clip, it must be
4842
            # put on all devices
4843 4844 4845 4846
            if (op.type == 'sum' or op.type == 'sqrt'
                    or op.type == 'fill_constant'
                    or op.type == 'elementwise_max'
                    or op.type == 'elementwise_div'):
4847
                device = f"{self._device}:all"
4848
            op._set_attr(self._op_device_key, device)
R
Roc 已提交
4849
        elif op.type == "alloc_float_status" or op.type == "clear_float_status":
4850
            op._set_attr(self._op_device_key, f"{self._device}:all")
4851 4852 4853 4854 4855 4856 4857 4858 4859 4860
            # NOTE(wangxi): NPU should only clear the float status
            # once at each batch step
            op._set_attr(self._op_role_key, self._op_role.LRSched)

            float_status_name = op.output_arg_names[0]
            float_status_var = block.var(float_status_name)
            # FIXME(wangxi): pipeline lr schedule will exec on sub_scope(0)
            # while update will exec on sub_scope(last_micro_step), should
            # set persistable to use global scope
            float_status_var.persistable = True
4861 4862
        else:
            other_known_ops = [
R
Roc 已提交
4863
                'update_loss_scaling', 'reduce_any', 'concat', 'sum',
4864
                'check_finite_and_unscale', 'memcpy'
4865 4866 4867 4868 4869
            ]
            assert op.type in other_known_ops, "For other ops without " \
                "op_device set, they must be one of {}, but it " \
                "is {}".format(other_known_ops, op.type)
            assert self._is_optimize_op(op)
4870
            op._set_attr(self._op_device_key, f"{self._device}:all")
4871 4872

    def _add_op_device_attr(self, block):
4873
        """
4874 4875
        Add op_device attrribute for ops in block that have 
        not that attribute set.
4876
        """
4877
        for idx, op in enumerate(list(block.ops)):
4878 4879 4880
            if (op.type == "create_py_reader" or op.type == "read"
                    or op.type == "create_double_buffer_reader"):
                # Copy read related ops to all section to make them exit
4881 4882 4883 4884
                # after each epoch.
                # We use "gpu:all" to represent the op should be put on all
                # sub-programs, such as lr-related ops. Note that: "gpu:all"
                # is only used by pipeline as an indicator.
4885
                op._set_attr(self._op_device_key, f"{self._device}:all")
4886 4887 4888 4889
                continue
            # op_device attribute has been set
            if self._get_op_device_attr(op): continue
            self._add_op_device_attr_for_op(op, idx, block)
H
hutuxian 已提交
4890

4891 4892
    def _check_validation(self, block):
        """
4893 4894 4895
        Check whether ops in a block have both the op_device and the 
        op_role attributes set.
        Then, return all devices in order.
4896
        """
4897 4898 4899 4900 4901 4902 4903 4904 4905 4906
        device_list = []
        # Section worker only supports the following op_role
        valid_op_role_value = [
            int(self._op_role.LRSched),
            int(self._op_role.Forward),
            int(self._op_role.Backward),
            int(self._op_role.Loss),
            int(self._op_role.Optimize),
            int(self._op_role.Backward) | int(self._op_role.Loss),
        ]
4907
        for op in block.ops:
4908
            if not op._has_kernel(op.type):
4909 4910
                assert op.type == "conditional_block" and (op.attr(
                    self._op_role_key) == int(self._op_role.LRSched)), (
4911 4912
                        "Now, the only supported op without kernel is "
                        "conditional_block, and its op role must be LRSched.")
4913 4914 4915
            assert op.has_attr(
                self._op_role_key), ("op ({}) has no {} attribute.".format(
                    op.type, self._op_role_key))
4916 4917
            op_role = op.attr(self._op_role_key)
            assert int(op_role) in valid_op_role_value, \
4918
                "op_role {} for op {} must be one of {}".format(
4919
                    op_role,
4920 4921
                    op.type,
                    valid_op_role_value)
4922

4923 4924 4925
            assert op.has_attr(
                self._op_device_key), ("op ({}) has no {} attribute.".format(
                    op.type, self._op_device_key))
4926 4927 4928 4929

            device = op.attr(self._op_device_key)
            assert device, ("op_device attribute for op "
                            "{} has not been set.".format(op.type))
4930
            if device == f"{self._device}:all": continue
4931

4932
            dev_type = device.split(':')[0]
B
Baibaifan 已提交
4933 4934 4935
            assert dev_type == "gpu" or dev_type == 'npu', (
                "Now only gpu and npu devices are supported "
                "for pipeline parallelism.")
4936 4937

            if device not in device_list:
4938
                device_list.append(device)
4939

4940
        return device_list
4941

4942
    def _insert_sendrecv_ops_for_boundaries(self, block):
4943
        """
4944
        Insert a pair of send and recv ops for every two
4945 4946
        consecutive ops on different devices.
        """
4947
        # A map from var to device where op takes it as input,
4948
        # avoiding multiple send and recv ops.
4949
        input_var_to_device = dict()
4950 4951 4952 4953 4954 4955 4956 4957 4958 4959
        # bugfix hybrid parallelism
        first_optimize_index = None
        for index, op in enumerate(list(block.ops)):
            if self._is_optimize_op(op):
                first_optimize_index = index
                break
        extra_index_info = {
            'index': 0,
            'first_optimize_index': first_optimize_index
        }
4960

4961
        for index, op in enumerate(list(block.ops)):
4962
            cur_device = op.attr(self._op_device_key)
4963
            if cur_device == f"{self._device}:all": continue
4964 4965
            for var_name in op.input_arg_names:
                var = block.var(var_name)
4966
                # skip data var
4967
                if var.is_data: continue
4968
                prev_device = None
4969 4970 4971

                prev_op = self._find_prev_op(index, var_name)
                if prev_op is None:
4972 4973
                    if var_name not in self._param_device_map:
                        continue
4974
                    prev_device = self._param_device_map[var_name]
4975

4976 4977 4978
                if not prev_device:
                    prev_device = prev_op.attr(self._op_device_key) \
                        if prev_op else None
4979

4980 4981
                if prev_device is None or prev_device == f"{self._device}:all":
                    continue
4982 4983

                if prev_device == cur_device: continue
4984

4985 4986 4987 4988 4989 4990 4991
                if var_name not in input_var_to_device:
                    input_var_to_device[var_name] = []
                if (cur_device, prev_device) in input_var_to_device[var_name]:
                    continue

                device_type = cur_device.split(':')[0] + ':'

4992 4993 4994 4995 4996 4997 4998 4999 5000 5001 5002 5003 5004 5005 5006 5007 5008 5009 5010
                def _check_stage(cur_id, prev_id):
                    # check send/recv stage valid
                    is_forward = self._is_forward_op(op)
                    is_backward = self._is_backward_op(op)
                    assert is_forward or is_backward, \
                        'send/recv in pipeline should only be inserted in forward or backward,' \
                        'please check the op_role of op={}'.format(op)

                    if is_forward:
                        assert prev_id < cur_id, \
                            "In forward, send/recv can only be passed forward, but now " \
                            "prev_stage={} great than cur_stage={}, please check op_device of op={}".format(
                                prev_id, cur_id, op)
                    elif is_backward:
                        assert prev_id > cur_id, \
                            "In backward, send/recv can only be passed backward, but now " \
                            "prev_stage={} less than cur_stage={}, please check op_device of op={}".format(
                                prev_id, cur_id, op)

5011 5012 5013 5014 5015 5016 5017 5018 5019 5020 5021 5022 5023 5024 5025 5026 5027 5028 5029 5030 5031 5032 5033
                def _insert_send_recv(cur_id, prev_id):
                    cur_dev = device_type + str(cur_id)
                    prev_dev = device_type + str(prev_id)
                    if (cur_dev, prev_dev) in input_var_to_device[var_name]:
                        return

                    if cur_id - prev_id > 1:
                        _insert_send_recv(cur_id - 1, prev_id)
                        _insert_send_recv(cur_id, cur_id - 1)
                        input_var_to_device[var_name].append(
                            (cur_dev, prev_dev))
                        return
                    elif cur_id - prev_id < -1:
                        _insert_send_recv(cur_id + 1, prev_id)
                        _insert_send_recv(cur_id, cur_id + 1)
                        input_var_to_device[var_name].append(
                            (cur_dev, prev_dev))
                        return

                    assert abs(cur_id - prev_id) == 1
                    input_var_to_device[var_name].append((cur_dev, prev_dev))

                    op_role = op.attr(self._op_role_key)
5034
                    var = block.vars[var_name]
5035 5036 5037
                    pair = (prev_id, cur_id)
                    # 1000 is just a magic number
                    pair_key = prev_id * 1000 + cur_id
5038 5039 5040 5041 5042 5043 5044
                    if pair not in self._pipeline_pair:
                        self._pipeline_pair.append(pair)
                        self._pp_ring_map[pair_key] = self.ring_id
                        ring_id = self.ring_id
                        self.ring_id += 1
                    else:
                        ring_id = self._pp_ring_map[pair_key]
5045

5046
                    if self.schedule_mode == 'F-then-B':  # F-then-B
F
fangshuixun007 已提交
5047
                        block._insert_op_without_sync(
5048
                            index=index + extra_index_info['index'],
5049 5050 5051
                            type='send_v2',
                            inputs={'X': var},
                            attrs={
5052
                                self._op_device_key: prev_dev,
5053 5054 5055 5056 5057
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 1,
                                'ring_id': ring_id
                            })
5058
                        extra_index_info['index'] += 1
5059 5060 5061
                        var_shape = list(var.shape)
                        var_shape[0] = self.micro_batch_size if var_shape[
                            0] < 0 else var_shape[0]
F
fangshuixun007 已提交
5062
                        block._insert_op_without_sync(
5063
                            index=index + extra_index_info['index'],
5064 5065 5066
                            type='recv_v2',
                            outputs={'Out': [var]},
                            attrs={
5067
                                'out_shape': var_shape,
5068
                                'dtype': var.dtype,
5069
                                self._op_device_key: cur_dev,
5070 5071 5072 5073 5074
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 0,
                                'ring_id': ring_id
                            })
5075
                        extra_index_info['index'] += 1
5076
                    elif self.schedule_mode == '1F1B':  # 1F1B
5077 5078 5079 5080
                        var_shape = list(var.shape)
                        var_shape[0] = self.micro_batch_size if var_shape[
                            0] < 0 else var_shape[0]

5081
                        numel = np.prod(var_shape)
5082 5083
                        use_mp = (self.mp_degree > 1) and (numel %
                                                           self.mp_degree == 0)
5084 5085 5086 5087 5088 5089 5090 5091 5092 5093 5094 5095 5096 5097 5098 5099 5100 5101 5102 5103 5104 5105 5106 5107 5108 5109

                        if 'subprog' in var.name:
                            # For recompute, if the checkpoints var is layer_norm_6.tmp_2
                            # this var will be sent twice, layer_norm_6.tmp_2 for forward pass,
                            # layer_norm_6.tmp_2.subprog_* for recompute pass.
                            # We can store the first sent var and copy the value to the
                            # second one to reduce one send/recv op.
                            # The origin_ckpt_name is layer_norm_6.tmp_2, which will be used
                            # to find the stored var for the forward pass.
                            origin_name = var.name.split('subprog')[0][0:-1]
                            associate_var = block.var(origin_name)
                            block._insert_op_without_sync(
                                index=index + extra_index_info['index'],
                                type='assign',
                                inputs={'X': [associate_var]},
                                outputs={'Out': [var]},
                                attrs={
                                    'out_shape': var_shape,
                                    'dtype': var.dtype,
                                    self._op_device_key: cur_dev,
                                    self._op_role_key: op_role,
                                    'use_calc_stream': True,
                                })
                            extra_index_info['index'] += 1
                            return

5110 5111
                        _check_stage(cur_id, prev_id)

5112 5113 5114 5115 5116 5117 5118 5119 5120 5121 5122
                        block._insert_op_without_sync(index=index +
                                                      extra_index_info['index'],
                                                      type='c_sync_calc_stream',
                                                      inputs={'X': [var]},
                                                      outputs={'Out': [var]},
                                                      attrs={
                                                          self._op_device_key:
                                                          prev_dev,
                                                          self._op_role_key:
                                                          op_role,
                                                      })
5123
                        extra_index_info['index'] += 1
5124 5125 5126 5127
                        prefix_name = var.name.split('@')[0]
                        prefix_var = block.var(prefix_name)
                        is_param = True if isinstance(prefix_var,
                                                      Parameter) else False
F
fangshuixun007 已提交
5128
                        block._insert_op_without_sync(
5129
                            index=index + extra_index_info['index'],
5130 5131
                            type='send_v2'
                            if not use_mp or is_param else 'partial_send',
5132 5133
                            inputs={'X': var},
                            attrs={
5134
                                self._op_device_key: prev_dev,
5135 5136 5137 5138
                                self._op_role_key: op_role,
                                'use_calc_stream': False,
                                'ring_id': ring_id,
                                'peer': 1,
5139 5140 5141
                                # if send_v2, num&id attr is not in op_attrs, will not insert
                                'num': self.mp_degree,
                                'id': self.mp_rank,
5142
                            })
5143
                        extra_index_info['index'] += 1
5144 5145 5146 5147 5148 5149 5150 5151
                        insert_index = None
                        if int(op_role) == int(self._op_role.Backward):
                            insert_index = extra_index_info[
                                'first_optimize_index']
                            new_op_role = self._op_role.Optimize
                        else:
                            insert_index = index
                            new_op_role = self._op_role.Backward
5152
                        sync_comm_op = block._insert_op_without_sync(
5153
                            index=insert_index + extra_index_info['index'],
5154 5155 5156 5157
                            type='c_sync_comm_stream',
                            inputs={'X': [var]},
                            outputs={'Out': [var]},
                            attrs={
5158
                                self._op_device_key: prev_dev,
5159
                                self._op_role_key: new_op_role,
5160 5161
                                'ring_id': ring_id,
                            })
5162
                        if int(op_role) == int(self._op_role.Forward):
5163
                            sync_comm_op._set_attr('pipeline_flag', '')
5164
                            extra_index_info['index'] += 1
F
fangshuixun007 已提交
5165
                        block._insert_op_without_sync(
5166
                            index=index + extra_index_info['index'],
5167 5168
                            type='recv_v2'
                            if not use_mp or is_param else 'partial_recv',
5169 5170 5171 5172
                            outputs={'Out': [var]},
                            attrs={
                                'out_shape': var_shape,
                                'dtype': var.dtype,
5173
                                self._op_device_key: cur_dev,
5174 5175 5176
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 0,
5177 5178 5179 5180
                                'ring_id': ring_id,
                                # if recv_v2, num&id attr is not in op_attrs, will not insert
                                'num': self.mp_degree,
                                'id': self.mp_rank,
5181
                            })
5182
                        extra_index_info['index'] += 1
5183
                        if use_mp and not is_param:
5184 5185 5186 5187 5188 5189 5190 5191 5192 5193 5194 5195 5196 5197 5198
                            block._insert_op_without_sync(
                                index=index + extra_index_info['index'],
                                type='partial_allgather',
                                inputs={'X': [var]},
                                outputs={'Out': [var]},
                                attrs={
                                    self._op_device_key: cur_dev,
                                    self._op_role_key: op_role,
                                    'use_calc_stream': True,
                                    'ring_id': 0,
                                    # if recv_v2, num&id attr is not in op_attrs, will not insert
                                    'nranks': self.mp_degree,
                                    'rank': self.mp_rank,
                                })
                            extra_index_info['index'] += 1
5199 5200 5201 5202 5203
                    else:
                        raise ValueError(
                            "Now only 'F-then-B' and '1F1B' are supported."
                            "The given value is {}.".format(self.schedule_mode))

5204 5205
                _insert_send_recv(int(cur_device.split(':')[1]),
                                  int(prev_device.split(':')[1]))
5206 5207
        block._sync_with_cpp()

5208
    def _insert_loss_scale(self, block):
5209
        """
5210
        Scale the loss corresponding to number of micro-batches.
5211
        """
5212
        if self._num_microbatches == 1: return
5213
        for index, op in reversed(tuple(enumerate(list(block.ops)))):
5214
            if self._is_loss_grad_op(op):
5215 5216 5217 5218 5219 5220 5221
                assert op.type == 'fill_constant', \
                    "loss_grad_op must be fill_constant op, " \
                    "but this op is {}".format(op.type)
                assert op.has_attr('value')
                loss_scale = float(op.attr('value'))
                loss_scale = loss_scale / self._num_microbatches
                op._set_attr('value', loss_scale)
5222 5223
                break

5224 5225 5226 5227 5228 5229
    def _rename_gradient_var_name(self, block):
        for index, op in enumerate(block.ops):
            if not self._is_optimize_op(op): continue
            input_names = op.input_arg_names
            output_names = op.output_arg_names
            in_out_names = input_names + output_names
L
lilong12 已提交
5230
            if op.type == 'cast' or op.type == "c_sync_comm_stream": continue
5231 5232 5233 5234 5235 5236 5237 5238
            # append "MERGED" to the names of parameter gradients,
            # and mofify the op_role_var attribute (by rename_arg func).
            for name in in_out_names:
                if not core.grad_var_suffix() in name: continue
                param_name = name.strip(core.grad_var_suffix())
                new_grad_name = name + "@MERGED"
                self._rename_arg(op, name, new_grad_name)

5239 5240 5241
    def _accumulate_gradients(self,
                              block,
                              pp_allreduce_in_optimize=False,
5242 5243
                              strategy=None,
                              shard=None):
5244 5245 5246 5247
        """
        Create a new merged gradient for each parameter and accumulate the
        corresponding gradient to it.
        """
5248 5249
        fp16_allreduce = strategy.fp16_allreduce if strategy else False
        if strategy and strategy.fuse_grad_merge:
5250
            fused_gradient_names = self._accumulate_gradients_with_fuse(
5251
                block, fp16_allreduce, strategy.fuse_grad_size_in_MB, shard)
5252 5253
            return fused_gradient_names

5254 5255 5256
        merged_gradient_names = []
        first_opt_op_idx = None

5257 5258 5259
        merged_suffix = '@MERGED@FP16' if fp16_allreduce else '@MERGED'
        dtype = paddle.float16 if fp16_allreduce else None

5260 5261 5262 5263 5264 5265 5266 5267
        for index, op in reversed(tuple(enumerate(list(block.ops)))):
            # remove the cast op of fp16 grad to fp32 grad
            if self._is_optimize_op(op) and op.type == 'cast':
                in_name = op.input_arg_names[0]
                out_name = op.output_arg_names[0]
                if out_name.strip('@GRAD') in self._param_device_map:
                    assert in_name.replace('.cast_fp16', '') == out_name
                    block._remove_op(index)
5268
                    continue
5269

5270
            if self._is_backward_op(op) and first_opt_op_idx is None:
5271
                first_opt_op_idx = index + 1
5272 5273
                # maybe have no optimize
                # if first_opt_op_idx == len(block.ops): return
5274

5275 5276
            if self._is_backward_op(op) and (self._op_role_var_key
                                             in op.attr_names):
5277 5278
                op_role_var = op.attr(self._op_role_var_key)
                if len(op_role_var) == 0: continue
5279 5280
                assert len(op_role_var) % 2 == 0
                for i in range(0, len(op_role_var), 2):
5281 5282 5283 5284
                    offset = 0
                    param_name = op_role_var[i]
                    if not block.has_var(param_name): continue
                    if '@BroadCast' in param_name: continue
5285

5286
                    param_grad_name = param_name + core.grad_var_suffix()
5287
                    merged_param_grad_name = param_grad_name + merged_suffix
5288 5289
                    if not block.has_var(merged_param_grad_name):
                        self._create_var(block, block.vars[param_name],
5290
                                         merged_param_grad_name, dtype)
5291
                    assert block.has_var(merged_param_grad_name)
5292

5293 5294 5295
                    param_grad_var = block.var(param_grad_name)
                    merged_param_grad_var = block.var(merged_param_grad_name)
                    merged_param_grad_var.persistable = True
5296
                    block._insert_op(
5297 5298 5299 5300
                        index=first_opt_op_idx + offset,
                        type='fill_constant',
                        inputs={},
                        outputs={'Out': [merged_param_grad_var]},
5301
                        attrs={
5302 5303 5304 5305 5306 5307
                            'shape':
                            merged_param_grad_var.shape,
                            'dtype':
                            merged_param_grad_var.dtype,
                            'value':
                            float(0),
5308
                            # a trick to run this op once per mini-batch
5309 5310
                            self._op_role_key:
                            self._op_role.Optimize.LRSched,
5311 5312
                        })
                    offset += 1
5313 5314
                    grad_name = op_role_var[i + 1]
                    grad_var = block.vars[grad_name]
5315 5316 5317 5318 5319 5320 5321 5322 5323

                    is_fp16_grad = 'cast_fp16' in grad_name
                    need_cast = (is_fp16_grad is not fp16_allreduce)

                    if need_cast:
                        # if fp16_allreduce:
                        #     cast grad to fp16 to accumulate to merged gradient
                        # else:
                        #     cast grad to fp32 to accumulate to merged gradient
5324
                        cast_grad_var_name = param_grad_name + '@TMP'
5325 5326
                        cast_grad_var = self._create_var(
                            block, param_grad_var, cast_grad_var_name, dtype)
5327
                        cast_grad_var.persistable = False
5328 5329 5330 5331 5332 5333 5334 5335 5336 5337 5338 5339
                        block._insert_op(index=first_opt_op_idx + offset,
                                         type='cast',
                                         inputs={'X': grad_var},
                                         outputs={'Out': cast_grad_var},
                                         attrs={
                                             'in_dtype':
                                             grad_var.dtype,
                                             'out_dtype':
                                             cast_grad_var.dtype,
                                             self._op_role_key:
                                             self._op_role.Backward,
                                         })
5340
                        offset += 1
5341 5342 5343 5344 5345 5346 5347
                        grad_var = cast_grad_var

                    block._insert_op(
                        index=first_opt_op_idx + offset,
                        type='sum',
                        inputs={'X': [merged_param_grad_var, grad_var]},
                        outputs={'Out': merged_param_grad_var},
5348 5349 5350
                        attrs={
                            self._op_role_key: self._op_role.Backward,
                        })
5351 5352 5353 5354 5355 5356 5357 5358 5359 5360 5361 5362 5363 5364 5365 5366 5367 5368 5369 5370 5371 5372 5373 5374 5375 5376 5377
                    offset += 1
                    merged_gradient_names.append(merged_param_grad_name)

        if not fp16_allreduce: return merged_gradient_names

        first_opt_op_idx = None
        for index, op in reversed(tuple(enumerate(list(block.ops)))):
            if self._is_backward_op(op) and first_opt_op_idx is None:
                first_opt_op_idx = index + 1
                break
        assert first_opt_op_idx is not None

        # insert cast op from fp16->fp32
        # FIXME(wangxi): maybe put in sharding is better, for some grad
        #                is not in sharding device.
        for fp16_grad_name in merged_gradient_names:
            grad_name = fp16_grad_name.replace('@FP16', '')
            param_name = fp16_grad_name.replace('@GRAD@MERGED@FP16', '')

            if not block.has_var(grad_name):
                self._create_var(block, block.vars[param_name], grad_name)
            assert block.has_var(grad_name)

            fp16_grad_var = block.var(fp16_grad_name)
            grad_var = block.var(grad_name)
            grad_var.persistable = False

5378 5379 5380 5381 5382 5383 5384 5385 5386
            block._insert_op(index=first_opt_op_idx,
                             type='cast',
                             inputs={'X': fp16_grad_var},
                             outputs={'Out': grad_var},
                             attrs={
                                 'in_dtype': fp16_grad_var.dtype,
                                 'out_dtype': grad_var.dtype,
                                 self._op_role_key: self._op_role.Optimize,
                             })
5387

5388
        return merged_gradient_names
5389

5390 5391 5392
    def _insert_accumulate_gradients_with_fuse(self, main_block, fp16,
                                               fused_size, grad_param_pairs,
                                               first_opt_op_idx):
5393 5394
        grad_param_pairs = self._sort_grad_param_by_dtype(
            main_block, grad_param_pairs)
5395

5396 5397 5398 5399 5400 5401 5402 5403 5404 5405 5406 5407 5408 5409 5410 5411
        grad_param_segments = []
        merged_suffix = '@MERGED@FP16' if fp16 else '@MERGED'
        dtype = paddle.float16 if fp16 else paddle.float32
        cur_size = 0.
        last_dtype = None
        # split the grad based on dtype and fused size
        for grad, param in grad_param_pairs:
            real_grad = main_block.var(grad)
            # create the gradient merged var for each grad
            merged_grad_var = main_block.create_var(
                name=param + core.grad_var_suffix() + merged_suffix,
                dtype=dtype,
                shape=real_grad.shape,
                persistable=True,
                stop_gradient=False)
            real_param = main_block.var(param)
5412 5413
            if hasattr(real_param, 'is_distributed'):
                merged_grad_var.is_distributed = real_param.is_distributed
5414 5415 5416 5417 5418 5419 5420 5421 5422 5423 5424 5425 5426 5427 5428 5429 5430 5431 5432 5433 5434 5435 5436
            tmp_size = self._get_var_size(real_grad)
            # two strategies for splitting the grad
            # 1. the current segment's size reach the user defined grad_size_in_MB
            # 2. the upcoming grad holds different dtype compared with grads in current segment
            if len(grad_param_segments) == 0 \
                    or cur_size + tmp_size > fused_size \
                    or real_grad.dtype != last_dtype:
                grad_param_segments.append(
                    ([real_grad], [real_param], [merged_grad_var]))
                last_dtype = real_grad.dtype
                cur_size = 0.
            else:
                grad_param_segments[-1][0].append(real_grad)
                grad_param_segments[-1][1].append(real_param)
                grad_param_segments[-1][2].append(merged_grad_var)
                cur_size += tmp_size

        fused_gradients = []
        fused_merged_gradients = []
        # create fused vars for grad and param
        for grad_param_segment in grad_param_segments:
            grad_segment = grad_param_segment[0]
            merged_grad_segment = grad_param_segment[2]
5437 5438 5439 5440 5441
            fused_grad = main_block.create_var(name='FusedGrad_{}'.format(
                grad_segment[0].name),
                                               dtype=grad_segment[0].dtype,
                                               persistable=False,
                                               stop_gradient=False)
5442 5443 5444 5445 5446 5447 5448 5449 5450 5451 5452 5453 5454 5455 5456 5457 5458 5459 5460 5461 5462 5463 5464 5465 5466 5467 5468 5469 5470 5471 5472 5473 5474 5475 5476
            # keep the '.cast_fp16' info in the fuse var name
            fused_merged_grad_name_prefix = 'FusedMergedGrad.cast_fp16.' if \
                merged_grad_segment[0].dtype == paddle.float16 else 'FusedMergedGrad'
            fused_merged_grad_name = fused_merged_grad_name_prefix + '_{}'.format(
                merged_grad_segment[0].name)
            fused_merged_grad = main_block.create_var(
                name=fused_merged_grad_name,
                dtype=merged_grad_segment[0].dtype,
                persistable=True,
                stop_gradient=False)
            fused_gradients.append(fused_grad)
            fused_merged_gradients.append(fused_merged_grad)

        assert len(fused_gradients) == len(grad_param_segments)
        assert len(fused_merged_gradients) == len(grad_param_segments)

        # insert coalesce op at the start of the backward pass
        # use param as the coalesce input to make sure the two Fused vars are in same shape
        first_back_op_idx = None
        for index, op in enumerate(main_block.ops):
            if self._is_backward_op(op) and first_back_op_idx is None:
                first_back_op_idx = index
                break
        assert first_back_op_idx is not None
        offset = 0
        for i in range(len(grad_param_segments)):
            fused_grad = fused_gradients[i]
            fused_merged_grad = fused_merged_gradients[i]
            grads = grad_param_segments[i][0]
            params = grad_param_segments[i][1]
            merged_grads = grad_param_segments[i][2]
            main_block._insert_op_without_sync(
                first_back_op_idx + offset,
                type="coalesce_tensor",
                inputs={"Input": params},
5477 5478 5479 5480
                outputs={
                    "Output": grads,
                    "FusedOutput": fused_grad
                },
5481 5482 5483 5484 5485 5486 5487 5488 5489 5490 5491 5492 5493 5494 5495 5496
                attrs={
                    # Explanation of user_defined_size_of_dtype:
                    # In coalesce op, the align size is 256 bytes
                    # the float takes 4 bytes while fp16 takes 2 bytes.
                    # To meet the requirement, 128 fp16 or 64 float will be aligned
                    # Think the total shape of the input tensors if [64],
                    # if the dtype is float, then the shape of the fuse var is [64]
                    # however if the dytpe if fp16, the shape of the fuse var is [128],
                    # which will cause the fused vars' shape vary between each other.
                    # To make sure the shape of the fused vars are identical,
                    # we set the dtype of float and fp16 both to 2.
                    # Under this way, the fused vars' shape for float and fp16 are all [128]
                    "user_defined_size_of_dtype": 2,
                    "copy_data": False,
                    "use_align": True,
                    "dtype": grads[0].dtype,
5497 5498 5499 5500 5501 5502 5503
                    self._op_role_key: self._op_role.Backward,
                    # On npu, the nan/inf check login is different with gpu.
                    # If there are some not initialized sections in the fused var,
                    # and the value in those sections are nan/inf, it will trigger the nan/inf check.
                    # To avoid these problematic triggers, set constant is needed for npu
                    "set_constant": core.is_compiled_with_npu(),
                    "constant": float(0.0),
5504 5505 5506 5507 5508 5509 5510 5511 5512 5513 5514 5515 5516 5517 5518 5519 5520 5521 5522 5523 5524 5525 5526 5527 5528 5529 5530 5531 5532 5533 5534 5535 5536 5537 5538 5539
                })
            offset += 1
            # For the gradient_merged_fused_var, given a init value during the coalesce op
            # this will remove a problematic fill_constant op. This op role of this coalesce
            # is set to be LRSched to make this coalesce (with init) only run once
            main_block._insert_op_without_sync(
                first_back_op_idx + offset,
                type="coalesce_tensor",
                inputs={"Input": params},
                outputs={
                    "Output": merged_grads,
                    "FusedOutput": fused_merged_grad
                },
                attrs={
                    "user_defined_size_of_dtype": 2,
                    "set_constant": True,
                    "constant": float(0.0),
                    "copy_data": False,
                    "use_align": True,
                    "dtype": merged_grads[0].dtype,
                    self._op_role_key: self._op_role.Optimize.LRSched
                })
            offset += 1

        # insert gradient merge relating ops
        first_opt_op_idx += offset
        offset = 0
        for i in range(len(fused_gradients)):
            fused_grad = fused_gradients[i]
            fused_merged_grad = fused_merged_gradients[i]
            is_fp16_grad = 'cast_fp16' in fused_grad.name
            need_cast = (is_fp16_grad is not fp16)
            if need_cast:
                # for fp16 allreduce, cast fp32 grad to fp16
                # for fp32 allreduce, cast fp16 grad to fp32
                cast_grad_var_name = fused_grad.name + '@TMP'
5540 5541 5542 5543 5544 5545 5546 5547 5548 5549 5550 5551 5552 5553
                cast_grad_var = main_block.create_var(name=cast_grad_var_name,
                                                      dtype=dtype,
                                                      persistable=False,
                                                      stop_gradient=False)
                main_block._insert_op(index=first_opt_op_idx + offset,
                                      type='cast',
                                      inputs={'X': fused_grad},
                                      outputs={'Out': cast_grad_var},
                                      attrs={
                                          'in_dtype': fused_grad.dtype,
                                          'out_dtype': cast_grad_var.dtype,
                                          self._op_role_key:
                                          self._op_role.Backward,
                                      })
5554 5555 5556 5557 5558 5559 5560 5561 5562 5563 5564 5565 5566 5567 5568 5569 5570 5571
                offset += 1
                fused_grad = cast_grad_var
            main_block._insert_op(
                index=first_opt_op_idx + offset,
                type='sum',
                inputs={'X': [fused_merged_grad, fused_grad]},
                outputs={'Out': fused_merged_grad},
                attrs={self._op_role_key: self._op_role.Backward})
            offset += 1

        if fp16:
            # if using fp16 allreduce, the optimizer needs fp32 grads, cast them back to fp32
            for grad, param in grad_param_pairs:
                real_grad = main_block.var(grad)
                fp16_grad_name = param + core.grad_var_suffix() + '@MERGED@FP16'
                assert main_block.has_var(fp16_grad_name)
                fp16_grad = main_block.var(fp16_grad_name)
                fp32_grad_name = param + core.grad_var_suffix() + '@MERGED'
5572 5573 5574 5575 5576 5577 5578 5579 5580 5581 5582 5583 5584 5585 5586
                fp32_grad = main_block.create_var(name=fp32_grad_name,
                                                  dtype=paddle.float32,
                                                  shape=real_grad.shape,
                                                  persistable=False,
                                                  stop_gradient=False)
                main_block._insert_op(index=first_opt_op_idx + offset,
                                      type='cast',
                                      inputs={'X': fp16_grad},
                                      outputs={'Out': fp32_grad},
                                      attrs={
                                          'in_dtype': paddle.float16,
                                          'out_dtype': paddle.float32,
                                          self._op_role_key:
                                          self._op_role.Optimize,
                                      })
5587 5588 5589 5590 5591 5592
                offset += 1

        # replace the var with it's name, which will be used for inserting allreduce
        for i in range(len(fused_merged_gradients)):
            fused_merged_gradients[i] = fused_merged_gradients[i].name

5593
        return fused_merged_gradients, first_opt_op_idx
5594

5595 5596 5597 5598 5599 5600 5601 5602 5603 5604 5605 5606 5607 5608 5609 5610 5611 5612 5613 5614 5615 5616 5617 5618
    def _accumulate_gradients_with_fuse(self,
                                        main_block,
                                        fp16,
                                        fused_size,
                                        shard=None):
        first_opt_op_idx = None
        grad_param_pairs = []
        # obtain all param/grad pairs that needed to be fused
        for index, op in reversed(tuple(enumerate(list(main_block.ops)))):
            # remove the cast op of fp16 grad to fp32 grad
            if self._is_optimize_op(op) and op.type == 'cast':
                in_name = op.input_arg_names[0]
                out_name = op.output_arg_names[0]
                if out_name.strip('@GRAD') in self._param_device_map:
                    assert in_name.replace('.cast_fp16', '') == out_name
                    main_block._remove_op(index)
                    continue

            if self._is_backward_op(op) and first_opt_op_idx is None:
                first_opt_op_idx = index + 1
                # no optimize phase
                if first_opt_op_idx == len(main_block.ops):
                    return

5619 5620
            if self._is_backward_op(op) and (self._op_role_var_key
                                             in op.attr_names):
5621 5622 5623 5624 5625 5626 5627 5628 5629 5630 5631 5632 5633 5634 5635 5636 5637 5638 5639 5640 5641 5642 5643 5644 5645 5646 5647 5648 5649 5650 5651 5652
                op_role_var = op.attr(self._op_role_var_key)
                if len(op_role_var) == 0:
                    continue
                assert len(op_role_var) % 2 == 0
                for i in range(0, len(op_role_var), 2):
                    param_name = op_role_var[i]
                    if not main_block.has_var(param_name):
                        continue
                    if '@BroadCast' in param_name:
                        continue
                    grad_param_pairs.append(
                        (op_role_var[i + 1], op_role_var[i]))

        if len(grad_param_pairs) == 0:
            return

        nranks = shard.worker_num if shard else 1
        device_to_pairs = [[] for _ in range(nranks)]
        for pair in grad_param_pairs:
            root_id = shard.device(pair[1]) if shard else 0
            assert 0 <= root_id < nranks
            device_to_pairs[root_id].append(pair)

        all_fused_merged_gradients = []
        for pairs in device_to_pairs:
            fused_merged_gradients, first_opt_op_idx = \
                self._insert_accumulate_gradients_with_fuse(
                    main_block, fp16, fused_size, pairs, first_opt_op_idx)
            all_fused_merged_gradients += fused_merged_gradients

        main_block._sync_with_cpp()
        return all_fused_merged_gradients
5653

5654 5655 5656 5657 5658 5659 5660 5661 5662 5663 5664 5665 5666 5667 5668 5669 5670 5671
    def _sort_grad_param_by_dtype(self, main_block, grad_param_pairs):
        # sort the grad param paris by the dtype
        fp16_pairs = []
        fp32_pairs = []
        other_pairs = []
        for pairs in grad_param_pairs:
            dtype = main_block.var(pairs[0]).dtype
            if dtype == paddle.float32:
                fp32_pairs.append(pairs)
            elif dtype == paddle.float16:
                fp16_pairs.append(pairs)
            else:
                other_pairs.append(pairs)
        sorted_pairs = fp16_pairs
        sorted_pairs.extend(fp32_pairs)
        sorted_pairs.extend(other_pairs)
        return sorted_pairs

5672 5673 5674 5675 5676 5677 5678 5679 5680 5681 5682 5683 5684 5685 5686
    def _get_var_size(self, var):
        dtype_to_size = {
            core.VarDesc.VarType.FP16: 2,
            core.VarDesc.VarType.FP32: 4,
            core.VarDesc.VarType.FP64: 8,
            core.VarDesc.VarType.INT16: 2,
            core.VarDesc.VarType.INT32: 4,
            core.VarDesc.VarType.INT64: 8,
            core.VarDesc.VarType.BOOL: 1,
            core.VarDesc.VarType.UINT8: 1,
        }
        assert -1 not in var.shape
        return reduce(lambda x, y: x * y,
                      var.shape) * dtype_to_size[var.dtype] / 1024.0 / 1024.0

5687 5688
    def _add_sub_blocks(self, main_block, program_list):
        main_program = main_block.program
5689
        for prog in program_list:
5690 5691 5692 5693 5694 5695
            for op in prog.block(0).ops:
                if not op.has_attr('sub_block'):
                    continue
                origin_sub_block_id = op.attr('sub_block').id
                origin_sub_block = main_program.block(origin_sub_block_id)
                new_sub_block = prog._create_block(parent_idx=0)
5696 5697
                for sub_op in origin_sub_block.ops:
                    op_desc = sub_op.desc
5698 5699 5700
                    ap_op = new_sub_block.desc.append_op()
                    ap_op.copy_from(op_desc)
                new_sub_block._sync_with_cpp()
5701
                self._create_vars(new_sub_block, origin_sub_block)
5702
                op._set_attr('sub_block', new_sub_block)
5703 5704 5705 5706 5707 5708 5709 5710 5711 5712 5713 5714 5715 5716 5717 5718

    def _get_device_info(self, block):
        for op in block.ops:
            if not op._has_kernel(op.type): continue
            op_device = op.attr(self._op_device_key)
            return op_device

    def _process_persistable_vars_in_multi_sections(self, main_program,
                                                    startup_prog, program_list):
        """
        Special Case: process persistable vars that exist in
        multiple sections, e.g., shared weight
        """
        # var_info = {var_name: [program1, program2...]},
        # persistable var only
        var_info = dict()
5719
        for prog in program_list:
5720 5721
            block = prog.block(0)
            for var_name in block.vars:
5722
                if var_name == "double_buffer_0": continue
5723 5724 5725 5726 5727 5728 5729 5730 5731 5732 5733 5734 5735 5736 5737 5738 5739
                var = block.var(var_name)
                if not var.persistable: continue
                if not var_name in var_info:
                    var_info[var_name] = []
                if not prog in var_info[var_name]:
                    var_info[var_name].append(prog)
        for var_name in list(var_info.keys()):
            if len(var_info[var_name]) == 1:
                var_info.pop(var_name)

        # write_info = {var_name: program}, where program is the only program
        # in which the var named var_name is written.
        write_info = dict()
        for var_name in var_info.keys():
            for prog in var_info[var_name]:
                block = prog.block(0)
                for op in block.ops:
5740
                    if op.type == "recv_v2" or op.type == "create_py_reader" or \
5741
                        op.type == "read" or op.type == "update_loss_scaling":
5742
                        continue
5743 5744 5745 5746 5747 5748 5749 5750 5751 5752 5753 5754 5755 5756 5757 5758 5759 5760 5761
                    # We have processed lr related vars
                    if op.attr(self._op_role_key) == int(
                            self._op_role.Optimize.LRSched):
                        continue
                    if var_name in op.desc.output_arg_names():
                        assert var_name not in write_info, (
                            "two sections write the same var({}): second "
                            "op {}.".format(var_name, op))
                        write_info[var_name] = prog
                        break

        for var_name in var_info.keys():
            # Case 1: read only variables, no special process
            if not var_name in write_info: continue

            # Case 2: one write multiple reads
            write_prog = write_info[var_name]
            write_block = write_prog.block(0)
            write_device = self._get_device_info(write_block)
5762
            write_dev_index = int(write_device.split(':')[1])
5763 5764 5765
            all_progs = var_info[var_name]
            for prog in all_progs:
                if prog == write_prog: continue
5766 5767 5768
                read_block = prog.block(0)
                read_device = self._get_device_info(read_block)
                read_dev_index = int(read_device.split(':')[1])
5769 5770 5771 5772 5773 5774 5775 5776 5777
                pair = (write_dev_index, read_dev_index)
                pair_key = write_dev_index * 1000 + read_dev_index
                if pair not in self._pipeline_pair:
                    self._pipeline_pair.append(pair)
                    self._pp_ring_map[pair_key] = self.ring_id
                    ring_id = self.ring_id
                    self.ring_id += 1
                else:
                    ring_id = self._pp_ring_map[pair_key]
5778 5779 5780

                write_block._insert_op(
                    index=0,
5781
                    type='send_v2',
5782 5783 5784
                    inputs={
                        'X': write_block.var(var_name),
                    },
5785
                    attrs={
5786 5787 5788 5789
                        self._op_device_key:
                        write_device,
                        'use_calc_stream':
                        False,
5790 5791
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
5792 5793 5794 5795 5796 5797
                        self._op_role_key:
                        self._op_role.LRSched,
                        'peer':
                        read_dev_index,
                        'ring_id':
                        ring_id
5798 5799 5800
                    })
                read_block._insert_op(
                    index=0,
5801
                    type='recv_v2',
5802 5803
                    outputs={'Out': [read_block.var(var_name)]},
                    attrs={
5804 5805 5806 5807 5808 5809 5810 5811
                        'out_shape':
                        read_block.var(var_name).shape,
                        'dtype':
                        read_block.var(var_name).dtype,
                        self._op_device_key:
                        read_device,
                        'use_calc_stream':
                        False,
5812 5813
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
5814 5815 5816 5817 5818 5819
                        self._op_role_key:
                        self._op_role.LRSched,
                        'peer':
                        write_dev_index,
                        'ring_id':
                        ring_id
5820
                    })
5821 5822 5823 5824 5825 5826
                read_block._insert_op(
                    index=1,
                    type='c_sync_comm_stream',
                    inputs={'X': [read_block.var(var_name)]},
                    outputs={'Out': [read_block.var(var_name)]},
                    attrs={
5827 5828
                        self._op_device_key:
                        read_device,
5829 5830
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
5831 5832 5833 5834
                        self._op_role_key:
                        self._op_role.LRSched,
                        'ring_id':
                        ring_id
5835 5836 5837 5838 5839 5840 5841 5842 5843
                    })

    def _is_gradient_clip_op(self, op):
        return op.desc.has_attr("op_namescope") \
            and op.desc.attr("op_namescope").startswith("/gradient_clip")

    def _is_regularization_op(self, op):
        return op.desc.has_attr("op_namescope") \
            and op.desc.attr("op_namescope").startswith("/regularization")
H
hutuxian 已提交
5844

5845 5846 5847 5848 5849
    def _is_weight_decay_op(self, op):
        # in AdamW namescope is /optimizer_*/weight decay/
        return op.desc.has_attr("op_namescope") \
            and 'weight decay' in op.desc.attr("op_namescope")

5850 5851 5852 5853 5854
    def _get_input_output_info(self, block):
        '''
        Get info of op input and output.
        '''
        # A map from output var to op which generate it.
5855
        output_var_to_op = defaultdict(list)
5856
        # A map from var to op which takes it as input.
5857
        input_var_to_op = defaultdict(list)
5858

5859
        for index, op in enumerate(block.ops):
5860
            for var_name in op.input_arg_names:
5861
                input_var_to_op[var_name].append([op, index])
5862
            for var_name in op.output_arg_names:
5863 5864 5865 5866 5867 5868 5869 5870 5871 5872 5873 5874
                output_var_to_op[var_name].append([op, index])

        return output_var_to_op, input_var_to_op

    def _optimize_forward_send_sync(self, program):
        """
        optimize forward send's sync_comm_stream schedule
        """
        if self.schedule_mode != '1F1B': return

        block = program.block(0)

5875
        recv_type = 'recv_v2' if self.mp_degree == 1 else 'partial_recv'
5876 5877
        backward_recv_index = None
        for index, op in enumerate(block.ops):
5878
            if op.type == recv_type and self._is_backward_op(op):
5879 5880 5881
                backward_recv_index = index
                break

5882
        # last pipeline stage
5883 5884 5885 5886 5887 5888 5889 5890 5891 5892 5893 5894 5895 5896 5897 5898 5899 5900 5901 5902 5903 5904 5905
        if backward_recv_index is None: return

        offset = 0
        for index, op in enumerate(list(block.ops)):
            if index >= backward_recv_index: break
            if op.type == 'c_sync_comm_stream' and op.has_attr('pipeline_flag'):
                var_name = op.input_arg_names[0]
                var = block.var(var_name)
                block._remove_op(index + offset, sync=False)
                offset -= 1
                # NOTE:
                # 1. When the backward recv is completed, it indicates
                # that the forward send is completed too. So we only need
                # to use the NOP op to prevent memory release.
                # 2. Because we removed sync_comm_op,
                # we will insert NOP after recv_op.
                block._insert_op_without_sync(
                    index=backward_recv_index,
                    type='nop',
                    inputs={'X': [var]},
                    outputs={'Out': [var]},
                    attrs={self._op_role_key: self._op_role.Backward})
        block._sync_with_cpp()
5906

5907 5908 5909 5910 5911 5912 5913 5914 5915 5916 5917 5918 5919 5920 5921 5922 5923 5924 5925 5926 5927 5928 5929 5930 5931 5932 5933 5934 5935 5936 5937 5938 5939 5940 5941 5942
    def _mv_head_recv(self, program):
        """
        A pass to move the recv op to the beginning of
        the forward/backward phase
        """
        forward_insert_index = 0
        backward_insert_index = None
        block = program.global_block()
        num_ops = len(program.global_block().ops)
        for i in range(num_ops):
            insert_index = None
            op = program.global_block().ops[i]
            op_role = int(op.attr(self._op_role_key))
            if op_role == int(
                    self._op_role.Backward) and backward_insert_index is None:
                backward_insert_index = i
            if op.type != "partial_recv" and op.type != "partial_allgather" and op.type != "nop" and op.type != "recv_v2":
                continue
            if op_role == int(self._op_role.Forward):
                if i == forward_insert_index:
                    forward_insert_index += 1
                    continue
                insert_index = forward_insert_index
            elif op_role == int(self._op_role.Backward):
                if i == backward_insert_index:
                    backward_insert_index += 1
                    continue
                insert_index = backward_insert_index
            else:
                raise ValueError("Unknown op_role: {}".format(op_role))
            op_inputs = dict()
            for name in op.input_names:
                op_inputs[name] = op.input(name)
            op_outputs = dict()
            for name in op.output_names:
                op_outputs[name] = op.output(name)
5943 5944 5945 5946 5947
            block._insert_op_without_sync(index=insert_index,
                                          type=op.type,
                                          inputs=op_inputs,
                                          outputs=op_outputs,
                                          attrs=op.all_attrs())
5948 5949 5950 5951 5952 5953 5954
            block._remove_op(i + 1)
            if op_role == int(self._op_role.Forward):
                forward_insert_index += 1
            elif op_role == int(self._op_role.Backward):
                backward_insert_index += 1
        block._sync_with_cpp()

5955 5956 5957 5958 5959 5960 5961 5962 5963 5964 5965 5966 5967 5968 5969 5970 5971 5972 5973 5974 5975 5976 5977 5978 5979 5980 5981 5982 5983
    def _check_pipeline_persist_var(self, program):
        """
        Pipeline may need multiple forward before
        """
        block = program.global_block()

        persist_output = set()
        used_in_backward = set()
        for op in block.ops:
            if self._is_forward_op(op):
                for var_name in op.output_arg_names:
                    var = block.vars[var_name]
                    if var.persistable:
                        persist_output.add(var_name)
            elif self._is_backward_op(op):
                for var_name in op.input_arg_names:
                    if var_name in persist_output:
                        used_in_backward.add(var_name)
        if len(used_in_backward) == 0:
            return
        warnings.warn(
            "The pipeline requires multiple forward calculations before backward, "
            "so when the persistable var is changed in the forward, it may cause "
            "errors in the backward calculation who using this persistable var. "
            "However, some backward op don't need this var(NoNeedBufferVars), "
            "there will be no error at this time.\n"
            "So please check these persistable vars which changed in "
            "forward and used in backward:\n{}".format(used_in_backward))

H
hutuxian 已提交
5984 5985 5986 5987 5988
    def minimize(self,
                 loss,
                 startup_program=None,
                 parameter_list=None,
                 no_grad_set=None):
5989
        main_block = loss.block
5990
        self.origin_main_block = main_block
5991
        main_program = main_block.program
5992 5993
        if startup_program is None:
            startup_program = default_startup_program()
5994

5995 5996
        pipeline_opt = main_program._pipeline_opt
        assert pipeline_opt, 'Please use pipeline with fleet.'
5997 5998 5999 6000 6001 6002 6003
        required_keys = [
            'local_rank',
            'schedule_mode',
            'micro_batch_size',
            'ring_id',
            'global_ring_id',
            'use_sharding',
6004 6005
            'mp_degree',
            'mp_rank',
6006 6007
        ]
        for key in required_keys:
6008
            assert key in pipeline_opt, \
6009
                'Please use pipeline with fleet to use {}.'.format(key)
6010 6011 6012 6013 6014 6015 6016 6017
        self.local_rank = pipeline_opt['local_rank']
        self.schedule_mode = pipeline_opt['schedule_mode']
        self.micro_batch_size = pipeline_opt['micro_batch_size']
        self.use_sharding = pipeline_opt['use_sharding']
        self.ring_id = pipeline_opt['ring_id']
        self.global_ring_id = pipeline_opt['global_ring_id']
        self.mp_degree = pipeline_opt['mp_degree']
        self.mp_rank = pipeline_opt['mp_rank']
6018
        self.scale_gradient = pipeline_opt.get('scale_gradient', False)
6019 6020
        assert self.mp_degree >= 1
        assert 0 <= self.mp_rank < self.mp_degree
6021 6022 6023 6024

        optimize_ops, params_grads = self._optimizer.minimize(
            loss, startup_program, parameter_list, no_grad_set)
        self._param_device_map = self._origin_optimizer._param_device_map
6025

6026 6027
        self.output_var_to_op, self.input_var_to_op = \
            self._get_input_output_info(main_block)
6028 6029 6030
        # Step1: add default op_device attribute for ops.
        self._add_op_device_attr(main_block)
        device_list = self._check_validation(main_block)
6031 6032 6033 6034 6035 6036 6037 6038 6039 6040 6041

        def device_cmp(device1, device2):
            dev1_id = int(device1.split(':')[1])
            dev2_id = int(device2.split(':')[1])
            if dev1_id < dev2_id:
                return -1
            elif dev1_id > dev2_id:
                return 1
            else:
                return 0

6042 6043 6044 6045 6046
        sorted_device_list = sorted(device_list, key=cmp_to_key(device_cmp))
        assert sorted_device_list == device_list, (
            "With pipeline parallelism, you must use gpu devices one after "
            "another in the order of their ids.")
        # Step2: add send and recv ops between section boundaries
6047
        self._insert_sendrecv_ops_for_boundaries(main_block)
6048

6049
        # Step3: split program into sections and add pairs of
6050 6051
        # send and recv ops for data var.
        main_program = main_block.program
6052
        program_list = self._split_program(main_program, device_list)
6053
        for p in program_list:
6054
            self._create_vars(p.global_block(), main_block)
6055

L
lilong12 已提交
6056 6057 6058 6059 6060 6061 6062 6063
        if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None):
            self.local_rank = int(os.getenv("PADDLE_MANUAL_PIPELINE_STAGE"))
            assert self.local_rank < len(device_list), (
                "Manually specified "
                "pipeline stage must be less than total number of pipeline "
                "stages.")
        else:
            self.local_rank %= len(device_list)
6064 6065 6066
        # Step3.5: optimize forward send sync_comm to overlap send and recv
        self._optimize_forward_send_sync(program_list[self.local_rank])

6067
        # Step4: Special Case: process persistable vars that exist in
6068
        # multiple sections
6069
        # FIXME
6070 6071
        # self._process_persistable_vars_in_multi_sections(
        #     main_program, startup_program, program_list)
6072

6073
        # Step5: Add sub blocks for section programs
6074 6075
        self._add_sub_blocks(main_block, program_list)

6076
        place_list = []
6077 6078
        for dev in device_list:
            dev_index = int(dev.split(":")[1])
6079 6080 6081 6082
            if core.is_compiled_with_cuda():
                place_list.append(core.CUDAPlace(dev_index % 1))
            elif core.is_compiled_with_npu():
                place_list.append(core.NPUPlace(dev_index % 1))
6083

6084
        # Step6: Split startup program
6085 6086
        new_startup_program = self._split_startup_program(
            startup_program, self.local_rank)
6087 6088 6089 6090

        startup_program._pipeline_opt = {
            "startup_program": new_startup_program,
        }
6091
        real_block = program_list[self.local_rank].global_block()
6092 6093
        if not self.scale_gradient:
            self._insert_loss_scale(real_block)
6094
        if not self.use_sharding:
6095
            # Step7: clear gradients before each mini-batch and
6096 6097 6098 6099 6100
            # accumulate gradients during backward
            self._rename_gradient_var_name(real_block)
            real_block._sync_with_cpp()
            self._accumulate_gradients(real_block)
            real_block._sync_with_cpp()
6101

6102 6103 6104 6105
        if core.is_compiled_with_cuda():
            place_id = int(os.getenv("FLAGS_selected_gpus", "0"))
        elif core.is_compiled_with_npu():
            place_id = int(os.getenv("FLAGS_selected_npus", "0"))
6106 6107 6108
        # A pass to move the recv op to the beginning of
        # the forward/backward phase
        self._mv_head_recv(program_list[self.local_rank])
6109 6110 6111 6112 6113

        # A pass to check pipeline persist var which changed in
        # forward and used in backward
        self._check_pipeline_persist_var(program_list[self.local_rank])

6114
        main_program._pipeline_opt = {
H
hutuxian 已提交
6115 6116
            "trainer": "PipelineTrainer",
            "device_worker": "Section",
6117
            "pipeline_stage": self.local_rank,
6118
            "num_pipeline_stages": len(device_list),
6119
            "schedule_mode": self.schedule_mode,
6120
            "inner_parallelism": len(device_list),
6121 6122
            "section_program": program_list[self.local_rank],
            "place": place_list[self.local_rank],
6123
            "place_id": place_id,
6124
            "sync_steps": -1,
L
lilong12 已提交
6125
            "num_microbatches": self._num_microbatches,
H
hutuxian 已提交
6126 6127
            "start_cpu_core_id": self._start_cpu_core_id,
        }
6128
        return optimize_ops, params_grads, program_list, self._pipeline_pair, self._pp_ring_map
M
mapingshuo 已提交
6129 6130


M
mapingshuo 已提交
6131 6132
class RecomputeOptimizer(Optimizer):
    """
6133
	:api_attr: Static Graph
S
swtkiwi 已提交
6134

M
mapingshuo 已提交
6135 6136 6137 6138 6139 6140 6141 6142 6143 6144 6145 6146 6147 6148 6149 6150 6151 6152 6153 6154 6155 6156 6157 6158 6159 6160 6161 6162 6163 6164 6165 6166 6167 6168 6169 6170 6171 6172 6173 6174 6175 6176 6177 6178 6179 6180 6181 6182 6183 6184 6185 6186 6187 6188 6189 6190 6191 6192 6193 6194
    Recompute Optimizer Wrapper

    Normally, a training step contains three sub-steps: first, run forward
    Operators to calculate the loss; second, run backward Operators to 
    calculate gradient of the parameters; third, apply optimization method
    to update the value of the parameters.

    In the forward computation process, all variables that are needed by 
    backward computation process will be kept in memory, which occupy a great
    amount of memory when the network becomes very deep.

    Recompute split the network to k segments. In each segment, It will 
    recompute the forward Operators, before running backward operators. It is
    very helpful for saving memory.
 
    The Variables that separate a network to segments are called as checkpoints,
    and users should set it manually. The usage is very simple:

    Args:
        optimizer (Optimizer): The optimizer that is applied to parameters.

    Examples:
        .. code-block:: python

            import paddle.fluid as fluid
            import numpy as np
            def gen_data():
                return {"x": np.random.random(size=(32, 32)).astype('float32'),
                "y": np.random.randint(2, size=(32, 1)).astype('int64')}
            def mlp(input_x, input_y, hid_dim=128, label_dim=2):
                print(input_x)
                fc_1 = fluid.layers.fc(input=input_x, size=hid_dim)
                prediction = fluid.layers.fc(input=[fc_1], size=label_dim, act='softmax')
                cost = fluid.layers.cross_entropy(input=prediction, label=input_y)
                sum_cost = fluid.layers.reduce_mean(cost)
                return sum_cost, fc_1, prediction
            input_x = fluid.layers.data(name="x", shape=[32], dtype='float32')
            input_y = fluid.layers.data(name="y", shape=[1], dtype='int64')
            cost, fc_1, pred = mlp(input_x, input_y)

            sgd = fluid.optimizer.Adam(learning_rate=0.01)
            sgd = fluid.optimizer.RecomputeOptimizer(sgd)
            sgd._set_checkpoints([fc_1, pred])
            sgd.minimize(cost)

            print("Finished optimize")
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            exe.run(fluid.default_startup_program())
            step = 10

            for i in range(step):
                cost_val = exe.run(feed=gen_data(),
                       program=fluid.default_main_program(),
                       fetch_list=[cost.name])
                print("step=%d cost=%f" % (i, cost_val[0]))

    """

    def __init__(self, optimizer):
J
Jiabin Yang 已提交
6195
        if framework._non_static_mode():
Z
zhongpu 已提交
6196
            raise Exception("In dygraph, don't support RecomputeOptimizer.")
M
mapingshuo 已提交
6197 6198
        self._optimizer = optimizer
        self._checkpoints = None
M
mapingshuo 已提交
6199 6200
        self._learning_rate = self._optimizer._learning_rate
        self._learning_rate_map = self._optimizer._learning_rate_map
J
JZ-LIANG 已提交
6201
        self.enable_offload = False
M
mapingshuo 已提交
6202 6203

    def _set_checkpoints(self, checkpoints):
6204 6205 6206 6207 6208 6209 6210 6211 6212
        """
        Args:
            checkpoints (list): List of Variable or string    
        """
        assert isinstance(
            checkpoints, list
        ), "_checkpoints should be a list of Variable or a list of String"
        for ckpt in checkpoints:
            assert (
6213 6214
                isinstance(ckpt, six.string_types)
                or isinstance(ckpt, Variable)
6215
            ), "_checkpoints should be a list of Variable or a list of String"
M
mapingshuo 已提交
6216 6217
        self._checkpoints = checkpoints

6218
    # should enable offload before calling backward
J
JZ-LIANG 已提交
6219 6220 6221
    def _enable_offload(self):
        self.enable_offload = True

6222 6223
    @framework.deprecate_stat_dict
    def load(self, state_dict):
M
mapingshuo 已提交
6224
        """
6225
	    :api_attr: Static Graph
S
swtkiwi 已提交
6226

M
mapingshuo 已提交
6227 6228 6229 6230
        load function is not supported by Recompute Optimizer for now.
        :return: None

        Args:
6231
            state_dict: the dict load by load_persistable method
M
mapingshuo 已提交
6232 6233 6234 6235 6236 6237 6238 6239 6240 6241 6242 6243 6244 6245 6246 6247 6248 6249 6250 6251 6252 6253 6254

        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
                import paddle.compat as cpt
                
                def mlp(input_x, input_y, hid_dim=128, label_dim=2):
                    fc_1 = fluid.layers.fc(input=input_x, size=hid_dim)
                    prediction = fluid.layers.fc(input=[fc_1], size=label_dim, act='softmax')
                    cost = fluid.layers.cross_entropy(input=prediction, label=input_y)
                    sum_cost = fluid.layers.reduce_mean(cost)
                    return sum_cost, fc_1, prediction
                
                input_x = fluid.layers.data(name="x", shape=[32], dtype='float32')
                input_y = fluid.layers.data(name="y", shape=[1], dtype='int64')
                cost, fc_1, pred = mlp(input_x, input_y)
                print("Finished FF")
                
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
                sgd._set_checkpoints([fc_1, pred])
                try:
6255 6256
                    state_dict = {}
                    sgd.load(state_dict)
M
mapingshuo 已提交
6257 6258 6259 6260 6261 6262 6263 6264 6265 6266 6267 6268 6269 6270 6271 6272 6273 6274 6275 6276 6277 6278 6279 6280 6281 6282 6283 6284 6285 6286 6287 6288 6289 6290 6291 6292 6293
                except NotImplementedError as e:
                    print(cpt.get_exception_message(e))
        """
        raise NotImplementedError(
            "load function is not supported by Recompute Optimizer for now")

    def apply_gradients(self, params_grads):
        """
        call apply_gradients function of self._optimizer.

        Args:
            params_grads (list): list of (param, grad) pair to do optimization.

        Returns:
            list: A list of operators appended to the current program.

        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
                import paddle.fluid.framework as framework

                def mlp(input_x, input_y, hid_dim=128, label_dim=2):
                    fc_1 = fluid.layers.fc(input=input_x, size=hid_dim)
                    prediction = fluid.layers.fc(input=[fc_1], size=label_dim, act='softmax')
                    cost = fluid.layers.cross_entropy(input=prediction, label=input_y)
                    sum_cost = fluid.layers.reduce_mean(cost)
                    return sum_cost, fc_1, prediction


                input_x = fluid.layers.data(name="x", shape=[32], dtype='float32')
                input_y = fluid.layers.data(name="y", shape=[1], dtype='int64')
                cost, fc_1, pred = mlp(input_x, input_y)
                print("Finished FF")

                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
6294
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
6295 6296 6297 6298
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
6299
                    no_grad_set=None)
M
mapingshuo 已提交
6300 6301 6302 6303 6304 6305 6306 6307 6308 6309

                program = cost.block.program
                with framework.program_guard(program, None):
                    optimize_ops = sgd.apply_gradients(params_grads)

                print("Finished apply gradients")
        """

        return self._optimizer.apply_gradients(params_grads=params_grads)

J
JZ-LIANG 已提交
6310 6311 6312 6313 6314 6315 6316 6317 6318 6319 6320 6321 6322 6323 6324 6325 6326 6327 6328 6329 6330 6331 6332 6333 6334 6335 6336 6337 6338 6339 6340 6341 6342 6343 6344 6345 6346 6347 6348 6349 6350 6351
    def _creat_vars(self, varname):
        pinned_var_name = unique_name.generate(varname + "@Pinned")
        fetched_var_name = unique_name.generate(varname + "@Fetch")

        pinned_var = self._main_program.global_block().create_var(
            name=pinned_var_name,
            shape=self.checkpoint_shape,
            dtype=self._main_program.global_block().var(varname).dtype,
            persistable=False,
            stop_gradient=True)

        fetch_var = self._main_program.global_block().create_var(
            name=fetched_var_name,
            shape=self.checkpoint_shape,
            dtype=self._main_program.global_block().var(varname).dtype,
            persistable=False,
            stop_gradient=False)

        return pinned_var_name, fetched_var_name

    def _append_fill_constant_ops(self, startup_program):
        """
        add fill_constant_ops to the end of the prog

        we should fill the pinned vars before runing the main_prog
        to instantiate their tensor hold_, which could tell us whether 
        the host memory could hold all the checkpoints from all the 
        GPU devices in this node. 
        """
        op_role = 0
        block = startup_program.global_block()
        fill_constant_vars = self.checkpoint_name2pinned_name.values()
        OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
        for varname in fill_constant_vars:
            var = self._main_program.global_block().var(varname)
            # NOTE (JZ-LIANG) to pre-allocate the CUDAPinned MEM
            pinned_var = block.create_var(
                name=varname,
                shape=self.checkpoint_shape,
                dtype=self._main_program.global_block().var(var.name).dtype,
                persistable=False,
                stop_gradient=True)
6352 6353 6354 6355 6356 6357 6358 6359 6360
            block.append_op(type='fill_constant',
                            outputs={'Out': varname},
                            attrs={
                                "shape": var.shape,
                                "dtype": var.dtype,
                                "value": 0.0,
                                "place_type": 2,
                                OP_ROLE_KEY: op_role,
                            })
J
JZ-LIANG 已提交
6361 6362 6363 6364

        return

    def _insert_async_memcpy_op(self, insert_idx, src_varname, dst_varname,
6365
                                op_role, dst_place_type):
J
JZ-LIANG 已提交
6366 6367 6368 6369 6370 6371 6372 6373
        OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
        self.block._insert_op_without_sync(
            insert_idx,
            type='memcpy',
            inputs={'X': [self._main_program.global_block().var(src_varname)]},
            outputs={
                'Out': [self._main_program.global_block().var(dst_varname)]
            },
6374 6375 6376 6377
            attrs={
                "dst_place_type": int(dst_place_type),
                OP_ROLE_KEY: op_role
            })
J
JZ-LIANG 已提交
6378 6379 6380 6381 6382 6383 6384

    def _insert_fetch_op(self, idx, varname):
        assert varname in self.checkpoint_name2pinned_name, "Try to fetch {} from Pinned Memory, but it is NOT a checkpoint".format(
            varname)

        pinned_varname = self.checkpoint_name2pinned_name[varname]
        fetch_varname = self.checkpoint_name2fetch_name[varname]
6385
        self._insert_async_memcpy_op(idx, pinned_varname, fetch_varname, 1, 1)
J
JZ-LIANG 已提交
6386 6387 6388 6389 6390

    def _insert_offload_op(self, idx, varname):
        assert varname in self.checkpoint_name2pinned_name, "Try to offload {} to Pinned Memory, but it is NOT a checkpoint".format(
            varname)
        pinned_varname = self.checkpoint_name2pinned_name[varname]
6391
        self._insert_async_memcpy_op(idx, varname, pinned_varname, 0, 2)
J
JZ-LIANG 已提交
6392 6393

    def _insert_sync_op(self, op_idx, checkpoint_name):
6394
        # single stream offload no need sync
J
JZ-LIANG 已提交
6395 6396 6397 6398 6399 6400 6401 6402 6403 6404 6405 6406 6407 6408 6409 6410 6411 6412 6413 6414 6415 6416 6417 6418 6419 6420 6421 6422
        pass

    def _record_fetch_op(self, idx):
        assert len(self.un_fetch_checkpoint_names
                   ) > 0, "Could NOT found checkpoint to fetch"
        checkpoint_name = self.un_fetch_checkpoint_names.pop(-1)
        logging.debug("Record fetch [{}]".format(checkpoint_name))
        self.idx2insertions[idx] = ("fetch", checkpoint_name)

        return checkpoint_name

    def _record_offload_op(self, idx, checkpoint_name):
        expected_checkpoint_name = self.un_offload_checkpoint_names.pop(0)
        assert checkpoint_name == expected_checkpoint_name, "expected to offload [{}] but got [{}]".format(
            expected_checkpoint_name, checkpoint_name)
        logging.debug("Record offload [{}]".format(checkpoint_name))
        self.idx2insertions[idx] = ("offload", checkpoint_name)

    def _record_sync_op(self, idx, checkpoint_name):
        assert checkpoint_name not in self.synced_checkpoints, "Try to sync the checkpoint [{}] twice".format(
            checkpoint_name)
        self.synced_checkpoints.add(checkpoint_name)
        logging.debug("Record offload sync [{}]".format(checkpoint_name))
        self.idx2insertions[idx] = ("sync", checkpoint_name)

    def _parse_backward(self):

        self.idx2insertions = {}
6423
        # don't offload the last checkpoints, to favor throughput
J
JZ-LIANG 已提交
6424 6425 6426 6427 6428 6429 6430 6431 6432 6433 6434 6435 6436 6437 6438 6439 6440 6441 6442 6443 6444 6445 6446 6447 6448 6449 6450 6451 6452 6453 6454 6455
        self.un_fetch_checkpoint_names = self.sorted_checkpoint_names[:]
        self.un_fetch_checkpoint_names.pop(-1)
        need_fetch_checkpoint_names = self.un_fetch_checkpoint_names[:]
        self.checkpoint_usage_count = {}
        for checkpoint_name in self.un_fetch_checkpoint_names:
            self.checkpoint_usage_count[checkpoint_name] = 0

        self.bw_strart_op_idx = len(self.block.ops)
        for idx, op in enumerate(self.block.ops):
            if int(op.desc.attr("op_role")) == 1:
                self.bw_strart_op_idx = idx
                break

        assert self.bw_strart_op_idx < len(
            self.block.ops), "Could NOT found backword op in prog"

        # fetch second to last checkpoint at the beginning of BW
        fetched_checkpoint_varname = self._record_fetch_op(
            self.bw_strart_op_idx)
        last_last_fetch_checkpoint = None

        for i, op in enumerate(self.block.ops[self.bw_strart_op_idx:]):
            idx = self.bw_strart_op_idx + i
            input_vars = op.desc.input_arg_names()

            for input_var in input_vars:
                if input_var in need_fetch_checkpoint_names:
                    if input_var not in self.un_fetch_checkpoint_names:
                        # fetch the  offloade checkpoint when the first usage of its previous one
                        if self.checkpoint_usage_count[input_var] == 0:
                            # TODO (JZ-LIANG) sync memcpy_stream if extra stream for memcpy
                            second_to_last_fetch_checkpoint = fetched_checkpoint_varname
6456
                            # there is NO fetch ahead the first checkpoint
J
JZ-LIANG 已提交
6457 6458 6459 6460
                            if input_var != self.sorted_checkpoint_names[0]:
                                fetched_checkpoint_varname = self._record_fetch_op(
                                    idx)

6461
                        # should check the current used checkpoint is ths last fetch one
J
JZ-LIANG 已提交
6462 6463 6464 6465 6466 6467 6468 6469 6470 6471 6472 6473 6474 6475 6476 6477 6478 6479 6480 6481 6482 6483 6484 6485 6486
                        assert second_to_last_fetch_checkpoint == input_var, "Current recompute segment should use [{}] BUT got [{}]".format(
                            second_to_last_fetch_checkpoint, input_var)
                        # rename
                        self.block.ops[idx]._rename_input(
                            input_var,
                            self.checkpoint_name2fetch_name[input_var])
                        self.checkpoint_usage_count[input_var] += 1
                    else:
                        raise ValueError(
                            "use checkpoint [{}] before fetch in BW".format(
                                input_var))

        assert len(self.un_fetch_checkpoint_names
                   ) == 0, "{} checkpoints have NOT been Recorded".format(
                       self.un_fetch_checkpoint_names)

    def _update_backward(self):
        if len(self.idx2insertions) == 0:
            return
        total_op = len(self.block.ops)
        for op_idx in reversed(range(self.bw_strart_op_idx, total_op)):
            if op_idx in self.idx2insertions:
                operation, checkpoint_name = self.idx2insertions[op_idx]
                if operation == "fetch":
                    self._insert_fetch_op(op_idx, checkpoint_name)
6487 6488
                    logging.debug(
                        "Insert [{}] fetch op.".format(checkpoint_name))
J
JZ-LIANG 已提交
6489 6490 6491 6492 6493 6494 6495 6496 6497 6498 6499 6500
                    del self.idx2insertions[op_idx]
                elif operation == "sync":
                    self._insert_sync_op(op_idx, checkpoint_name)
                    logging.debug("Sync [{}] fetch op.".format(checkpoint_name))
        self.block._sync_with_cpp()
        assert len(
            self.idx2insertions) == 0, "{} checkpoints left un-Fecthed".format(
                [ele[1] for ele in self.idx2insertions.values()])

    def _parse_forward(self):

        self.idx2insertions = {}
6501
        # don't offload the last checkpoints, faster, less memory saving
J
JZ-LIANG 已提交
6502 6503 6504 6505 6506 6507 6508 6509 6510 6511 6512 6513 6514 6515 6516 6517 6518 6519 6520 6521
        self.un_offload_checkpoint_names = self.sorted_checkpoint_names[:]
        last_checkpoint = self.un_offload_checkpoint_names.pop(-1)
        need_offload_checkpoint_names = self.un_offload_checkpoint_names[:]
        self.checkpoint_usage_count_and_idx = {}
        for checkpoint_name in self.un_offload_checkpoint_names:
            self.checkpoint_usage_count_and_idx[checkpoint_name] = {
                'count': 0,
                'idx': -1
            }
        self.synced_checkpoints = set()
        self.fw_strart_op_idx = len(self.block.ops)
        for idx, op in enumerate(self.block.ops):
            if int(op.desc.attr("op_role")) == 0:
                self.fw_strart_op_idx = idx
                break

        assert self.fw_strart_op_idx < len(
            self.block.ops), "Could NOT found Forward op in prog"
        last_offload_checkpoint = None

6522 6523
        for i, op in enumerate(
                self.block.ops[self.fw_strart_op_idx:self.bw_strart_op_idx]):
J
JZ-LIANG 已提交
6524 6525 6526 6527 6528 6529 6530 6531 6532 6533 6534 6535 6536 6537 6538 6539 6540 6541 6542 6543 6544 6545 6546 6547 6548 6549 6550 6551 6552 6553 6554

            idx = self.fw_strart_op_idx + i
            output_vars = op.desc.output_arg_names()
            input_vars = op.desc.input_arg_names()

            for output_var in output_vars:
                if output_var in need_offload_checkpoint_names:
                    assert len(
                        output_vars
                    ) == 1, "chekpoint should be the only Output of a certain op, but [{}] is from [{}]".format(
                        output_var, op)

                    if output_var in self.un_offload_checkpoint_names:
                        # insert sync op if last checkpoint has not been sync
                        if last_offload_checkpoint != None:
                            if self.checkpoint_usage_count_and_idx[
                                    last_offload_checkpoint]['count'] == 0:
                                self._record_sync_op(idx,
                                                     last_offload_checkpoint)
                            else:
                                last_usage_idx = self.checkpoint_usage_count_and_idx[
                                    last_offload_checkpoint]['idx']
                                assert last_usage_idx > 0, "last_usage_idx of checkpoint [{}] should large than 0".format(
                                    last_offload_checkpoint)
                                self._record_sync_op(last_usage_idx + 1,
                                                     last_offload_checkpoint)
                        # insert offload op after the checkpoint's generation op
                        self._record_offload_op(idx + 1, output_var)
                        last_offload_checkpoint = output_var
                    else:
                        raise ValueError(
6555 6556
                            "There should be just ONE op that output checkpoint [{}]"
                            .format(output_var))
J
JZ-LIANG 已提交
6557 6558 6559 6560 6561 6562 6563 6564 6565 6566 6567 6568 6569 6570 6571 6572 6573 6574 6575 6576 6577
                # need to sync the last need to offload checkpoint before the last checkpoint as output op
                if output_var == last_checkpoint:
                    assert len(
                        output_vars
                    ) == 1, "chekpoint should be the only Output of a certain op, but [{}] is from [{}]".format(
                        output_var, op)
                    assert last_offload_checkpoint == self.sorted_checkpoint_names[
                        -2], "the last offload chekpoint before [{}] is suppose to be [{}], but got [{}]".format(
                            last_checkpoint, self.sorted_checkpoint_names[-2],
                            last_offload_checkpoint)
                    # sync if last checkpoint has not been sync
                    if self.checkpoint_usage_count_and_idx[
                            last_offload_checkpoint]['idx'] == 0:
                        self._record_sync_op(idx, last_offload_checkpoint)
                    else:
                        last_usage_idx = self.checkpoint_usage_count_and_idx[
                            last_offload_checkpoint]['idx']
                        assert last_usage_idx > 0, "last_usage_idx of checkpoint [{}] should large than 0".format(
                            last_offload_checkpoint)
                        self._record_sync_op(last_usage_idx + 1,
                                             last_offload_checkpoint)
6578
            # record checkpoint usage
J
JZ-LIANG 已提交
6579 6580 6581 6582 6583 6584 6585 6586 6587 6588 6589 6590 6591 6592 6593 6594 6595 6596 6597 6598 6599 6600 6601 6602
            for input_var in input_vars:
                if input_var in need_offload_checkpoint_names:
                    assert input_var not in self.synced_checkpoints, "checkpoint [{}] used after sync".format(
                        input_var)
                    self.checkpoint_usage_count_and_idx[input_var]['count'] += 1
                    self.checkpoint_usage_count_and_idx[input_var]['idx'] = idx

        assert len(self.un_offload_checkpoint_names
                   ) == 0, "{} checkpoints have NOT been Recorded".format(
                       self.un_fetch_checkpoint_names)
        assert len(self.synced_checkpoints) == len(
            need_offload_checkpoint_names
        ), "{} checkpoints have NOT been Recorded".format(
            set(need_offload_checkpoint_names) - set(self.synced_checkpoints))

    def _update_forward(self):
        if len(self.idx2insertions) == 0:
            return
        for op_idx in reversed(
                range(self.fw_strart_op_idx, self.bw_strart_op_idx)):
            if op_idx in self.idx2insertions:
                operation, checkpoint_name = self.idx2insertions[op_idx]
                if operation == "offload":
                    self._insert_offload_op(op_idx, checkpoint_name)
6603 6604
                    logging.debug(
                        "Insert [{}] offload op.".format(checkpoint_name))
J
JZ-LIANG 已提交
6605 6606 6607
                    del self.idx2insertions[op_idx]
                elif operation == "sync":
                    self._insert_sync_op(op_idx, checkpoint_name)
6608 6609
                    logging.debug(
                        "Insert [{}] offload_sync op.".format(checkpoint_name))
J
JZ-LIANG 已提交
6610 6611 6612 6613 6614 6615 6616 6617 6618 6619 6620 6621 6622 6623 6624 6625 6626 6627 6628 6629 6630 6631
                    del self.idx2insertions[op_idx]

        self.block._sync_with_cpp()
        assert len(self.idx2insertions
                   ) == 0, "{} checkpoints left un-Offloaded".format(
                       [ele[1] for ele in self.idx2insertions.values()])

    def _check_offload_fetch(self):
        # TODO(JZ-LIANG) the single stream offload need no sync
        pass

    def _offload(self, loss, startup_program=None):
        """
        core steps for recompute offload
        1. create pinned vars and temp vars 
        2. parse & update Forward pass: offload, sync
        3. parse & update Backward pass: rename, fetch, sync
        4. verify the correctness
        """
        self._main_program = loss.block.program
        self.block = loss.block
        if startup_program == None:
J
JZ-LIANG 已提交
6632
            startup_program = paddle.static.default_startup_program()
J
JZ-LIANG 已提交
6633 6634 6635

        with program_guard(self._main_program, startup_program):
            assert len(self.checkpoint_shape) > 0, (
6636 6637
                "checkpoints shape {} should be an non empty list like: [12, 512, 1024]"
                .format(self.checkpoint_shape))
J
JZ-LIANG 已提交
6638
            assert all([ele > 0 for ele in self.checkpoint_shape]), (
6639 6640
                "all ele in checkpoints shape {} should be a determined integer larger than 0"
                .format(self.checkpoint_shape))
J
JZ-LIANG 已提交
6641 6642 6643 6644 6645 6646 6647 6648 6649 6650 6651 6652 6653 6654 6655 6656 6657 6658 6659 6660 6661 6662
            self.checkpoint_name2pinned_name = dict()
            self.checkpoint_name2fetch_name = dict()
            for checkpoint_varname in self.sorted_checkpoint_names:
                pinned_var_name, fetch_var_name = self._creat_vars(
                    checkpoint_varname)
                self.checkpoint_name2pinned_name[
                    checkpoint_varname] = pinned_var_name
                self.checkpoint_name2fetch_name[
                    checkpoint_varname] = fetch_var_name
            self._append_fill_constant_ops(startup_program)
            # TODO (JZ-LIANG) to provide two offload stragtegy in future
            # step 2. parse & update FW: rename, offload, sync
            self._parse_backward()
            self._update_backward()
            # step 3. parse & update BW: rename, offload, sync
            self._parse_forward()
            self._update_forward()
            # step 4. verify the correctness
            self._check_offload_fetch()

        return

M
mapingshuo 已提交
6663 6664 6665 6666 6667
    def backward(self,
                 loss,
                 startup_program=None,
                 parameter_list=None,
                 no_grad_set=None,
6668
                 callbacks=None):
M
mapingshuo 已提交
6669 6670 6671 6672 6673 6674 6675
        """
        call append_backward with checkpoints.

        Args:
            loss (Variable): loss variable to run optimizations.
            startup_program (Program): startup_program for initializing parameters
                in `parameter_list`.
6676 6677
            parameter_list (list): list of Variables or Variable.names to update.
            no_grad_set (set|None): set of Variables or Variables.names should be ignored.
M
mapingshuo 已提交
6678 6679 6680 6681 6682 6683 6684 6685 6686 6687 6688 6689 6690 6691 6692 6693 6694 6695 6696 6697 6698 6699 6700 6701
            callbacks (list|None): list of callables to run when appending backward
                operator for one parameter.
            checkpoints (list): list of Variables as checkpoints

        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
    
                def mlp(input_x, input_y, hid_dim=128, label_dim=2):
                    fc_1 = fluid.layers.fc(input=input_x, size=hid_dim)
                    prediction = fluid.layers.fc(input=[fc_1], size=label_dim, act='softmax')
                    cost = fluid.layers.cross_entropy(input=prediction, label=input_y)
                    sum_cost = fluid.layers.reduce_mean(cost)
                    return sum_cost, fc_1, prediction
    
    
                input_x = fluid.layers.data(name="x", shape=[32], dtype='float32')
                input_y = fluid.layers.data(name="y", shape=[1], dtype='int64')
                cost, fc_1, pred = mlp(input_x, input_y)
                print("Finished FF")
    
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
6702
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
6703 6704 6705 6706
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
6707
                    no_grad_set=None)
M
mapingshuo 已提交
6708 6709
                print("Finished backward")
        """
6710 6711
        assert (self._checkpoints
                is not None), "You should call _set_checkpoints first"
M
mapingshuo 已提交
6712

J
Jiabin Yang 已提交
6713
        if framework._non_static_mode():
M
mapingshuo 已提交
6714 6715 6716 6717 6718 6719
            raise NotImplementedError(
                "DyGraph current does not support recompute")

        self._dtype = loss.dtype
        program = loss.block.program
        with program_guard(program, startup_program):
6720 6721 6722 6723 6724 6725 6726
            checkpoint_vars = []
            for ckpt in self._checkpoints:
                if isinstance(ckpt, Variable):
                    checkpoint_vars.append(ckpt)
                else:
                    checkpoint_vars.append(loss.block.var(ckpt))

J
JZ-LIANG 已提交
6727 6728 6729 6730 6731 6732 6733 6734
            # allow return to non-recompute when checkpoints is empty
            if len(checkpoint_vars) > 0:
                params_grads, sorted_checkpoint_names = append_backward(
                    loss,
                    parameter_list,
                    no_grad_set,
                    checkpoints=checkpoint_vars)
            else:
6735 6736 6737 6738
                params_grads = append_backward(loss,
                                               parameter_list,
                                               no_grad_set,
                                               checkpoints=checkpoint_vars)
J
JZ-LIANG 已提交
6739 6740 6741 6742 6743

        if self.enable_offload:
            self.sorted_checkpoint_names = sorted_checkpoint_names
            self._offload(loss, startup_program=startup_program)

M
mapingshuo 已提交
6744 6745 6746 6747 6748 6749 6750 6751 6752 6753 6754 6755 6756 6757 6758 6759 6760 6761 6762
        return params_grads

    def apply_optimize(self, loss, startup_program, params_grads):
        """
        call the apply_optimize function of self._optimizer
        Args:
            loss (Variable): loss variable to run optimizations.
            startup_program (Program): startup_program for initializing parameters
                in `parameter_list`.
            params_grads (list): list of (param, grad) pair to do optimization.
        Examples:
            .. code-block:: python
                import paddle.fluid as fluid
                
                def mlp(input_x, input_y, hid_dim=128, label_dim=2):
                    fc_1 = fluid.layers.fc(input=input_x, size=hid_dim)
                    prediction = fluid.layers.fc(input=[fc_1], size=label_dim, act='softmax')
                    cost = fluid.layers.cross_entropy(input=prediction, label=input_y)
                    sum_cost = fluid.layers.reduce_mean(cost)
M
mapingshuo 已提交
6763
                    return sum_cost, fc_1, prediction                
M
mapingshuo 已提交
6764 6765 6766 6767 6768 6769 6770 6771
                
                input_x = fluid.layers.data(name="x", shape=[32], dtype='float32')
                input_y = fluid.layers.data(name="y", shape=[1], dtype='int64')
                cost, fc_1, pred = mlp(input_x, input_y)
                print("Finished FF")
                
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
6772
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
6773 6774 6775 6776
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
6777
                    no_grad_set=None)
M
mapingshuo 已提交
6778 6779 6780 6781 6782 6783 6784
                
                optimize_ops = sgd.apply_optimize(
                    cost, startup_program=None, params_grads=params_grads)
                
                print("Finished apply_optimize")
        """

Y
Yuang Liu 已提交
6785 6786 6787
        func = self._optimizer.apply_optimize if hasattr(
            self._optimizer,
            'apply_optimize') else self._optimizer._apply_optimize
6788 6789 6790
        return func(loss,
                    startup_program=startup_program,
                    params_grads=params_grads)
M
mapingshuo 已提交
6791 6792 6793 6794 6795

    def minimize(self,
                 loss,
                 startup_program=None,
                 parameter_list=None,
6796
                 no_grad_set=None):
6797
        assert isinstance(loss, Variable), "The loss should be an Variable."
6798 6799
        assert (self._checkpoints
                is not None), "You should call _set_checkpoints first"
J
Jiabin Yang 已提交
6800
        if framework._non_static_mode():
M
mapingshuo 已提交
6801 6802
            raise NotImplementedError(
                "DyGraph current does not support recompute")
6803 6804 6805 6806
        params_grads = self.backward(loss,
                                     startup_program=startup_program,
                                     parameter_list=parameter_list,
                                     no_grad_set=no_grad_set)
M
mapingshuo 已提交
6807

6808 6809 6810
        optimize_ops = self.apply_optimize(loss,
                                           startup_program=startup_program,
                                           params_grads=params_grads)
M
mapingshuo 已提交
6811 6812 6813 6814

        return optimize_ops, params_grads


M
mapingshuo 已提交
6815
class LookaheadOptimizer(object):
6816
    r"""
6817
	:api_attr: Static Graph
S
swtkiwi 已提交
6818

M
mapingshuo 已提交
6819 6820 6821 6822 6823 6824 6825 6826 6827 6828 6829 6830 6831 6832 6833 6834 6835 6836 6837 6838 6839 6840 6841 6842 6843
    This implements the Lookahead optimizer of the
    paper : https://arxiv.org/abs/1907.08610.

    Lookahead keeps two sets of params: the fast_params and
    the slow_params. inner_optimizer update fast_params every 
    training step. Lookahead updates the slow_params and fast_params 
    every k training steps as follows:

    .. math::
        
        slow\_param_t &= slow\_param_{t-1} + \\alpha * (fast\_param_{t-1} - slow\_param_{t-1})
	
	fast\_param_t &=  slow\_param_t

    Args:
        inner_optimizer (Optimizer): The optimizer that update fast params step by step. 
        alpha (float): The learning rate of Lookahead.
        k (int): The slow params is updated every k steps.

    Examples:
        .. code-block:: python

            import paddle
            import paddle.fluid as fluid
            import numpy as np
6844
            import numpy.random as random
M
mapingshuo 已提交
6845

6846 6847 6848 6849 6850 6851
            paddle.enable_static()
        
            x = fluid.layers.data(name='x', shape=[2], dtype='float32')
            label = fluid.layers.data(name="label", shape=[1], dtype="int64")
            y = fluid.layers.fc(input=[x], size=2, act="softmax")
            loss = fluid.layers.cross_entropy(input=y, label=label)
6852
            loss = paddle.mean(x=loss)
6853 6854 6855 6856 6857 6858 6859 6860 6861
            sgd = fluid.optimizer.SGD(learning_rate=0.01)
            optimizer = fluid.optimizer.LookaheadOptimizer(sgd,
                                                alpha=0.5,
                                                k=5)
            optimizer.minimize(loss)
            main_program = fluid.default_main_program()
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            exe.run(fluid.default_startup_program())
M
mapingshuo 已提交
6862

6863 6864 6865 6866 6867 6868 6869 6870 6871 6872
            def train_reader(limit=5):
                for i in range(limit):
                    yield random.random([2]).astype('float32'), random.random([1]).astype('int64')
            
            feeder = fluid.DataFeeder(feed_list=[x, label], place=place)
            reader = paddle.batch(paddle.reader.shuffle(train_reader, buf_size=50000),batch_size=1)
            
            for batch_data in reader():
                exe.run(fluid.default_main_program(),
                feed=feeder.feed(batch_data))
M
mapingshuo 已提交
6873 6874 6875 6876 6877

    """

    def __init__(self, inner_optimizer, alpha=0.5, k=5):

J
Jiabin Yang 已提交
6878
        if framework._non_static_mode():
Z
zhongpu 已提交
6879
            raise Exception("In dygraph, don't support LookaheadOptimizer.")
M
mapingshuo 已提交
6880 6881 6882 6883 6884 6885 6886 6887 6888 6889 6890 6891 6892 6893 6894 6895 6896 6897 6898 6899 6900 6901 6902 6903 6904 6905 6906 6907
        assert (inner_optimizer is not None), "inner optimizer can not be None"
        assert (
            0.0 <= alpha <= 1.0
        ), "alpha should be larger or equal to 0.0, and less or equal than 1.0"
        assert (isinstance(k, int) and k > 0), "k should be a positive integer"

        self.inner_optimizer = inner_optimizer
        self.alpha = alpha
        self.k = k
        self.type = "lookahead"

    def minimize(self, loss, startup_program=None):

        # Apply inner optimizer to the main_program
        mini_out = self.inner_optimizer.minimize(
            loss, startup_program=startup_program)

        # Get startup_program and main_program
        if startup_program is None:
            startup_program = default_startup_program()
        main_block = loss.block

        # add some vars to the main_program
        params = [param.name for param in main_block.all_parameters()]
        param_to_slow = {}
        for param in params:
            fast_var = main_block.var(param)
            assert (fast_var is not None)
6908 6909 6910 6911
            slow_var = main_block.create_var(name=param + "@SLOW",
                                             shape=fast_var.shape,
                                             dtype=fast_var.dtype,
                                             persistable=True)
M
mapingshuo 已提交
6912 6913 6914 6915 6916 6917 6918
            param_to_slow[param] = slow_var

        # add some vars to the startup_program
        startup_block = startup_program.global_block()
        for param in params:
            fast_var = startup_block.var(param)
            assert (fast_var is not None)
6919 6920 6921 6922
            slow_var = startup_block.create_var(name=param + "@SLOW",
                                                shape=fast_var.shape,
                                                dtype=fast_var.dtype,
                                                persistable=True)
M
mapingshuo 已提交
6923

6924 6925 6926
            startup_block.append_op(type="assign",
                                    inputs={"X": fast_var},
                                    outputs={"Out": slow_var})
M
mapingshuo 已提交
6927

6928 6929
        with framework.program_guard(main_block.program, startup_program):
            # Add Var k to main prog and startup prog
6930 6931 6932 6933 6934
            k = layers.create_global_var(name="lookahead_k",
                                         shape=[1],
                                         value=int(self.k),
                                         dtype='int32',
                                         persistable=True)
M
mapingshuo 已提交
6935

6936
            # Add Var alpha to main prog and startup prog
6937 6938 6939 6940 6941
            alpha = layers.create_global_var(name="lookahead_alpha",
                                             shape=[1],
                                             value=float(self.alpha),
                                             dtype='float32',
                                             persistable=True)
M
mapingshuo 已提交
6942

6943
            # Add Var step
6944 6945 6946 6947 6948
            step = layers.create_global_var(name="lookahead_step",
                                            shape=[1],
                                            value=int(0),
                                            dtype='int32',
                                            persistable=True)
6949 6950 6951
            layers.increment(x=step, value=1.0, in_place=True)

            # lookahead
6952 6953 6954
            zero_var = layers.fill_constant(shape=[1],
                                            dtype='float32',
                                            value=0.0)
6955

6956 6957 6958
            one_var = layers.fill_constant(shape=[1],
                                           dtype='float32',
                                           value=1.0)
6959 6960 6961

            mod = layers.elementwise_mod(step, k)
            with layers.control_flow.Switch() as switch:
6962 6963 6964 6965 6966
                with switch.case(step == one_var):
                    for param_name in params:
                        fast_var = main_block.var(param_name)
                        slow_var = param_to_slow[param_name]
                        layers.assign(input=fast_var, output=slow_var)
6967 6968 6969 6970 6971 6972 6973 6974 6975 6976 6977 6978 6979
                with switch.case(mod == zero_var):
                    for param_name in params:
                        fast_var = main_block.var(param_name)
                        slow_var = param_to_slow[param_name]
                        tmp_var = layers.elementwise_add(
                            layers.elementwise_mul(fast_var, alpha),
                            layers.elementwise_mul(
                                slow_var,
                                layers.elementwise_sub(one_var, alpha)))
                        layers.assign(input=tmp_var, output=slow_var)
                        layers.assign(input=tmp_var, output=fast_var)
                with switch.default():
                    pass
M
mapingshuo 已提交
6980
        return mini_out
6981 6982 6983 6984 6985 6986 6987 6988 6989 6990 6991 6992 6993 6994 6995 6996 6997 6998 6999 7000 7001 7002 7003 7004 7005 7006 7007 7008 7009 7010 7011 7012 7013 7014 7015 7016 7017 7018 7019 7020 7021 7022 7023 7024 7025 7026 7027 7028 7029 7030 7031 7032 7033 7034 7035 7036 7037


class GradientMergeOptimizer(object):
    """
    Gradient Merge, also called as Gradient Accumulation,
    is a training strategy for larger batches. With this strategy,
    the parameter will not be updated until specific steps.

    For each step, the forward network and the backward network
    will run to calculate the gradient of the parameters.

    For every k step, the optimization network will run,
    applying a specific optimization method (such as SGD, Adam)
    to the parameters.

    Args:
        inner_optimizer (Optimizer): The specific optimization (such as SGD, Adam)
            which update the parameters
        k_steps (int): the update period of the parameters
        avg (bool): whether to average the gradients of each mini-batch,
            the default value is `True`

    Examples:
        .. code-block:: python

        import paddle.fluid as fluid
        import numpy as np

        def gen_data(batch_size):
            return {"x": np.random.random(size=(batch_size, 32)).astype('float32'),
                    "y": np.random.random(size=(batch_size, 1)).astype('int64')}

        def mlp(input_x, input_y, hid_dim=128, label_dim=2):
            fc_1 = fluid.layers.fc(input=input_x, size=hid_dim)
            prediction = fluid.layers.fc(input=[fc_1], size=label_dim, act='softmax')
            cost = fluid.layers.cross_entropy(input=prediction, label=input_y)
            sum_cost = fluid.layers.reduce_mean(cost)
            return sum_cost, fc_1, prediction

        input_x = fluid.layers.data(name="x", shape=[32], dtype='float32')
        input_y = fluid.layers.data(name="y", shape=[1], dtype='int64')
        cost, fc_1, pred = mlp(input_x, input_y)
        sgd = fluid.optimizer.Adam(learning_rate=0.01)
        sgd = fluid.optimizer.GradientMergeOptimizer(sgd, k_steps=4, avg=True)
        sgd.minimize(cost)

        place = fluid.CPUPlace()
        exe = fluid.Executor(place)
        exe.run(fluid.default_startup_program())

        for i in range(10):
            cost_val = exe.run(feed=gen_data(32),
                       program=fluid.default_main_program(),
                       fetch_list=[cost.name])
            print("step=%d, cost=%f" % (i, cost_val[0]))
    """

7038 7039
    GRAD_MERGE_COND_NAME = "grad_merge_cond_name"

7040
    def __init__(self, inner_optimizer, k_steps=1, avg=True):
J
Jiabin Yang 已提交
7041
        if framework._non_static_mode():
7042 7043 7044 7045 7046 7047
            raise Exception(
                "In dygraph, we don't support GradientMergeOptimizer."
                "You can do Gradient merge by yourself with k-times forward + backward, "
                "and one-time optimizer.minimize()")

        assert (inner_optimizer is not None), "inner optimizer can not be None"
7048 7049
        assert (isinstance(k_steps, int)
                and k_steps > 0), "k_steps should be a positive integer"
7050 7051 7052 7053 7054

        self.inner_optimizer = inner_optimizer
        self.k_steps = k_steps
        self.type = "gradient_merge"
        self.avg = avg
7055
        self._optimize_ops = None
7056

7057 7058 7059 7060 7061 7062
    def _set_k_steps(self, k_steps):
        self.k_steps = k_steps

    def _set_avg(self, avg):
        self.avg = avg

7063
    def backward(self,
7064 7065 7066
                 loss,
                 startup_program=None,
                 parameter_list=None,
7067 7068
                 no_grad_set=None,
                 callbacks=None):
7069 7070 7071 7072 7073 7074 7075 7076 7077 7078
        assert isinstance(loss, Variable), "The loss should be an Variable."
        assert (
            parameter_list is None
        ), "The parameter_list should be None when using GradientMergeOptimizer"
        assert (
            no_grad_set is None
        ), "The no_grad_set should be None when using GradientMergeOptimizer"

        params_grads = self.inner_optimizer.backward(
            loss, startup_program=startup_program)
7079 7080 7081 7082 7083 7084 7085 7086 7087 7088 7089 7090 7091 7092 7093 7094 7095 7096 7097 7098 7099 7100 7101 7102 7103 7104 7105 7106 7107 7108 7109 7110 7111 7112 7113 7114 7115 7116 7117 7118 7119 7120 7121 7122 7123 7124 7125 7126 7127 7128 7129 7130 7131 7132 7133 7134 7135
        return params_grads

    def apply_optimize(self, loss, startup_program, params_grads):
        program = loss.block.program
        with program_guard(program, startup_program):
            optimize_ops = self.apply_gradients(params_grads)
        return optimize_ops

    def _is_the_backward_op(self, op):
        op_maker = core.op_proto_and_checker_maker
        backward = core.op_proto_and_checker_maker.OpRole.Backward
        if op_maker.kOpRoleVarAttrName() in op.attr_names and \
                int(op.all_attrs()[op_maker.kOpRoleAttrName()]) == int(backward):
            return True
        return False

    def _remove_op_role_var(self, param, grad):
        op_maker = core.op_proto_and_checker_maker
        op = grad.op
        assert self._is_the_backward_op(op), \
            'grad.op={} is not the backward op which produces the grad={}' \
            .format(op, grad.name)

        block = grad.block
        var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
        assert param.name in var_attr, \
            'when using GradientMergeOptimizer, param={} must be in var_attr={}' \
            .format(param.name, var_attr)
        assert grad.name in var_attr, \
            'when using GradientMergeOptimizer, grad={} must be in var_attr={}' \
            .format(param.name, var_attr)

        # remove (param, grad) from op_role_var
        var_attr.remove(param.name)
        var_attr.remove(grad.name)
        if len(var_attr) > 1:
            op._set_attr(op_maker.kOpRoleVarAttrName(), var_attr)
        else:
            op._remove_attr(op_maker.kOpRoleVarAttrName())

    def _add_gm_op_role_var(self, op, param, grad, cond):
        grad.op = op
        op_maker = core.op_proto_and_checker_maker
        backward = op_maker.OpRole.Backward

        # NOTE(wangxi). When distributed, we will insert grad_merge_all_reduce_op_handle
        # in multi_devices_graph_pass, which will allreduce(grad) if cond is True, else
        # do nothing.
        # In this way, the gradient can be merged first, and then communicate when the
        # condition is met, reducing the number of communications to increase the
        # speed.
        op._set_attr(self.GRAD_MERGE_COND_NAME, cond.name)
        op._set_attr(op_maker.kOpRoleAttrName(), backward)
        op._set_attr(op_maker.kOpRoleVarAttrName(), [param.name, grad.name])

    def _get_gm_cond_var(self, main_block):
        # Add const var
7136 7137 7138 7139 7140 7141 7142 7143 7144 7145 7146 7147 7148
        k_step_var = layers.create_global_var(name="gradient_merge_k",
                                              shape=[1],
                                              value=int(self.k_steps),
                                              dtype='int32',
                                              persistable=True,
                                              force_cpu=True)

        zero_var = layers.create_global_var(name="gradient_merge_zero",
                                            shape=[1],
                                            value=int(0),
                                            dtype='int32',
                                            persistable=True,
                                            force_cpu=True)
7149 7150

        # Add step var & cond var
7151 7152 7153 7154 7155 7156
        step_var = layers.create_global_var(name="gradient_merge_step",
                                            shape=[1],
                                            value=int(0),
                                            dtype='int32',
                                            persistable=True,
                                            force_cpu=True)
7157

7158 7159 7160
        cond_var = main_block.create_var(name="gradient_merge_cond",
                                         shape=[1],
                                         dtype='bool')
7161 7162 7163 7164

        with device_guard("cpu"):
            # step_var = (step_var + 1) % k_step
            layers.increment(x=step_var, value=1.0, in_place=True)
7165 7166 7167 7168 7169 7170 7171 7172 7173 7174
            main_block.append_op(type='elementwise_mod',
                                 inputs={
                                     'X': step_var,
                                     'Y': k_step_var
                                 },
                                 outputs={'Out': step_var},
                                 attrs={
                                     'axis': -1,
                                     'use_mkldnn': False
                                 })
7175 7176

            # cond_var = (step_var == 0)
7177 7178 7179 7180 7181 7182
            main_block.append_op(type='equal',
                                 inputs={
                                     'X': step_var,
                                     'Y': zero_var
                                 },
                                 outputs={'Out': cond_var})
7183 7184 7185 7186 7187 7188 7189 7190 7191 7192

        return cond_var

    def apply_gradients(self, params_grads):
        main_program = default_main_program()
        startup_program = default_startup_program()
        main_block = main_program.global_block()
        startup_block = startup_program.global_block()

        cond = self._get_gm_cond_var(main_block)
7193 7194

        #TODO(mapingshuo) support sparse embedding
7195 7196
        # step1: remove grad.op's op_role_var
        for param, grad in params_grads:
7197
            assert (
7198
                param.type != core.VarDesc.VarType.SELECTED_ROWS
7199 7200
            ), "SELECTED_ROWS is not supported in GradientMergeOptimizer for now"

7201
            self._remove_op_role_var(param, grad)
7202

7203
        param_to_grad = {k.name: v for (k, v) in params_grads}
7204 7205 7206
        param_names = param_to_grad.keys()
        param_to_gradient_merge = {}

7207 7208 7209 7210 7211
        new_params_grads = []
        # step2: create gradient_merge var and init with 0
        # and update op_role_var
        for param, grad in params_grads:
            param_name = param.name
7212 7213
            param_var = main_block.var(param_name)
            assert (param_var is not None)
7214 7215 7216 7217 7218
            gradient_merge_var = main_block.create_var(name=param_name +
                                                       "@GRAD@GradientMerge",
                                                       shape=param_var.shape,
                                                       dtype=param_var.dtype,
                                                       persistable=True)
7219
            param_to_gradient_merge[param_name] = gradient_merge_var
7220

7221 7222 7223 7224 7225
            startup_gradient_merge_var = startup_block.create_var(
                name=param_name + "@GRAD@GradientMerge",
                shape=param_var.shape,
                dtype=param_var.dtype,
                persistable=True)
7226 7227 7228 7229 7230 7231 7232
            startup_block.append_op(type="fill_constant",
                                    outputs={"Out": startup_gradient_merge_var},
                                    attrs={
                                        "shape": param_var.shape,
                                        "dtype": param_var.dtype,
                                        "value": float(0),
                                    })
7233

7234 7235 7236
            # grad_merge += grad
            new_grad_op = main_block.append_op(
                type="elementwise_add",
7237 7238 7239 7240
                inputs={
                    'X': grad,
                    'Y': gradient_merge_var
                },
7241
                outputs={'Out': gradient_merge_var},
7242 7243 7244 7245
                attrs={
                    'axis': -1,
                    'use_mkldnn': False
                })
7246 7247 7248 7249 7250 7251 7252 7253 7254 7255
            self._add_gm_op_role_var(new_grad_op, param, gradient_merge_var,
                                     cond)
            new_params_grads.append([param, gradient_merge_var])

        def true_apply_gradient():
            cur_block_idx = main_program.current_block_idx
            cur_block = main_program.current_block()

            # cur_block's forward_block & backward_block is itself
            cur_block._set_forward_block_idx(cur_block_idx)
7256
            op_maker = core.op_proto_and_checker_maker
7257 7258 7259 7260

            if self.avg:
                for param, new_grad in new_params_grads:
                    # grad /= k_steps
7261 7262 7263 7264 7265 7266 7267 7268
                    cur_block.append_op(type='scale',
                                        inputs={'X': new_grad},
                                        outputs={'Out': new_grad},
                                        attrs={
                                            'scale': 1.0 / self.k_steps,
                                            'bias': 0.0,
                                            'bias_after_scale': False
                                        })
7269 7270
                    new_grad.op._set_attr(op_maker.kOpRoleAttrName(),
                                          op_maker.OpRole.Backward)
7271

7272 7273 7274 7275 7276 7277
            for param, new_grad in new_params_grads:
                # NOTE. regularization will append ops to grad.block,
                # while new_grad's real block is global_block,
                # but we want append regularization ops to cur_block,
                # so we set new_grad.block = cur_block
                new_grad.block = cur_block
7278

7279 7280
            self._optimize_ops = self.inner_optimizer.apply_gradients(
                new_params_grads)
7281

7282 7283
            # clear gradient_merge_vars
            for param, new_grad in new_params_grads:
7284 7285 7286 7287
                layers.fill_constant(shape=new_grad.shape,
                                     dtype=new_grad.dtype,
                                     value=0.0,
                                     out=new_grad)
7288 7289
                new_grad.op._set_attr(op_maker.kOpRoleAttrName(),
                                      op_maker.OpRole.Optimize)
7290 7291 7292 7293 7294 7295 7296 7297 7298 7299 7300 7301 7302

        # step3. apply gradient
        layers.cond(cond, true_fn=true_apply_gradient, false_fn=None)

        return self._optimize_ops

    def minimize(self,
                 loss,
                 startup_program=None,
                 parameter_list=None,
                 no_grad_set=None):
        assert isinstance(loss, Variable), "The loss should be an Variable."

7303 7304 7305 7306
        params_grads = self.backward(loss,
                                     startup_program=startup_program,
                                     parameter_list=parameter_list,
                                     no_grad_set=no_grad_set)
7307

7308 7309 7310
        optimize_ops = self.apply_optimize(loss,
                                           startup_program=startup_program,
                                           params_grads=params_grads)
7311 7312

        return optimize_ops, params_grads