optimizer.py 301.9 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
import numpy as np
16
import os
17
import logging
18
from collections import defaultdict
19

20
import paddle
Q
Qiao Longfei 已提交
21
from paddle.fluid.distribute_lookup_table import find_distributed_lookup_table
22 23 24 25 26 27 28 29 30
from paddle.fluid.framework import (
    Program,
    Variable,
    Parameter,
    name_scope,
    default_main_program,
    default_startup_program,
    device_guard,
)
31

32 33
from . import framework
from . import layers
34
from . import unique_name
35 36 37 38 39 40 41 42 43 44 45 46 47
from .backward import (
    append_backward,
    _some_in_set_,
    _append_grad_suffix_,
    _get_no_grad_set_name,
)
from .clip import (
    GradientClipBase,
    GradientClipByNorm,
    error_clip_callback,
    append_gradient_clip_ops,
    ClipGradByGlobalNorm,
)
48 49 50
from .framework import program_guard
from .initializer import Constant
from .layer_helper import LayerHelper
51
from .dygraph import base as imperative_base
52
from .dygraph import no_grad
53 54 55 56
from .dygraph.learning_rate_scheduler import (
    LearningRateDecay,
    _LearningRateEpochDecay,
)
57 58 59
from paddle.fluid import core
from paddle.fluid.layers import tensor
from functools import reduce
60
from functools import cmp_to_key
61
from .wrapped_decorator import signature_safe_contextmanager
62
import warnings
63
from paddle import _C_ops, _legacy_C_ops
64 65 66 67 68
from ..fluid.framework import (
    _in_legacy_dygraph,
    in_dygraph_mode,
    _current_expected_place,
)
69

70
__all__ = [
71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97
    'SGD',
    'Momentum',
    'Adagrad',
    'Adam',
    'Adamax',
    'Dpsgd',
    'DecayedAdagrad',
    'Ftrl',
    'SGDOptimizer',
    'MomentumOptimizer',
    'AdagradOptimizer',
    'AdamOptimizer',
    'AdamaxOptimizer',
    'DpsgdOptimizer',
    'DecayedAdagradOptimizer',
    'RMSPropOptimizer',
    'FtrlOptimizer',
    'Adadelta',
    'AdadeltaOptimizer',
    'ModelAverage',
    'LarsMomentum',
    'LarsMomentumOptimizer',
    'LambOptimizer',
    'ExponentialMovingAverage',
    'PipelineOptimizer',
    'LookaheadOptimizer',
    'RecomputeOptimizer',
98
]
Q
Qiao Longfei 已提交
99 100


101
class Optimizer:
Q
Qiao Longfei 已提交
102 103 104
    """Optimizer Base class.

    Define the common interface of an optimizer.
105 106
    User should not use this class directly,
    but need to use one of it's implementation.
Q
Qiao Longfei 已提交
107 108
    """

109
    @imperative_base.no_grad
110 111 112 113 114 115 116 117 118 119
    def __init__(
        self,
        learning_rate,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        flatten_param_grads=False,
        align_size=-1,
        name=None,
    ):
120 121
        """
        Args:
122 123
            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,
124 125
                and the grad_clip ops / optimizer ops will be fuse to one operator.
        """
126
        # Because of the loop import, so place it in the function body
127
        from paddle.optimizer.lr import LRScheduler
128 129 130 131

        self._parameter_list = (
            list(parameter_list) if parameter_list is not None else None
        )
132
        self._name = name
J
Jiabin Yang 已提交
133
        if framework._non_static_mode():
134 135 136
            if not isinstance(
                learning_rate, (float, LearningRateDecay, LRScheduler)
            ):
M
minqiyang 已提交
137
                raise TypeError(
138
                    "learning rate should be float or LRScheduler, got %s here"
139 140
                    % type(learning_rate)
                )
141
            if self._parameter_list is None:
142 143 144
                raise AttributeError(
                    "parameter_list argument given to the Optimizer should not be None in dygraph mode."
                )
145 146 147 148 149 150
            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!"
151 152
                            % regularization.__str__()
                        )
153
                        break
M
minqiyang 已提交
154
        else:
155 156 157
            if not isinstance(
                learning_rate, (float, framework.Variable, LRScheduler)
            ):
M
minqiyang 已提交
158
                raise TypeError(
159
                    "learning rate should be float or LRScheduler, got %s here"
160 161
                    % type(learning_rate)
                )
M
minqiyang 已提交
162

163 164 165 166 167
        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 已提交
168
        self.regularization = regularization
169
        self._grad_clip = grad_clip
170
        self._learning_rate = learning_rate
171 172
        self._flatten_param_grads = flatten_param_grads
        self._align_size = align_size
L
Leo Chen 已提交
173

D
dzhwinter 已提交
174
        self._dtype = None
L
Leo Chen 已提交
175 176 177 178
        # Infer the dtype form parameter
        if self._parameter_list:
            self._dtype = self._parameter_list[0].dtype

179
        # each program should have a independent learning rate
180
        # program -> Variable(learning_rate)
Q
qiaolongfei 已提交
181
        self._learning_rate_map = dict()
182
        if isinstance(self._learning_rate, framework.Variable):
183
            self._learning_rate_map[
184 185
                framework.default_main_program()
            ] = self._learning_rate
186 187 188 189 190
        # 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())
191 192
        # global_accumulator dict, {accum_name : acc_variable, ...}
        self._global_accumulators = {}
193
        self.helper = LayerHelper(self.__class__.__name__)
194
        self._opti_name_list = []
H
hong 已提交
195
        self._accumulators_holder = {}
196
        self._param_device_map = dict()
197 198
        # 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.
199
        # And these variables should not be the parameters of Optimizer's construnctor (because not commonly used).
200 201
        # Use _auxiliary_vars together with _set_auxiliary_var/_get_auxiliary_var to achieve that.
        self._auxiliary_vars = dict()
H
hong 已提交
202 203 204 205

    @framework.dygraph_only
    def state_dict(self):
        '''
T
tianshuo78520a 已提交
206 207
        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 已提交
208 209 210

        Args: None
        Return:
T
tianshuo78520a 已提交
211
            state_dict(dict) : dict contains all the variable used by optimizer
212

H
hong 已提交
213 214 215 216
        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
217
                import paddle
218 219

                with fluid.dygraph.guard():
220
                    emb = paddle.nn.Embedding(10, 10)
221 222 223

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

        '''
226
        from paddle.optimizer.lr import LRScheduler
227

H
hong 已提交
228 229 230 231
        state_dict = {}
        for k, v in self._accumulators.items():
            for para_name, var_tmp in v.items():
                state_dict[var_tmp.name] = var_tmp
232 233
        for k, v in self._global_accumulators.items():
            state_dict[v.name] = v
H
hong 已提交
234
        # global step if use lr decay
235
        if isinstance(self._learning_rate, LRScheduler):
236 237
            state_dict["LR_Scheduler"] = self._learning_rate.state_dict()
            return state_dict
H
hong 已提交
238
        if isinstance(self._learning_rate, LearningRateDecay):
239 240 241 242
            state_dict["LR_Scheduler"] = self._learning_rate.state_dict()

            if not isinstance(self._learning_rate, _LearningRateEpochDecay):
                var_tmp = None
243 244 245
                var_temp = framework._varbase_creator(
                    None, name='global_step', dtype='int32'
                )
246

247 248 249
                tensor.fill_constant(
                    [1], "int32", self._learning_rate.step_num, out=var_temp
                )
H
hong 已提交
250

251
                state_dict['global_step'] = var_temp
H
hong 已提交
252 253 254
        return state_dict

    @framework.dygraph_only
255
    def set_state_dict(self, state_dict):
H
hong 已提交
256
        '''
T
tianshuo78520a 已提交
257
        Load optimizer state dict. For Adam optimizer, contains beta1, beta2, momentum etc. If LearningRateDecay have been used, global_step will be changed.
H
hong 已提交
258

259
        Args:
H
hong 已提交
260 261 262
            state_dict(dict) : Dict contains all the Variable needed by optimizer
        Return:
            None
263

H
hong 已提交
264 265
        Examples:
            .. code-block:: python
266

267 268
                import paddle
                import paddle.fluid as fluid
269 270 271

                paddle.disable_static()

272
                emb = paddle.nn.Embedding(10, 10)
273

274
                state_dict = emb.state_dict()
275
                fluid.save_dygraph(state_dict, "paddle_dy")
276

277
                scheduler = paddle.optimizer.lr.NoamDecay(
278 279 280 281
                    d_model=0.01, warmup_steps=100, verbose=True)
                adam = paddle.optimizer.Adam(
                    learning_rate=scheduler,
                    parameters=emb.parameters())
282
                state_dict = adam.state_dict()
283
                fluid.save_dygraph(state_dict, "paddle_dy")
284

285
                para_state_dict, opti_state_dict = fluid.load_dygraph("paddle_dy")
H
hong 已提交
286
        '''
287
        from paddle.optimizer.lr import LRScheduler
288

289
        if isinstance(self._learning_rate, LRScheduler):
290
            self._learning_rate.set_dict(state_dict["LR_Scheduler"])
H
hong 已提交
291 292

        if isinstance(self._learning_rate, LearningRateDecay):
293 294 295
            self._learning_rate.set_dict(state_dict["LR_Scheduler"])

            if not isinstance(self._learning_rate, _LearningRateEpochDecay):
296 297 298
                assert (
                    'global_step' in state_dict
                ), 'Global step not in state dict, Dygraph use LearningRateDecay, global_step must in state_dict'
299 300 301 302 303
                global_step = state_dict['global_step']

                if isinstance(global_step, Variable):
                    step_np = global_step
                    step_np = np.array(step_np.value().get_tensor())
304 305 306 307 308
                    assert step_np.shape == (
                        1,
                    ), "global step shape is (1,), the shape is {}".format(
                        step_np.shape
                    )
309 310 311

                    self._learning_rate.step_num = int(step_np[0])
                elif isinstance(global_step, np.ndarray):
312 313 314 315 316
                    assert global_step.shape == (
                        1,
                    ), "global step shape is (1,), the shape is {}".format(
                        global_step.shape
                    )
317 318 319 320
                    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 ",
321 322
                        type(global_step),
                    )
H
hong 已提交
323

324 325 326 327 328 329 330 331 332 333 334 335
        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:
336 337 338
                raise RuntimeError(
                    "State dict type {} not supprt".format(str(type(load_para)))
                )
339

340 341 342 343 344
            assert (
                model_np.shape == load_para_np.shape
            ), "Parameter shape not match, Dygraph Parameter [ {} ] need tensor with shape {} but load tensor with shape {}".format(
                param.name, model_np.shape, load_para_np.shape
            )
345

346 347 348 349 350
            assert (
                model_np.dtype == load_para_np.dtype
            ), "Parameter dtype not match, Dygraph Parameter [ {} ] need tensor with dtype {}  but load tensor with dtype {}".format(
                param.name, model_np.dtype, load_para_np.dtype
            )
351 352 353

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

H
hong 已提交
354 355 356
        self._accumulators_holder = state_dict
        for k, v in self._accumulators.items():
            for para_name, var_tmp in v.items():
357 358 359
                assert (
                    var_tmp.name in state_dict
                ), "optimizer variable {} not found".format(var_tmp.name)
360
                _load_state_para(state_dict, var_tmp)
H
hong 已提交
361

362
        for k, v in self._global_accumulators.items():
363 364 365
            assert (
                v.name in state_dict
            ), "optimizer variable {} not found".format(v.name)
366
            _load_state_para(state_dict, v)
367

368 369 370
    # [aliases] Compatible with old method names
    set_dict = set_state_dict

371 372
    def get_opti_var_name_list(self):
        return self._opti_name_list
Q
Qiao Longfei 已提交
373

374 375 376 377 378 379 380 381 382
    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 已提交
383
    def _create_global_learning_rate(self):
384
        from paddle.optimizer.lr import LRScheduler
385

386
        if isinstance(self._learning_rate, LRScheduler):
387 388 389 390 391 392 393 394 395 396
            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,
397 398
                    dtype='float32' if self._dtype is None else self._dtype,
                )
399 400 401
                main_prog = framework.default_main_program()
                main_prog.lr_sheduler = self._learning_rate
                main_prog.lr_var = lr_var
402
                self._learning_rate_map[
403 404
                    framework.default_main_program()
                ] = lr_var
405 406 407

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

412 413 414
        if imperative_base.enabled():
            # create learning rate Variable
            if isinstance(self._learning_rate, float):
M
minqiyang 已提交
415 416 417 418 419
                lr = self._global_learning_rate()

                if isinstance(lr, framework.Variable):
                    return
                else:
420 421
                    self._learning_rate_map[
                        framework.default_main_program()
422
                    ] = paddle.static.create_global_var(
M
minqiyang 已提交
423 424 425 426
                        name=unique_name.generate("learning_rate"),
                        shape=[1],
                        value=float(self._learning_rate),
                        dtype='float32' if self._dtype is None else self._dtype,
427 428
                        persistable=True,
                    )
429
            # get learning rate Variable from LearningRateDecay
M
minqiyang 已提交
430
            elif isinstance(self._learning_rate, LearningRateDecay):
431
                self._learning_rate_map[
432 433
                    framework.default_main_program()
                ] = self._learning_rate()
434
            else:
Q
qiaolongfei 已提交
435
                raise TypeError(
436 437
                    "optimizer's learning rate must be float or LearningRateDecay"
                )
438
        else:
439 440 441 442
            lr = self._global_learning_rate()

            if isinstance(lr, framework.Variable):
                return
M
minqiyang 已提交
443 444 445 446 447 448
            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 已提交
449

450
            # create learning rate in the current main program
451
            self._learning_rate_map[
452
                framework.default_main_program()
453
            ] = paddle.static.create_global_var(
454 455 456 457 458 459
                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,
            )
460

461 462 463 464
    @framework.dygraph_only
    def set_lr(self, value):
        """
        :api_attr: imperative
465

466 467 468 469 470 471 472 473
        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
474

475 476 477
        Examples:
            .. code-block:: python

478
                import paddle
479
                import paddle.fluid as fluid
480
                import paddle
481

482
                with fluid.dygraph.guard():
483
                    linear = paddle.nn.Linear(10, 10)
484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501

                    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
502
                    lr_var = paddle.static.create_global_var(
503 504 505 506 507 508 509 510 511 512 513 514 515
                        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."
516 517
                % (type(value))
            )
518 519 520 521 522 523 524 525
        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:
526 527
                if in_dygraph_mode():
                    place = _current_expected_place()
528 529 530 531 532 533 534
                    _C_ops.full_(
                        current_lr,
                        list(current_lr.shape),
                        float(value),
                        current_lr.dtype,
                        place,
                    )
535 536

                elif _in_legacy_dygraph():
537 538 539 540 541 542 543 544 545
                    _legacy_C_ops.fill_constant(
                        current_lr,
                        'value',
                        float(value),
                        'dtype',
                        current_lr.dtype,
                        'shape',
                        list(current_lr.shape),
                    )
546
                else:
547 548 549 550 551 552 553 554 555 556 557 558 559
                    global_block = (
                        framework.default_main_program().global_block()
                    )
                    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,
                    )
560
        else:
561 562 563
            assert (
                len(value.shape) == 1 and value.shape[0] == 1
            ), "optimizer's learning rate must be 1-D Tensor with shape[1]"
564 565
            self._learning_rate_map[framework.default_main_program()] = value

566 567 568
    @framework.dygraph_only
    def current_step_lr(self):
        """
569
        :api_attr: imperative
570

571 572 573 574 575 576 577 578 579 580 581
        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
582
                import paddle
583 584 585

                # example1: LearningRateDecay is not used, return value is all the same
                with fluid.dygraph.guard():
586
                    emb = paddle.nn.Embedding(10, 10)
587 588 589 590 591 592 593
                    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")
594
                    linear = paddle.nn.Linear(10, 10)
595 596
                    inp = fluid.dygraph.to_variable(inp)
                    out = linear(inp)
597
                    loss = paddle.mean(out)
598

599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615
                    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()
616
        if isinstance(current_lr, framework.Variable):
617 618 619 620
            return self._global_learning_rate().numpy()[0]

        if isinstance(self._learning_rate, float):
            return self._learning_rate
621 622 623
        elif isinstance(self._learning_rate, _LearningRateEpochDecay):
            step_lr = self._learning_rate()
            return step_lr.numpy()[0]
624 625 626 627 628 629 630
        else:
            step_lr = self._learning_rate.step()
            if isinstance(step_lr, (float, int)):
                return step_lr
            else:
                return step_lr.numpy()[0]

Y
yuyang18 已提交
631
    def _global_learning_rate(self, program=None):
Q
Qiao Longfei 已提交
632 633 634 635
        """
        get global decayed learning rate
        :return:
        """
636 637
        if program is None:
            program = framework.default_main_program()
Q
qiaolongfei 已提交
638
        return self._learning_rate_map.get(program, None)
Q
Qiao Longfei 已提交
639

Q
Qiao Longfei 已提交
640
    def _append_optimize_op(self, block, param_and_grad):
641
        """append optimize operator to block and return all the added optimize_op"""
Q
Qiao Longfei 已提交
642 643
        raise NotImplementedError()

644 645 646 647
    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 已提交
648 649
        if type(param_lr) == Variable:
            return param_lr
Q
qiaolongfei 已提交
650
        else:
W
Wu Yi 已提交
651
            if param_lr == 1.0:
Y
yuyang18 已提交
652
                return self._global_learning_rate()
W
Wu Yi 已提交
653
            else:
X
Xin Pan 已提交
654
                with default_main_program()._lr_schedule_guard(
655 656
                    is_with_opt=True
                ), framework.name_scope('scale_with_param_lr'):
657
                    return self._global_learning_rate() * param_lr
658 659 660 661 662 663 664

    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 已提交
665
        """
666 667
        pass

668
    def _finish_update(self, block, parameters_and_grads):
669 670 671 672 673 674 675 676
        """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 已提交
677
            None
678 679 680
        """
        pass

681 682 683 684 685 686 687 688 689 690
    def _add_accumulator(
        self,
        name,
        param,
        dtype=None,
        fill_value=0.0,
        shape=None,
        type=None,
        device=None,
    ):
691 692 693 694 695 696 697 698 699
        """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 已提交
700 701
        if self._name is not None:
            name = self._name + "_" + name
702 703 704 705
        if (
            name in self._accumulators
            and param.name in self._accumulators[name]
        ):
J
Jiabin Yang 已提交
706
            if framework._non_static_mode():
X
polish  
Xin Pan 已提交
707
                return self._accumulators[name][param.name]
708 709
            raise Exception(
                "Accumulator {} already exists for parameter {}".format(
710 711 712
                    name, param.name
                )
            )
713
        if shape is None:
714
            shape = param.shape
Q
Qiao Longfei 已提交
715
        assert isinstance(self.helper, LayerHelper)
716 717 718 719 720

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

Q
Qiao Longfei 已提交
721
        var = self.helper.create_global_variable(
722
            name=var_name,
Q
Qiao Longfei 已提交
723
            persistable=True,
F
fengjiayi 已提交
724
            dtype=dtype or param.dtype,
725
            type=core.VarDesc.VarType.LOD_TENSOR
726 727
            if framework._non_static_mode()
            else (param.type if type is None else type),
H
hong 已提交
728
            shape=shape,
729 730
            belong_to_optimizer=True,
        )
731 732 733 734
        if device is None:
            device = self._get_device_for_param(param.name)
        with device_guard(device):
            self.helper.set_variable_initializer(
735 736
                var, initializer=Constant(value=float(fill_value))
            )
H
hong 已提交
737

J
Jiabin Yang 已提交
738
        if framework._non_static_mode():
H
hong 已提交
739
            if len(self._accumulators_holder) > 0:
740 741 742 743 744
                assert (
                    var_name in self._accumulators_holder
                ), "Optimizer set error, {} should in state dict".format(
                    var_name
                )
H
hong 已提交
745 746
                var.set_value(self._accumulators_holder[var_name])

Q
Qiao Longfei 已提交
747
        self._accumulators[name][param.name] = var
748
        return var
749

750 751 752 753 754 755 756 757 758
    def _add_global_accumulator(
        self,
        name,
        dtype=None,
        fill_value=0.0,
        shape=None,
        type=None,
        device=None,
    ):
759 760 761 762 763 764 765 766 767 768 769 770 771
        """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
772
        if name in self._global_accumulators:
J
Jiabin Yang 已提交
773
            if framework._non_static_mode():
774 775
                return self._global_accumulators[name]
            raise Exception("Global accumulator {} already exists".format(name))
776
        if shape is None:
777 778 779 780 781 782 783 784 785 786 787 788 789
            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,
790 791
            belong_to_optimizer=True,
        )
792 793 794 795
        if device is None:
            device = 'cpu'
        with device_guard(device):
            self.helper.set_variable_initializer(
796 797
                var, initializer=Constant(value=float(fill_value))
            )
798

J
Jiabin Yang 已提交
799
        if framework._non_static_mode():
800
            if len(self._accumulators_holder) > 0:
801 802 803 804 805
                assert (
                    var_name in self._accumulators_holder
                ), "Optimizer set error, {} should in state dict".format(
                    var_name
                )
806 807 808 809 810
                var.set_value(self._accumulators_holder[var_name])

        self._global_accumulators[name] = var
        return var

811 812 813 814 815 816 817 818
    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:
819
            accumulator variable
820
        """
W
whs 已提交
821 822
        if self._name is not None:
            name = self._name + "_" + name
823 824 825 826
        if (
            name not in self._accumulators
            or param.name not in self._accumulators[name]
        ):
827 828
            raise Exception(
                "Accumulator {} does not exist for parameter {}".format(
829 830 831
                    name, param.name
                )
            )
832 833
        return self._accumulators[name][param.name]

834 835 836 837 838 839 840 841 842 843 844
    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
845
        if name not in self._global_accumulators:
846 847 848
            raise Exception("Global accumulator {} does not exist".format(name))
        return self._global_accumulators[name]

849 850 851 852 853
    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
854 855
                device_attr_name = (
                    core.op_proto_and_checker_maker.kOpDeviceAttrName()
856 857 858 859 860
                )
                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(
861 862
                            device_attr_name
                        )
863
                        break
864 865 866 867 868 869 870

    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

871
    def _create_optimization_pass(self, parameters_and_grads):
Q
Qiao Longfei 已提交
872 873 874
        """Add optimization operators to update gradients to variables.

        Args:
Q
qiaolongfei 已提交
875
          parameters_and_grads(list(tuple(Variable, Variable))):
876
            a list of (variable, gradient) pair to update.
Q
Qiao Longfei 已提交
877 878

        Returns:
879
          return_op_list: a list of operators that will complete one step of
880 881 882
            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 已提交
883
        """
884 885 886 887 888
        # 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
889
        # for parameters and extend _finish_update method to add custom ops.
890

891
        # Allways called under program_guard use global block as loss block
892 893 894
        # But if current block is in control flow, append optimize op in the
        # grad block of current block

895
        global_block = framework.default_main_program().global_block()
896 897 898
        target_block = global_block
        current_block = framework.default_main_program().current_block()
        if current_block.idx != global_block.idx:
899 900 901
            assert (
                current_block.backward_block_idx != -1
            ), "current block is not global_block, but it doesn't have backward block."
902
            target_block = framework.default_main_program().blocks[
903 904
                current_block.backward_block_idx
            ]
905 906

        start = len(target_block.ops)
907

908
        self._update_param_device_map(parameters_and_grads, target_block)
C
chengduo 已提交
909
        self._create_accumulators(
910 911
            target_block, [p[0] for p in parameters_and_grads if p[0].trainable]
        )
912 913
        self._create_global_learning_rate()

J
Jiabin Yang 已提交
914
        if framework._non_static_mode():
915 916 917
            for param_and_grad in parameters_and_grads:
                if param_and_grad[1] is None:
                    continue
918 919
                if param_and_grad[0].trainable is True:
                    self._append_optimize_op(target_block, param_and_grad)
920 921 922 923 924
        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(
925 926
                    param_and_grad
                ), name_scope("optimizer"):
927
                    if param_and_grad[0].trainable is True:
928
                        device = self._get_device_for_param(
929 930
                            param_and_grad[0].name
                        )
931 932
                        with device_guard(device):
                            optimize_op = self._append_optimize_op(
933 934
                                target_block, param_and_grad
                            )
935 936 937

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

940 941
        end = len(target_block.ops)
        return target_block._slice_ops(start, end)
942 943

    def _process_distribute_lookuptable(self, param_grads):
Q
Qiao Longfei 已提交
944 945 946 947 948 949 950 951 952
        """
        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
        """
953 954
        program = framework.default_main_program()
        global_block = framework.default_main_program().global_block()
Q
Qiao Longfei 已提交
955 956 957 958 959 960 961 962
        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(
963 964
                        "multi dist table var found, only support one now!"
                    )
Q
Qiao Longfei 已提交
965 966 967 968 969 970
                table_param = p
                table_grad = g
            else:
                new_param_grads.append((p, g))
        sgd_op = None
        if table_param is not None:
971
            param_and_grad = [table_param, table_grad]
972 973 974
            with table_param.block.program._optimized_guard(
                param_and_grad
            ), framework.name_scope("optimizer"):
975 976 977 978 979 980 981
                self._create_global_learning_rate()
                # create the optimize op
                sgd_op = global_block.append_op(
                    type='sgd',
                    inputs={
                        "Param": table_param,
                        "Grad": table_grad,
982
                        "LearningRate": self._create_param_lr(param_and_grad),
983
                    },
984 985
                    outputs={"ParamOut": param_and_grad[0]},
                )
Q
Qiao Longfei 已提交
986 987
        return new_param_grads, (table_param, table_grad), sgd_op

988 989 990 991 992 993 994 995
    def backward(
        self,
        loss,
        startup_program=None,
        parameter_list=None,
        no_grad_set=None,
        callbacks=None,
    ):
996
        """
997
        The first part of ``minimize``, do auto-diff to append backward operations for
998 999 1000
        the current program.

        Args:
1001 1002 1003 1004
            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 已提交
1005
            parameter_list (Iterable, optional): Iterable of ``Variable`` or ``Variable.name`` to update
1006 1007
                to minimize ``loss``. The default value is None, at this time all parameters
                will be updated.
1008
            no_grad_set (set, optional): Set of ``Variable``  or ``Variable.name`` that don't need
1009 1010 1011
                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 已提交
1012

1013
        Return:
1014 1015
            list: list of (param, grad) variable pairs, param is ``Parameter``,
                grad is the gradient value corresponding to the parameter.
M
minqiyang 已提交
1016

1017
        Examples:
1018
            See examples in ``apply_gradients``.
1019
        """
1020
        act_no_grad_set = None
J
Jiabin Yang 已提交
1021
        if framework._non_static_mode():
1022
            pass
L
Leo Chen 已提交
1023 1024
        else:
            act_no_grad_set = self._get_no_grad_set(loss, no_grad_set)
G
gongweibao 已提交
1025

L
Leo Chen 已提交
1026 1027 1028 1029
        # Infer dtype by loss if None
        if self._dtype is None:
            self._dtype = loss.dtype

J
Jiabin Yang 已提交
1030
        if framework._non_static_mode():
1031 1032 1033
            parameter_list = (
                parameter_list if parameter_list else self._parameter_list
            )
1034

C
chengduo 已提交
1035
            params_grads = []
1036
            for param in parameter_list:
C
chengduo 已提交
1037 1038
                if not param.trainable:
                    continue
1039
                if param._grad_ivar() is not None:
C
chengduo 已提交
1040
                    # create gradient variable
1041
                    grad_var = param._grad_ivar()
C
chengduo 已提交
1042
                    params_grads.append((param, grad_var))
1043
        else:
C
chengduo 已提交
1044 1045 1046
            if callbacks is None:
                callbacks = [error_clip_callback]
            else:
1047
                assert isinstance(callbacks, list)
C
chengduo 已提交
1048
            program = loss.block.program
1049 1050
            assert len(loss.shape) == 1 and loss.shape[0] == 1, (
                "The loss.shape should be (1L,), but the current loss.shape is {}. "
1051
                "Maybe that you should call paddle.mean to process the current loss.".format(
1052 1053 1054 1055 1056 1057
                    loss.shape
                )
            )
            parameter_list = (
                parameter_list if parameter_list else self._parameter_list
            )
C
chengduo 已提交
1058
            with program_guard(program, startup_program):
1059 1060 1061
                params_grads = append_backward(
                    loss, parameter_list, act_no_grad_set, callbacks
                )
C
chengduo 已提交
1062
        return params_grads
1063

1064
    def _create_regularization_of_grad(self, param, grad, regularization=None):
1065
        """Create and add backward regularization Operators
1066

1067 1068 1069
        Function helper of append_regularization_ops.
        """
        # If no gradient or no regularization is specified,  then we don't need to do anything
1070
        if grad is None or (
1071 1072 1073 1074 1075 1076
            (
                not hasattr(param, 'regularizer')
                or (hasattr(param, 'regularizer') and param.regularizer is None)
            )
            and regularization is None
        ):
1077 1078 1079 1080 1081 1082 1083 1084 1085 1086
            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 已提交
1087
        if framework._non_static_mode():
1088
            return _legacy_C_ops.sum([grad, regularization_term])
1089

1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100
        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,
1101 1102
                type=core.VarDesc.VarType.LOD_TENSOR,
            )
1103 1104 1105

        inputs = {"X": [grad, regularization_term]}
        outputs = {"Out": [new_grad]}
1106
        grad.block.append_op(type='sum', inputs=inputs, outputs=outputs)
1107 1108 1109

        return new_grad

1110 1111 1112
    def append_regularization_ops(
        self, parameters_and_grads, regularization=None
    ):
1113
        r"""Create and add backward regularization Operators
1114

1115 1116 1117 1118
        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.
1119

1120 1121 1122 1123 1124
        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.
1125

1126 1127 1128
        Returns:
            list[(Variable, Variable)]: list of (parameters, gradients) \
            pair with the regularized gradient
1129

1130 1131 1132 1133
        Raises:
            Exception: Unknown regularization type
        """
        params_and_grads = []
J
Jiabin Yang 已提交
1134
        if framework._non_static_mode():
1135
            for param, grad in parameters_and_grads:
1136
                new_grad = self._create_regularization_of_grad(
1137 1138
                    param, grad, regularization
                )
1139 1140 1141 1142 1143
                params_and_grads.append((param, new_grad))
        else:
            repeate_regularizer = False
            with framework.name_scope('regularization'):
                for param, grad in parameters_and_grads:
1144 1145 1146 1147 1148
                    if (
                        not repeate_regularizer
                        and getattr(param, 'regularizer', None) is not None
                        and regularization is not None
                    ):
1149 1150 1151 1152
                        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!"
1153 1154
                            % regularization.__str__()
                        )
1155 1156
                    with param.block.program._optimized_guard([param, grad]):
                        new_grad = self._create_regularization_of_grad(
1157 1158
                            param, grad, regularization
                        )
1159 1160 1161
                        params_and_grads.append((param, new_grad))
        return params_and_grads

1162 1163 1164 1165 1166 1167 1168
    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
1169 1170 1171 1172
            if (
                getattr(p, 'need_clip', True) is False
                or getattr(p, 'regularizer', None) is not None
            ):
1173 1174
                warnings.warn(
                    "flatten_param_grads=True will be discarded since paramter '{}''s need_clip is False or "
1175 1176
                    "the regularizer is set".format(p.name)
                )
1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190
                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)],
1191 1192
            belong_to_optimizer=True,
        )
1193 1194 1195 1196 1197 1198 1199 1200 1201 1202

        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)],
1203 1204
            belong_to_optimizer=True,
        )
1205 1206

        with program_guard(default_main_program()):
1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235
            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,
                },
            )

            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,
                },
            )
1236

1237
        # NOTE(zhiqiu): the initializer should be set after coalesce_tensor op,
1238
        # so the shape of flatten_param and flatten_grad will be inferred.
1239 1240 1241 1242 1243 1244
        self.helper.set_variable_initializer(
            flatten_param, initializer=Constant(0.0)
        )
        self.helper.set_variable_initializer(
            flatten_grad, initializer=Constant(0.0)
        )
1245 1246 1247

        return [(flatten_param, flatten_grad)]

1248 1249 1250 1251 1252 1253 1254
    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 已提交
1255

1256 1257
        Returns:
            list: A list of operators appended to the current program.
M
minqiyang 已提交
1258

1259 1260 1261
        Examples:
            .. code-block:: python

1262
                import paddle.fluid as fluid
1263 1264 1265 1266 1267 1268 1269 1270 1271
                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)

1272 1273
        # NOTE(zhiqiu): currently, only support ClipGradByGlobalNorm and without regularization.
        if self._flatten_param_grads and self.regularization is None:
1274
            if self._grad_clip is None or isinstance(
1275 1276
                self._grad_clip, ClipGradByGlobalNorm
            ):
1277 1278
                params_grads = self.flatten_param_grads(params_grads)

1279
        # 'optimizer(grad_clip)' or 'set_gradient_clip'
1280 1281 1282 1283
        if self._grad_clip is not None:
            params_grads = self._grad_clip(params_grads)
        else:
            params_grads = append_gradient_clip_ops(params_grads)
1284 1285

        # Add regularization if any
1286 1287 1288
        params_grads = self.append_regularization_ops(
            params_grads, self.regularization
        )
1289 1290 1291 1292

        optimize_ops = self._create_optimization_pass(params_grads)
        return optimize_ops

C
chengduo 已提交
1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304
    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 已提交
1305
        if framework._non_static_mode():
1306 1307 1308 1309
            with program_guard(
                framework.default_main_program(),
                framework.default_startup_program(),
            ):
1310 1311
                if self._grad_clip is not None:
                    params_grads = self._grad_clip(params_grads)
1312
                params_grads = self.append_regularization_ops(
1313 1314
                    params_grads, self.regularization
                )
C
chengduo 已提交
1315 1316 1317 1318 1319 1320 1321
                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 已提交
1322
    def _get_no_grad_set(self, loss, no_grad_set=None):
1323
        no_grad_set = _get_no_grad_set_name(no_grad_set)
G
gongweibao 已提交
1324 1325
        parameters = loss.block.program.global_block().all_parameters()
        param_no_trainable = set(
1326 1327
            [param.name for param in parameters if param.trainable is False]
        )
G
gongweibao 已提交
1328 1329 1330 1331 1332
        # If the parameter is no trainable, it should not have a gradient.
        no_grad_set.update(param_no_trainable)

        return no_grad_set

1333 1334 1335 1336
    @framework.dygraph_only
    def clear_gradients(self):
        """
        Clear the gradients of all optimized parameters for model.
1337 1338

        If not, new gradient will accumulat on previous gradient.
1339

1340 1341
        Returns:
            None
1342

1343 1344 1345 1346
        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
1347
                import paddle
1348 1349 1350 1351 1352
                import numpy as np

                with fluid.dygraph.guard():
                    value = np.arange(26).reshape(2, 13).astype("float32")
                    a = fluid.dygraph.to_variable(value)
1353
                    linear = paddle.nn.Linear(13, 5)
1354
                    # This can be any optimizer supported by dygraph.
1355
                    adam = fluid.optimizer.Adam(learning_rate = 0.01,
1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366
                                                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()

1367
    @imperative_base.no_grad
1368 1369 1370
    def minimize(
        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
    ):
1371
        """
1372
        Add operations to minimize ``loss`` by updating ``parameter_list``.
M
minqiyang 已提交
1373

1374
        Args:
1375 1376 1377 1378
            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 已提交
1379
            parameter_list (Iterable, optional): Iterable of ``Variable`` or ``Variable.name`` to update
1380 1381
                to minimize ``loss``. The default value is None, at this time all parameters
                will be updated.
1382
            no_grad_set (set, optional): Set of ``Variable``  or ``Variable.name`` that don't need
1383
                to be updated. The default value is None.
Q
Qiao Longfei 已提交
1384

1385
        Returns:
1386 1387 1388
            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.
1389 1390
            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
1391
            ``fetch_list`` before run, see details in ``Executor``.
1392 1393 1394

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

1398 1399 1400
        parameter_list = (
            parameter_list if parameter_list else self._parameter_list
        )
1401

1402 1403 1404 1405 1406 1407
        params_grads = self.backward(
            loss,
            startup_program=startup_program,
            parameter_list=parameter_list,
            no_grad_set=no_grad_set,
        )
1408

1409 1410 1411
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
M
minqiyang 已提交
1412

Q
Qiao Longfei 已提交
1413
        return optimize_ops, params_grads
Q
Qiao Longfei 已提交
1414 1415 1416


class SGDOptimizer(Optimizer):
1417
    r"""
Q
qiaolongfei 已提交
1418 1419 1420 1421 1422 1423
    Optimizer of the stochastic gradient descent algorithm.

    .. math::

        param\_out = param - learning\_rate * grad

1424 1425 1426
    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 已提交
1427
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1428 1429
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
1430 1431 1432 1433 1434
        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.
1435 1436 1437
        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` ,
1438
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
1439 1440
        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 已提交
1441 1442 1443 1444

    Examples:
        .. code-block:: python

1445 1446 1447 1448
            import paddle
            import paddle.fluid as fluid
            import numpy as np

1449
            paddle.enable_static()
1450 1451 1452 1453 1454 1455
            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)
1456
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
1457
                avg_cost = paddle.mean(cost)
1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470

                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 已提交
1471 1472
    """

1473 1474 1475 1476 1477 1478 1479 1480 1481
    def __init__(
        self,
        learning_rate,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        multi_precision=False,
        name=None,
    ):
Q
Qiao Longfei 已提交
1482
        assert learning_rate is not None
1483
        super().__init__(
1484 1485 1486 1487 1488 1489
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
Q
Qiao Longfei 已提交
1490
        self.type = "sgd"
1491
        self._use_mkldnn = False
1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502
        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)
1503
            var = paddle.static.create_global_var(
1504 1505 1506 1507 1508 1509
                name=var_name,
                shape=param.shape,
                value=0,
                dtype='float32',
                persistable=True,
            )
1510
            block = self.helper.startup_program.global_block()
1511 1512 1513 1514 1515 1516 1517 1518 1519
            block.append_op(
                type="cast",
                inputs={"X": [param]},
                outputs={"Out": [var]},
                attrs={
                    "in_dtype": param.dtype,
                    "out_dtype": core.VarDesc.VarType.FP32,
                },
            )
1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532
            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
1533 1534 1535 1536
            if (
                p.dtype == core.VarDesc.VarType.FP16
                and not self._multi_precision
            ):
1537 1538 1539 1540
                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 已提交
1541

1542
    @no_grad
1543
    def _append_optimize_op(self, block, param_and_grad):
1544

1545 1546 1547 1548 1549 1550 1551 1552 1553
        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
        )
1554

1555
        lr = self._create_param_lr(param_and_grad)
Z
zyfncg 已提交
1556
        if in_dygraph_mode():
1557 1558 1559 1560 1561 1562 1563
            _C_ops.sgd_(
                param_and_grad[0],
                lr,
                param_and_grad[1],
                master_weight,
                find_master,
            )
Z
zyfncg 已提交
1564 1565
            return None
        if _in_legacy_dygraph():
1566 1567 1568 1569 1570 1571 1572 1573
            _legacy_C_ops.sgd(
                param_and_grad[0],
                lr,
                param_and_grad[1],
                master_weight,
                param_and_grad[0],
                master_weight,
            )
1574
            return None
1575

1576
        assert isinstance(block, framework.Block)
Q
Qiao Longfei 已提交
1577
        # create the optimize op
1578 1579 1580
        inputs = {
            "Param": param_and_grad[0],
            "Grad": param_and_grad[1],
1581
            "LearningRate": lr,
1582 1583 1584 1585 1586 1587 1588 1589 1590 1591
        }

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

        attrs = {"multi_precision": find_master}

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

1592 1593 1594 1595 1596 1597 1598
        sgd_op = block.append_op(
            type=self.type,
            inputs=inputs,
            outputs=outputs,
            attrs=attrs,
            stop_gradient=True,
        )
Q
Qiao Longfei 已提交
1599 1600

        return sgd_op
1601 1602 1603


class MomentumOptimizer(Optimizer):
1604
    r"""
Q
qiaolongfei 已提交
1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617

    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):

1618
        &\quad   param = param - (gradient + mu * velocity) * learning\_rate
Q
qiaolongfei 已提交
1619 1620 1621

        & else:

Q
qiaolongfei 已提交
1622
        &\quad   param = param - learning\_rate * velocity
Q
qiaolongfei 已提交
1623

1624 1625 1626 1627
    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 已提交
1628
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1629 1630
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
1631
        use_nesterov (bool, optional): Enables Nesterov momentum, default is false.
1632 1633 1634 1635 1636
        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.
1637 1638 1639
        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` ,
1640
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
1641 1642
        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 已提交
1643 1644 1645 1646

    Examples:
        .. code-block:: python

1647 1648 1649 1650
            import paddle
            import paddle.fluid as fluid
            import numpy as np

1651
            paddle.enable_static()
1652 1653 1654 1655 1656 1657
            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)
1658
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
1659
                avg_cost = paddle.mean(cost)
1660 1661 1662 1663 1664 1665 1666 1667 1668 1669 1670 1671 1672

                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)

1673 1674 1675
    """
    _velocity_acc_str = "velocity"

1676 1677 1678 1679 1680 1681 1682 1683 1684 1685
    def __init__(
        self,
        learning_rate,
        momentum,
        parameter_list=None,
        use_nesterov=False,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
1686 1687
        assert learning_rate is not None
        assert momentum is not None
1688
        super().__init__(
1689 1690 1691 1692 1693 1694
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
1695 1696
        self.type = "momentum"
        self._momentum = momentum
1697
        self._use_nesterov = bool(use_nesterov)
1698 1699 1700 1701 1702

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

        for p in parameters:
Q
Qiao Longfei 已提交
1703
            self._add_accumulator(self._velocity_acc_str, p)
1704 1705 1706 1707

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

1708 1709 1710
        velocity_acc = self._get_accumulator(
            self._velocity_acc_str, param_and_grad[0]
        )
1711
        lr = self._create_param_lr(param_and_grad)
1712
        master_weight = None
J
Jiabin Yang 已提交
1713
        if framework._non_static_mode():
1714
            _, _, _ = _legacy_C_ops.momentum(
1715 1716 1717 1718 1719 1720 1721 1722 1723 1724 1725 1726 1727
                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,
            )
1728
            return None
1729

1730
        attrs = {"mu": self._momentum, "use_nesterov": self._use_nesterov}
1731 1732 1733 1734
        inputs = {
            "Param": [param_and_grad[0]],
            "Grad": [param_and_grad[1]],
            "Velocity": [velocity_acc],
1735
            "LearningRate": [lr],
1736 1737 1738 1739
        }

        outputs = {
            "ParamOut": [param_and_grad[0]],
1740
            "VelocityOut": [velocity_acc],
1741
        }
1742
        # create the momentum optimize op
1743 1744 1745 1746 1747 1748 1749
        momentum_op = block.append_op(
            type=self.type,
            inputs=inputs,
            outputs=outputs,
            attrs=attrs,
            stop_gradient=True,
        )
1750 1751

        return momentum_op
1752 1753


1754
class LarsMomentumOptimizer(Optimizer):
1755
    r"""
1756 1757 1758 1759 1760 1761 1762 1763 1764
    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||}

1765
        & velocity = mu * velocity + local\_learning\_rate * (gradient + lars\_weight\_decay * param + epsilon)
1766 1767 1768

        & param = param - velocity

1769 1770 1771 1772 1773 1774
    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 已提交
1775
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1776 1777
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
1778 1779 1780 1781 1782
        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.
1783 1784 1785
        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` ,
1786
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
1787 1788
        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.
1789 1790
        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.
1791 1792 1793
        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`.
1794

1795 1796 1797
    Examples:
        .. code-block:: python

1798
            import paddle
1799 1800 1801
            import paddle.fluid as fluid
            import numpy as np

1802
            paddle.enable_static()
1803 1804 1805 1806
            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)
1807
            out = paddle.sum(out)
1808 1809 1810 1811 1812 1813 1814 1815
            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])
1816 1817 1818
    """
    _velocity_acc_str = "velocity"

1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832 1833
    def __init__(
        self,
        learning_rate,
        momentum,
        lars_coeff=0.001,
        lars_weight_decay=0.0005,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
        exclude_from_weight_decay=None,
        epsilon=0,
        multi_precision=False,
        rescale_grad=1.0,
    ):
1834 1835
        assert learning_rate is not None
        assert momentum is not None
1836
        super().__init__(
1837 1838 1839 1840 1841 1842
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
1843 1844 1845 1846
        self.type = "lars_momentum"
        self._momentum = momentum
        self._lars_coeff = float(lars_coeff)
        self._lars_weight_decay = float(lars_weight_decay)
1847 1848 1849 1850 1851
        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
1852 1853 1854 1855 1856
        self._multi_precision = multi_precision
        self._rescale_grad = float(rescale_grad)
        self._master_weights = {}

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

1862 1863
            var_name = param.name + '_fp32_master'
            var_name = unique_name.generate(var_name)
1864
            var = paddle.static.create_global_var(
1865 1866 1867 1868 1869 1870
                name=var_name,
                shape=param.shape,
                value=0,
                dtype='float32',
                persistable=True,
            )
1871
            block = self.helper.startup_program.global_block()
1872 1873 1874 1875 1876 1877 1878 1879 1880
            block.append_op(
                type="cast",
                inputs={"X": [param]},
                outputs={"Out": [var]},
                attrs={
                    "in_dtype": param.dtype,
                    "out_dtype": core.VarDesc.VarType.FP32,
                },
            )
1881
            self._master_weights[param.name] = var
1882 1883 1884 1885 1886 1887 1888 1889 1890 1891 1892 1893
        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
1894 1895 1896 1897 1898 1899
        find_master = (
            self._multi_precision and param.dtype == core.VarDesc.VarType.FP16
        )
        target_param = (
            self._master_weights[param.name] if find_master else param
        )
1900
        target_name = target_param.name
1901 1902 1903 1904
        if (
            name not in self._accumulators
            or target_name not in self._accumulators[name]
        ):
1905 1906
            raise Exception(
                "Accumulator {} does not exist for parameter {}".format(
1907 1908 1909
                    name, target_name
                )
            )
1910
        return self._accumulators[name][target_name]
1911 1912 1913 1914 1915

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

        for p in parameters:
1916 1917 1918 1919
            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
1920 1921 1922 1923
            if (
                p.dtype == core.VarDesc.VarType.FP16
                and not self._multi_precision
            ):
1924 1925 1926 1927
                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."
                )
1928 1929 1930 1931
            self._add_accumulator(self._velocity_acc_str, p)

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)
1932 1933 1934 1935 1936 1937 1938 1939
        _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

1940 1941 1942
        velocity_acc = self._get_accumulator(
            self._velocity_acc_str, param_and_grad[0]
        )
1943 1944
        lr = self._create_param_lr(param_and_grad)

1945 1946 1947 1948 1949 1950 1951 1952 1953
        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
        )
1954 1955 1956

        attrs = {
            "mu": self._momentum,
1957
            "lars_coeff": self._lars_coeff,
L
limingshu 已提交
1958
            "lars_weight_decay": [_lars_weight_decay],
1959
            "multi_precision": find_master,
L
limingshu 已提交
1960
            "epsilon": self._epsilon,
1961
            "rescale_grad": self._rescale_grad,
1962 1963 1964 1965 1966 1967
        }

        inputs = {
            "Param": param_and_grad[0],
            "Grad": param_and_grad[1],
            "Velocity": velocity_acc,
1968
            "LearningRate": lr,
1969 1970 1971 1972 1973 1974 1975 1976
        }

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

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

J
Jiabin Yang 已提交
1977
        if framework._non_static_mode():
1978
            tmp, tmp2 = _legacy_C_ops.lars_momentum(
1979 1980 1981 1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995 1996 1997
                [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,
            )
1998 1999
        else:
            # create the momentum optimize op
2000 2001 2002 2003 2004 2005 2006
            momentum_op = block.append_op(
                type=self.type,
                inputs=inputs,
                outputs=outputs,
                attrs=attrs,
                stop_gradient=True,
            )
2007

2008
            return momentum_op
2009 2010


2011
class AdagradOptimizer(Optimizer):
2012
    r"""
2013 2014
    The Adaptive Gradient optimizer (Adagrad for short) can adaptively assign
    different learning rates to individual parameters.
Q
qiaolongfei 已提交
2015

2016
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2017 2018 2019 2020 2021 2022 2023

    .. math::

        moment\_out &= moment + grad * grad

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

2024 2025 2026 2027 2028 2029
    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 已提交
2030 2031 2032
    for numerical stability to avoid the division by zero error.

    Args:
2033 2034 2035 2036
        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 已提交
2037
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2038 2039
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2040 2041 2042 2043 2044
        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.
2045 2046 2047
        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` ,
2048
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2049 2050 2051 2052 2053
        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 已提交
2054 2055 2056 2057

    Examples:
        .. code-block:: python

2058
            import paddle
2059
            import numpy as np
2060
            import paddle.fluid as fluid
2061

2062
            paddle.enable_static()
2063
            np_inp = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
2064
            inp = fluid.data(name="inp", shape=[2, 2])
2065
            out = fluid.layers.fc(inp, size=3)
2066
            out = paddle.sum(out)
2067
            optimizer = fluid.optimizer.AdagradOptimizer(learning_rate=0.2)
2068 2069 2070 2071 2072 2073 2074
            optimizer.minimize(out)

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

2078 2079 2080 2081 2082 2083 2084 2085 2086 2087
    def __init__(
        self,
        learning_rate,
        epsilon=1.0e-6,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
        initial_accumulator_value=0.0,
    ):
2088 2089
        assert learning_rate is not None
        assert epsilon is not None
2090
        super().__init__(
2091 2092 2093 2094 2095 2096
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
2097 2098
        self.type = "adagrad"
        self._epsilon = epsilon
2099
        self.initial_accumulator_value = initial_accumulator_value
2100 2101 2102 2103 2104

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

        for p in parameters:
2105 2106 2107 2108 2109
            self._add_accumulator(
                self._moment_acc_str,
                p,
                fill_value=self.initial_accumulator_value,
            )
2110 2111 2112 2113

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

2114 2115 2116
        moment_acc = self._get_accumulator(
            self._moment_acc_str, param_and_grad[0]
        )
C
caozhou 已提交
2117
        if in_dygraph_mode():
2118 2119 2120 2121 2122 2123 2124
            _C_ops.adagrad_(
                param_and_grad[0],
                param_and_grad[1],
                moment_acc,
                self._create_param_lr(param_and_grad),
                self._epsilon,
            )
C
caozhou 已提交
2125 2126
            return None
        elif _in_legacy_dygraph():
2127 2128 2129 2130 2131 2132 2133 2134 2135 2136
            _legacy_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 已提交
2137
            return None
2138 2139 2140 2141 2142 2143 2144 2145
        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,
2146
                    "LearningRate": self._create_param_lr(param_and_grad),
2147 2148 2149
                },
                outputs={
                    "ParamOut": param_and_grad[0],
2150
                    "MomentOut": moment_acc,
2151 2152
                },
                attrs={"epsilon": self._epsilon},
2153 2154
                stop_gradient=True,
            )
2155

2156
            return adagrad_op
2157 2158 2159


class AdamOptimizer(Optimizer):
2160
    r"""
T
tianshuo78520a 已提交
2161
    The Adam optimizer uses an optimization described at the end
2162 2163 2164
    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.
2165

2166
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2167 2168 2169 2170 2171 2172 2173 2174 2175 2176 2177 2178 2179 2180

    .. 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}

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

Q
qiaolongfei 已提交
2183
    Args:
2184 2185
        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.
2186 2187
        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.
2188
            The default value is 0.9.
2189 2190
        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.
2191
            The default value is 0.999.
2192 2193
        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.
2194
            The default value is 1e-08.
H
hong 已提交
2195
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2196 2197
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2198 2199 2200 2201 2202
        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.
2203 2204 2205
        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` ,
2206
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2207 2208 2209 2210 2211 2212 2213 2214 2215 2216
        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.
2217
        use_global_beta_pow (bool, optional): Whether to use global beta_pow. If true, Adam will use global beta_pow
2218
            for whole model instead of creating beta_pow for each parameter. Default is false.
2219 2220
        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
2221
            use same align_size as allocator.
Q
qiaolongfei 已提交
2222 2223 2224 2225

    Examples:
        .. code-block:: python

2226 2227 2228
            import paddle
            import paddle.fluid as fluid

2229
            paddle.enable_static()
2230 2231 2232
            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
2233 2234
                x = fluid.data(name='x', shape=[None, 13], dtype='float32')
                y = fluid.data(name='y', shape=[None, 1], dtype='float32')
2235
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
2236
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
2237
                avg_cost = paddle.mean(cost)
2238 2239 2240 2241 2242 2243 2244 2245 2246 2247 2248 2249

                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 已提交
2250

2251 2252 2253 2254 2255 2256 2257
        .. 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

2258
            paddle.enable_static()
2259 2260 2261 2262 2263 2264
            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)
2265
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
2266
                avg_cost = paddle.mean(cost)
2267 2268

                # define beta decay variable
2269
                def get_decayed_betas(beta1_init, beta2_init, decay_steps, decay_rate, epsilon_init):
2270 2271
                    global_step = lr_scheduler._decay_step_counter()

2272
                    beta1 = paddle.static.create_global_var(
2273 2274 2275 2276 2277 2278
                        shape=[1],
                        value=float(beta1_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="beta1")
2279
                    beta2 = paddle.static.create_global_var(
2280 2281 2282 2283 2284 2285
                        shape=[1],
                        value=float(beta2_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="beta2")
2286
                    epsilon = paddle.static.create_global_var(
2287 2288 2289 2290 2291 2292
                        shape=[1],
                        value=float(epsilon_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="epsilon")
2293 2294 2295 2296 2297 2298 2299

                    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)

2300
                    return beta1, beta2, epsilon
2301

2302
                beta1, beta2, epsilon = get_decayed_betas(0.9, 0.99, 1e5, 0.9, 1e-8)
2303 2304
                adam_optimizer = fluid.optimizer.AdamOptimizer(
                                                    learning_rate=0.01,
2305
                                                    beta1=beta1,
2306 2307
                                                    beta2=beta2,
                                                    epsilon=epsilon)
2308 2309 2310 2311 2312 2313 2314 2315 2316 2317
                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)
2318 2319 2320
    """
    _moment1_acc_str = "moment1"
    _moment2_acc_str = "moment2"
Q
qiaolongfei 已提交
2321 2322
    _beta1_pow_acc_str = "beta1_pow_acc"
    _beta2_pow_acc_str = "beta2_pow_acc"
2323

2324 2325 2326 2327 2328 2329 2330 2331 2332 2333 2334 2335 2336 2337 2338
    def __init__(
        self,
        learning_rate=0.001,
        beta1=0.9,
        beta2=0.999,
        epsilon=1e-8,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
        lazy_mode=False,
        use_global_beta_pow=False,
        flatten_param_grads=False,
        align_size=-1,
    ):
2339 2340 2341 2342
        assert learning_rate is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
2343
        super().__init__(
2344 2345 2346 2347 2348 2349 2350 2351
            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,
        )
2352 2353 2354 2355
        self.type = "adam"
        self._beta1 = beta1
        self._beta2 = beta2
        self._epsilon = epsilon
Q
Qiao Longfei 已提交
2356
        self._lazy_mode = lazy_mode
2357
        self._use_global_beta_pow = use_global_beta_pow
2358 2359 2360 2361 2362 2363

    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 已提交
2364 2365
            self._add_accumulator(self._moment1_acc_str, p)
            self._add_accumulator(self._moment2_acc_str, p)
2366 2367 2368 2369
            if not self._use_global_beta_pow:
                self._add_accumulator(
                    name=self._beta1_pow_acc_str,
                    param=p,
2370 2371 2372
                    fill_value=0.9
                    if isinstance(self._beta1, Variable)
                    else self._beta1,
2373
                    shape=[1],
2374 2375 2376
                    type=core.VarDesc.VarType.LOD_TENSOR,
                    device='cpu',
                )
2377 2378 2379
                self._add_accumulator(
                    name=self._beta2_pow_acc_str,
                    param=p,
2380 2381 2382
                    fill_value=0.999
                    if isinstance(self._beta2, Variable)
                    else self._beta2,
2383
                    shape=[1],
2384 2385 2386
                    type=core.VarDesc.VarType.LOD_TENSOR,
                    device='cpu',
                )
2387 2388
        if self._use_global_beta_pow:
            self._add_global_accumulator(
Q
qiaolongfei 已提交
2389
                name=self._beta1_pow_acc_str,
2390 2391 2392
                fill_value=0.9
                if isinstance(self._beta1, Variable)
                else self._beta1,
2393
                shape=[1],
2394 2395 2396
                type=core.VarDesc.VarType.LOD_TENSOR,
                device='cpu',
            )
2397
            self._add_global_accumulator(
Q
qiaolongfei 已提交
2398
                name=self._beta2_pow_acc_str,
2399 2400 2401
                fill_value=0.999
                if isinstance(self._beta2, Variable)
                else self._beta2,
2402
                shape=[1],
2403 2404 2405
                type=core.VarDesc.VarType.LOD_TENSOR,
                device='cpu',
            )
2406 2407 2408 2409

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

2410 2411 2412 2413 2414 2415
        moment1 = self._get_accumulator(
            self._moment1_acc_str, param_and_grad[0]
        )
        moment2 = self._get_accumulator(
            self._moment2_acc_str, param_and_grad[0]
        )
2416 2417
        if self._use_global_beta_pow:
            beta1_pow_acc = self._get_global_accumulator(
2418 2419
                self._beta1_pow_acc_str
            )
2420
            beta2_pow_acc = self._get_global_accumulator(
2421 2422
                self._beta2_pow_acc_str
            )
2423
        else:
2424 2425 2426 2427 2428 2429
            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]
            )
2430
        lr = self._create_param_lr(param_and_grad)
2431
        # create the adam optimize op
2432

J
Jiabin Yang 已提交
2433
        if framework._non_static_mode():
2434 2435 2436 2437 2438 2439 2440 2441 2442 2443
            _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)
            )
2444
            master_weight = None
2445
            _, _, _, _, _, _ = _legacy_C_ops.adam(
2446 2447 2448 2449 2450 2451 2452 2453 2454 2455 2456 2457 2458 2459 2460 2461 2462 2463 2464 2465 2466 2467 2468 2469 2470 2471 2472
                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,
                'epsilon',
                self._epsilon,
                'lazy_mode',
                self._lazy_mode,
                'min_row_size_to_use_multithread',
                1000,
                'beta1',
                _beta1,
                'beta2',
                _beta2,
                'use_global_beta_pow',
                self._use_global_beta_pow,
            )
2473 2474 2475

            return None

2476
        inputs = {
2477 2478
            "Param": [param_and_grad[0]],
            "Grad": [param_and_grad[1]],
2479
            "LearningRate": [lr],
2480 2481 2482
            "Moment1": [moment1],
            "Moment2": [moment2],
            "Beta1Pow": [beta1_pow_acc],
2483
            "Beta2Pow": [beta2_pow_acc],
2484
        }
2485 2486 2487 2488 2489 2490 2491

        # 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

2492
        outputs = {
2493 2494 2495 2496 2497
            "ParamOut": [param_and_grad[0]],
            "Moment1Out": [moment1],
            "Moment2Out": [moment2],
            "Beta1PowOut": [beta1_pow_acc],
            "Beta2PowOut": [beta2_pow_acc],
2498 2499 2500
        }
        attrs = {
            "lazy_mode": self._lazy_mode,
2501
            "min_row_size_to_use_multithread": 1000,
2502
            'use_global_beta_pow': self._use_global_beta_pow,
2503 2504 2505 2506 2507 2508 2509 2510 2511 2512
        }

        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
2513 2514 2515 2516
        if isinstance(self._epsilon, Variable):
            inputs['EpsilonTensor'] = self._epsilon
        else:
            attrs['epsilon'] = self._epsilon
2517

2518 2519 2520 2521 2522 2523 2524
        adam_op = block.append_op(
            type=self.type,
            inputs=inputs,
            outputs=outputs,
            attrs=attrs,
            stop_gradient=True,
        )
2525 2526 2527

        return adam_op

2528
    def _finish_update(self, block, parameters_and_grads):
2529
        r"""Update beta1_pow and beta2_pow accumulator"""
2530 2531 2532
        assert isinstance(block, framework.Block)
        if self._use_global_beta_pow:
            beta1_pow_acc = self._get_global_accumulator(
2533 2534
                self._beta1_pow_acc_str
            )
2535
            beta2_pow_acc = self._get_global_accumulator(
2536 2537
                self._beta2_pow_acc_str
            )
2538 2539 2540

            with block.program._optimized_guard([]):
                inputs = {"X": beta1_pow_acc}
2541
                outputs = {"Out": beta1_pow_acc}
2542 2543
                attrs = {}
                if isinstance(self._beta1, Variable):
2544 2545
                    inputs["Y"] = self._beta1
                    # use elementwise_mul for better performance
2546 2547 2548 2549 2550 2551 2552
                    block.append_op(
                        type="elementwise_mul",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
2553 2554
                else:
                    attrs['scale'] = self._beta1
2555 2556 2557 2558 2559 2560 2561
                    block.append_op(
                        type="scale",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
2562 2563

                inputs = {"X": beta2_pow_acc}
2564
                outputs = {"Out": beta2_pow_acc}
2565 2566
                attrs = {}
                if isinstance(self._beta2, Variable):
2567 2568
                    inputs["Y"] = self._beta2
                    # use elementwise_mul for better performance
2569 2570 2571 2572 2573 2574 2575
                    block.append_op(
                        type="elementwise_mul",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
2576 2577
                else:
                    attrs['scale'] = self._beta2
2578 2579 2580 2581 2582 2583 2584
                    block.append_op(
                        type="scale",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
2585

2586 2587

class AdamaxOptimizer(Optimizer):
2588
    r"""
2589
    The Adamax optimizer is implemented based on the Adamax Optimization
2590 2591 2592
    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 已提交
2593

2594
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2595 2596 2597 2598 2599 2600 2601 2602 2603 2604 2605 2606 2607

    .. 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}

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

2610 2611 2612 2613 2614 2615 2616 2617 2618 2619 2620 2621
    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 已提交
2622
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2623 2624
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2625 2626 2627 2628 2629
        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.
2630 2631 2632
        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` ,
2633
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2634 2635 2636 2637 2638 2639
        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 已提交
2640

2641 2642 2643 2644 2645
    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          import numpy
2
201716010711 已提交
2646 2647
          import paddle
          paddle.enable_static()
2648 2649 2650 2651 2652 2653 2654 2655

          # 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):
2656
              data = fluid.data(name='X', shape=[None, 1], dtype='float32')
2657
              hidden = fluid.layers.fc(input=data, size=10)
2
201716010711 已提交
2658
              loss = paddle.mean(hidden)
2659
              adam = fluid.optimizer.AdamaxOptimizer(learning_rate=0.2)
2660 2661 2662 2663 2664 2665 2666 2667 2668
              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])
2669 2670 2671
    """
    _moment_acc_str = "moment"
    _inf_norm_acc_str = "inf_norm"
Q
qiaolongfei 已提交
2672
    _beta1_pow_acc_str = "beta1_pow_acc"
2673

2674 2675 2676 2677 2678 2679 2680 2681 2682 2683 2684
    def __init__(
        self,
        learning_rate=0.001,
        beta1=0.9,
        beta2=0.999,
        epsilon=1e-8,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
2685 2686 2687 2688
        assert learning_rate is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
2689
        super().__init__(
2690 2691 2692 2693 2694 2695
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
2696 2697 2698 2699 2700 2701 2702 2703
        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 已提交
2704 2705
            self._add_accumulator(self._moment_acc_str, p)
            self._add_accumulator(self._inf_norm_acc_str, p)
2706 2707 2708 2709 2710 2711
            self._add_accumulator(
                name=self._beta1_pow_acc_str,
                param=p,
                fill_value=self._beta1,
                shape=[1],
            )
2712 2713 2714 2715 2716

    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])
2717 2718 2719 2720 2721 2722
        inf_norm = self._get_accumulator(
            self._inf_norm_acc_str, param_and_grad[0]
        )
        beta1_pow_acc = self._get_accumulator(
            self._beta1_pow_acc_str, param_and_grad[0]
        )
2723 2724

        if framework.in_dygraph_mode():
2725 2726 2727 2728 2729 2730 2731 2732 2733 2734 2735
            _C_ops.adamax_(
                param_and_grad[0],
                param_and_grad[1],
                self._create_param_lr(param_and_grad),
                moment,
                inf_norm,
                beta1_pow_acc,
                self._beta1,
                self._beta2,
                self._epsilon,
            )
2736
        elif framework._in_legacy_dygraph():
2737 2738 2739 2740 2741 2742 2743 2744 2745 2746 2747 2748 2749 2750 2751 2752 2753
            _legacy_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,
            )
2754 2755 2756 2757 2758 2759 2760 2761 2762 2763
        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,
2764
                    "Beta1Pow": beta1_pow_acc,
2765 2766 2767 2768
                },
                outputs={
                    "ParamOut": param_and_grad[0],
                    "MomentOut": moment,
2769
                    "InfNormOut": inf_norm,
2770 2771 2772 2773
                },
                attrs={
                    "beta1": self._beta1,
                    "beta2": self._beta2,
2774
                    "epsilon": self._epsilon,
2775
                },
2776 2777
                stop_gradient=True,
            )
2778

2779
            return adamax_op
2780

2781
    def _finish_update(self, block, parameters_and_grads):
2782
        """Update Beta1 Power accumulator"""
2783
        assert isinstance(block, framework.Block)
2784
        for param, grad in parameters_and_grads:
C
chengduo 已提交
2785
            if grad is None or param.trainable is False:
2786
                continue
2787 2788 2789 2790 2791 2792
            with param.block.program._optimized_guard(
                [param, grad]
            ), name_scope('adamx'):
                beta1_pow_acc = self._get_accumulator(
                    self._beta1_pow_acc_str, param
                )
J
Jiabin Yang 已提交
2793
                if framework._non_static_mode():
2794
                    if framework.in_dygraph_mode():
2795 2796 2797
                        tmp = _C_ops.scale(
                            beta1_pow_acc, self._beta1, 0.0, True
                        )
2798
                    else:
2799 2800 2801
                        tmp = _legacy_C_ops.scale(
                            beta1_pow_acc, "scale", self._beta1
                        )
2802 2803
                    beta1_pow_acc.copy_(tmp, False)
                else:
2804 2805 2806 2807 2808 2809 2810
                    block.append_op(
                        type="scale",
                        inputs={"X": beta1_pow_acc},
                        outputs={"Out": beta1_pow_acc},
                        attrs={"scale": self._beta1},
                        stop_gradient=True,
                    )
2811 2812


2813
class DpsgdOptimizer(Optimizer):
2814
    r"""
2815 2816 2817 2818 2819 2820 2821 2822
    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
2
201716010711 已提交
2823 2824
          import paddle
          paddle.enable_static()
2825 2826 2827 2828 2829 2830 2831 2832 2833 2834

          # 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)
2
201716010711 已提交
2835
              loss = paddle.mean(hidden)
2836 2837 2838 2839 2840 2841 2842 2843 2844 2845 2846 2847 2848 2849 2850 2851 2852
              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 已提交
2853
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2854 2855
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2856 2857 2858 2859
    Notes:
       Currently, DpsgdOptimizer doesn't support sparse parameter optimization.
    """

2860 2861 2862 2863 2864 2865 2866 2867
    def __init__(
        self,
        learning_rate=0.001,
        clip=0.9,
        batch_size=0.999,
        sigma=1e-8,
        parameter_list=None,
    ):
2868 2869 2870 2871
        assert learning_rate is not None
        assert clip is not None
        assert batch_size is not None
        assert sigma is not None
2872
        super().__init__(
2873 2874
            learning_rate=learning_rate, parameter_list=parameter_list
        )
2875 2876 2877 2878
        self.type = "dpsgd"
        self._clip = clip
        self._batch_size = batch_size
        self._sigma = sigma
Z
zhongpu 已提交
2879 2880 2881 2882 2883 2884 2885
        '''
        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
2886 2887 2888 2889 2890

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

        # create the dpsgd optimize op
2891
        if self._seed is None:
Z
zhongpu 已提交
2892 2893
            self._seed = 0

J
Jiabin Yang 已提交
2894
        if framework._non_static_mode():
2895 2896 2897 2898 2899 2900 2901 2902 2903 2904 2905 2906 2907 2908
            _legacy_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,
            )
2909
        else:
2910 2911 2912 2913 2914 2915 2916 2917 2918 2919 2920 2921 2922 2923 2924 2925
            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,
            )
2926

2927
            return dpsgd_op
2928 2929


2930
class DecayedAdagradOptimizer(Optimizer):
2931
    r"""
2932 2933 2934
    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.
2935

2936
    The parameter ``param_out`` update rule with gradient ``grad``:
2937 2938 2939 2940 2941 2942 2943

    .. math::

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

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

2944 2945 2946 2947
    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
2948 2949 2950
    stability to avoid the division by zero error.

    Args:
2951 2952 2953 2954 2955
        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 已提交
2956
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2957 2958
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2959 2960 2961 2962 2963
        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.
2964 2965 2966
        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` ,
2967
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2968 2969 2970 2971 2972 2973
        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.**
2974 2975 2976 2977

    Examples:
        .. code-block:: python

2978 2979
            import paddle.fluid as fluid

2980 2981 2982 2983
            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)
2984
            optimizer.minimize(cost)
2985 2986 2987
    """
    _moment_acc_str = "moment"

2988 2989 2990 2991 2992 2993 2994 2995 2996 2997
    def __init__(
        self,
        learning_rate,
        decay=0.95,
        epsilon=1.0e-6,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
2998 2999 3000 3001
        assert learning_rate is not None
        assert decay is not None
        assert epsilon is not None

3002
        super().__init__(
3003 3004 3005 3006 3007 3008
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
3009 3010 3011 3012 3013 3014 3015 3016 3017 3018 3019 3020 3021
        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)

3022 3023 3024
        moment_acc = self._get_accumulator(
            self._moment_acc_str, param_and_grad[0]
        )
3025

J
Jiabin Yang 已提交
3026
        if framework._non_static_mode():
3027 3028 3029 3030 3031 3032 3033 3034 3035 3036 3037 3038
            _legacy_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,
            )
3039 3040 3041 3042 3043 3044 3045 3046
        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,
3047
                    "LearningRate": self._create_param_lr(param_and_grad),
3048 3049 3050
                },
                outputs={
                    "ParamOut": param_and_grad[0],
3051
                    "MomentOut": moment_acc,
3052
                },
3053 3054 3055
                attrs={"epsilon": self._epsilon, "decay": self._decay},
                stop_gradient=True,
            )
3056

3057
            return decayed_adagrad_op
3058 3059


3060
class AdadeltaOptimizer(Optimizer):
3061
    r"""
Z
Zeng Jinle 已提交
3062
    **Notes: This API does not support sparse parameter optimization.**
Q
qiaolongfei 已提交
3063

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

    The update is done as follows:
3068

Z
Zeng Jinle 已提交
3069 3070
    .. math::

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

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

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

    Args:
Z
Zeng Jinle 已提交
3078 3079 3080
        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 已提交
3081
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3082 3083
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3084 3085 3086 3087 3088
        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.
3089 3090 3091
        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` ,
3092
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3093 3094 3095
        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` .
3096 3097 3098 3099

    Examples:
        .. code-block:: python

3100
            import paddle.fluid as fluid
Z
Zeng Jinle 已提交
3101

3102
            image = fluid.data(name='image', shape=[None, 28], dtype='float32')
Z
Zeng Jinle 已提交
3103 3104
            fc = fluid.layers.fc(image, size=10)
            cost = fluid.layers.reduce_mean(fc)
3105 3106
            optimizer = fluid.optimizer.Adadelta(
                learning_rate=0.0003, epsilon=1.0e-6, rho=0.95)
C
chengduo 已提交
3107

Z
Zeng Jinle 已提交
3108 3109 3110 3111
            # 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)
3112
    """
3113

3114 3115 3116
    _avg_squared_grad_acc_str = "_avg_squared_grad"
    _avg_squared_update_acc_str = "_avg_squared_update"

3117 3118 3119 3120 3121 3122 3123 3124 3125 3126
    def __init__(
        self,
        learning_rate,
        epsilon=1.0e-6,
        rho=0.95,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
3127 3128 3129 3130 3131 3132
        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.")
3133
        super().__init__(
3134 3135 3136 3137 3138 3139
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
3140 3141 3142 3143 3144
        self.type = "adadelta"
        self._epsilon = epsilon
        self._rho = rho

    def _create_accumulators(self, block, parameters):
3145 3146
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")
3147 3148 3149 3150 3151 3152

        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):
3153 3154
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")
3155 3156

        avg_squared_grad_acc = self._get_accumulator(
3157 3158
            self._avg_squared_grad_acc_str, param_and_grad[0]
        )
3159
        avg_squared_update_acc = self._get_accumulator(
3160 3161
            self._avg_squared_update_acc_str, param_and_grad[0]
        )
3162

3163
        if framework.in_dygraph_mode():
3164 3165 3166 3167 3168 3169 3170 3171
            _C_ops.adadelta_(
                param_and_grad[0],
                param_and_grad[1],
                avg_squared_grad_acc,
                avg_squared_update_acc,
                self._rho,
                self._epsilon,
            )
3172
        elif framework._in_legacy_dygraph():
3173 3174 3175 3176 3177 3178 3179 3180 3181 3182 3183 3184 3185
            _legacy_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,
            )
3186 3187
        else:
            # Create the adadelta optimizer op
3188 3189 3190 3191 3192 3193 3194 3195 3196 3197 3198 3199 3200 3201 3202 3203
            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,
            )
3204

3205
            return adadelta_op
3206 3207


Q
qingqing01 已提交
3208
class RMSPropOptimizer(Optimizer):
3209
    r"""
Q
qingqing01 已提交
3210 3211 3212 3213 3214 3215 3216 3217
    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 已提交
3218
        r(w, t) & = \\rho r(w, t-1) + (1 - \\rho)(\\nabla Q_{i}(w))^2
Q
qingqing01 已提交
3219 3220 3221 3222

        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 已提交
3223
    each weight. Then dividing the gradient by :math:`sqrt{v(w,t)}`.
Q
qingqing01 已提交
3224 3225 3226 3227 3228 3229

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

    ..  math::

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

3232 3233 3234 3235 3236 3237 3238 3239 3240 3241 3242 3243 3244 3245
        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 已提交
3246 3247 3248 3249
            \\epsilon}} \\nabla Q_{i}(w)

        w & = w - v(w, t)

Q
qiaolongfei 已提交
3250
    where, :math:`\\rho` is a hyperparameter and typical values are 0.9, 0.95
Q
qingqing01 已提交
3251 3252 3253 3254 3255
    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.


3256 3257 3258
    Parameters:
        learning_rate(float): Global learning rate.
        rho(float): rho is :math: `\\rho` in equation, default is 0.95.
Q
qingqing01 已提交
3259
        epsilon(float): :math: `\\epsilon` in equation is smoothing term to
3260
            avoid division by zero, default is 1e-6.
Q
qiaolongfei 已提交
3261
        momentum(float): :math:`\\beta` in equation is the momentum term,
3262
            default is 0.0.
3263 3264 3265 3266
        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 已提交
3267
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3268 3269
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3270 3271 3272 3273 3274
        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.
3275 3276 3277
        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` ,
3278
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3279 3280
        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 已提交
3281 3282 3283 3284 3285 3286 3287

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

    Examples:
          .. code-block:: python

3288 3289 3290 3291
            import paddle
            import paddle.fluid as fluid
            import numpy as np

3292
            paddle.enable_static()
3293 3294 3295 3296 3297 3298
            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)
3299
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
3300
                avg_cost = paddle.mean(cost)
3301 3302 3303 3304 3305 3306 3307 3308 3309 3310 3311 3312 3313

                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 已提交
3314 3315 3316 3317
    """

    _momentum_acc_str = "momentum"
    _mean_square_acc_str = "mean_square"
3318
    _mean_grad_acc_str = "mean_grad"
Q
qingqing01 已提交
3319

3320 3321 3322 3323 3324 3325 3326 3327 3328 3329 3330 3331
    def __init__(
        self,
        learning_rate,
        rho=0.95,
        epsilon=1.0e-6,
        momentum=0.0,
        centered=False,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
3332
        super().__init__(
3333 3334 3335 3336 3337 3338
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
Q
qingqing01 已提交
3339 3340 3341 3342 3343 3344 3345 3346 3347 3348 3349 3350 3351
        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
3352
        self._centered = centered
Q
qingqing01 已提交
3353 3354 3355 3356 3357 3358 3359 3360

    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)
3361
            self._add_accumulator(self._mean_grad_acc_str, p)
Q
qingqing01 已提交
3362 3363 3364 3365 3366

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

3367 3368 3369 3370 3371 3372 3373 3374 3375
        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]
        )
        mean_grad_acc = self._get_accumulator(
            self._mean_grad_acc_str, param_and_grad[0]
        )
C
caozhou 已提交
3376
        if in_dygraph_mode():
3377 3378 3379 3380 3381 3382 3383 3384 3385 3386 3387 3388
            _C_ops.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,
            )
C
caozhou 已提交
3389 3390
            return None
        elif _in_legacy_dygraph():
3391 3392 3393 3394 3395 3396 3397 3398 3399 3400 3401 3402 3403 3404 3405 3406 3407 3408 3409
            _legacy_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 已提交
3410
            return None
3411 3412 3413 3414 3415 3416 3417 3418 3419 3420 3421 3422 3423 3424 3425
        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,
3426
                    "MeanGradOut": mean_grad_acc,
3427 3428 3429 3430 3431
                },
                attrs={
                    "epsilon": self._epsilon,
                    "decay": self._rho,
                    "momentum": self._momentum,
3432
                    "centered": self._centered,
3433
                },
3434 3435
                stop_gradient=True,
            )
Q
qingqing01 已提交
3436

3437
            return rmsprop_op
Q
qingqing01 已提交
3438 3439


Q
qiaolongfei 已提交
3440
class FtrlOptimizer(Optimizer):
3441
    r"""
Q
qiaolongfei 已提交
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 3468 3469 3470 3471 3472 3473 3474 3475 3476 3477 3478 3479
    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

3480 3481 3482 3483 3484
    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 已提交
3485
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3486 3487
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3488 3489 3490 3491 3492
        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.
3493 3494 3495
        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` ,
3496
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3497 3498
        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 已提交
3499 3500 3501 3502 3503 3504 3505

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

    Examples:
          .. code-block:: python

3506 3507 3508 3509
            import paddle
            import paddle.fluid as fluid
            import numpy as np

3510 3511
            paddle.enable_static()

3512 3513 3514 3515 3516 3517
            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)
3518
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
3519
                avg_cost = paddle.mean(cost)
3520 3521 3522 3523 3524 3525 3526 3527 3528 3529 3530 3531

                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 已提交
3532

3533
    NOTE:
C
chengduo 已提交
3534
       Currently, FtrlOptimizer doesn't support sparse parameter optimization.
Q
qiaolongfei 已提交
3535 3536 3537 3538 3539
    """

    _squared_acc_str = "squared"
    _linear_acc_str = "linear"

3540 3541 3542 3543 3544 3545 3546 3547 3548 3549 3550
    def __init__(
        self,
        learning_rate,
        l1=0.0,
        l2=0.0,
        lr_power=-0.5,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
3551
        super().__init__(
3552 3553 3554 3555 3556 3557
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
Q
qiaolongfei 已提交
3558 3559 3560 3561 3562 3563 3564 3565 3566 3567 3568 3569 3570 3571 3572 3573 3574 3575 3576 3577
        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.")

3578 3579 3580 3581 3582 3583
        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 已提交
3584
        if framework._non_static_mode():
3585 3586 3587 3588 3589 3590 3591 3592 3593 3594 3595 3596 3597 3598 3599 3600
            _legacy_C_ops.ftrl(
                param_and_grad[0],
                squared_acc,
                linear_acc,
                param_and_grad[1],
                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,
            )
3601 3602

        else:
3603 3604 3605 3606 3607 3608 3609 3610 3611 3612 3613 3614 3615 3616 3617 3618 3619 3620 3621 3622 3623
            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 已提交
3624

3625
            return ftrl_op
Q
qiaolongfei 已提交
3626 3627


Y
Yibing Liu 已提交
3628
class LambOptimizer(AdamOptimizer):
3629
    r"""
Y
Yibing Liu 已提交
3630 3631
    LAMB (Layer-wise Adaptive Moments optimizer for Batching training) Optimizer.

3632 3633 3634
    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
    correction. For more information, please refer to `Large Batch Optimization for
Y
Yibing Liu 已提交
3635
    Deep Learning: Training BERT in 76 minutes <https://arxiv.org/abs/1904.00962>`_ .
Y
Yibing Liu 已提交
3636 3637 3638 3639 3640

    The updating of parameters follows:

    ..  math::

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

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

3645 3646 3647 3648
        m_t &= \\frac{m_t}{\\beta_1^t}

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

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

Y
Yibing Liu 已提交
3651
        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 已提交
3652 3653


3654
    where :math:`m` is the 1st moment, and :math:`v` the 2nd moment, :math:`\\eta` the
Y
Yibing Liu 已提交
3655 3656 3657
    learning rate, :math:`\\lambda` the LAMB weight decay rate.

    Args:
Y
Yibing Liu 已提交
3658 3659 3660 3661 3662 3663 3664 3665
        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 已提交
3666
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3667 3668
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3669 3670 3671 3672 3673
        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.
3674 3675
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of
            some derived class of ``GradientClipBase`` . There are three cliping strategies
3676 3677 3678
            ( :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.
3679 3680
        exclude_from_weight_decay_fn (function|None): Exclude a parameter from weight
            decay when **exclude_from_weight_decay_fn(parameter)** returns true.
Y
Yibing Liu 已提交
3681
            Default None.
3682
        name(str|None): For detailed information, please refer to
Y
Yibing Liu 已提交
3683
            :ref:`api_guide_Name` . Usually name is no need to set and None by default.
Y
Yibing Liu 已提交
3684 3685 3686

    Examples:
        .. code-block:: python
3687

2
201716010711 已提交
3688
            import paddle
3689
            import paddle.fluid as fluid
2
201716010711 已提交
3690
            paddle.enable_static()
Y
Yibing Liu 已提交
3691

Y
Yibing Liu 已提交
3692
            data = fluid.data(name='x', shape=[-1, 5], dtype='float32')
Y
Yibing Liu 已提交
3693
            hidden = fluid.layers.fc(input=data, size=10)
2
201716010711 已提交
3694
            cost = paddle.mean(hidden)
Y
Yibing Liu 已提交
3695

Y
Yibing Liu 已提交
3696 3697 3698 3699 3700
            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 已提交
3701 3702 3703 3704 3705 3706 3707
            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"

3708 3709 3710 3711 3712 3713 3714 3715 3716 3717 3718 3719 3720
    def __init__(
        self,
        learning_rate=0.001,
        lamb_weight_decay=0.01,
        beta1=0.9,
        beta2=0.999,
        epsilon=1e-6,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        exclude_from_weight_decay_fn=None,
        name=None,
    ):
Y
Yibing Liu 已提交
3721 3722 3723 3724 3725
        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
3726
        super().__init__(
3727 3728 3729 3730 3731 3732 3733 3734 3735
            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 已提交
3736 3737
        self.type = "lamb"
        self._weight_decay = lamb_weight_decay
Y
Yibing Liu 已提交
3738
        self._exclude_from_weight_decay_fn = exclude_from_weight_decay_fn
Y
Yibing Liu 已提交
3739 3740 3741

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)
3742
        block.program._use_lamb = True
Y
Yibing Liu 已提交
3743

3744 3745 3746 3747 3748 3749 3750 3751 3752 3753 3754 3755 3756 3757 3758 3759 3760
        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]
        )

        if (
            self._exclude_from_weight_decay_fn is not None
            and self._exclude_from_weight_decay_fn(param_and_grad[0])
        ):
Y
Yibing Liu 已提交
3761 3762 3763
            weight_decay = 0.0
        else:
            weight_decay = self._weight_decay
3764
        lr = self._create_param_lr(param_and_grad)
3765
        master_weight = None
J
Jiabin Yang 已提交
3766
        if framework._non_static_mode():
3767 3768 3769 3770 3771 3772 3773 3774 3775 3776 3777 3778 3779 3780 3781 3782 3783 3784 3785 3786 3787 3788 3789 3790
            _legacy_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,
            )
3791
            return None
Y
Yibing Liu 已提交
3792

Y
Yibing Liu 已提交
3793
        # create the lamb optimize op
3794 3795 3796 3797 3798 3799 3800 3801 3802 3803 3804 3805 3806 3807 3808 3809 3810 3811 3812 3813 3814 3815 3816 3817 3818 3819
        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 已提交
3820 3821 3822 3823

        return lamb_op


3824 3825 3826 3827 3828 3829 3830 3831 3832 3833 3834 3835 3836
# 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
3837
Dpsgd = DpsgdOptimizer
3838
DecayedAdagrad = DecayedAdagradOptimizer
3839
Adadelta = AdadeltaOptimizer
Q
qingqing01 已提交
3840
RMSProp = RMSPropOptimizer
Q
qiaolongfei 已提交
3841
Ftrl = FtrlOptimizer
3842
LarsMomentum = LarsMomentumOptimizer
Y
Yibing Liu 已提交
3843
Lamb = LambOptimizer
3844 3845 3846


class ModelAverage(Optimizer):
3847
    r"""
3848
	:api_attr: Static Graph
S
swtkiwi 已提交
3849

3850 3851 3852 3853 3854 3855 3856 3857 3858 3859 3860 3861 3862 3863 3864 3865 3866 3867
    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:

    ::
3868

3869 3870 3871 3872 3873 3874 3875 3876 3877
        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.
3878 3879

    Args:
3880 3881 3882
        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.
3883 3884 3885 3886 3887
        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.
3888 3889 3890
        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.
3891

3892
    Examples:
Q
qiaolongfei 已提交
3893 3894 3895

      .. code-block:: python

2
201716010711 已提交
3896
        import paddle
3897 3898
        import paddle.fluid as fluid
        import numpy
2
201716010711 已提交
3899
        paddle.enable_static()
3900 3901 3902 3903

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

3905 3906 3907 3908
        train_program = fluid.Program()
        startup_program = fluid.Program()
        with fluid.program_guard(train_program, startup_program):
            # build net
3909
            data = fluid.data(name='X', shape=[None, 1], dtype='float32')
3910
            hidden = fluid.layers.fc(input=data, size=10)
2
201716010711 已提交
3911
            loss = paddle.mean(hidden)
3912 3913 3914 3915 3916 3917
            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,
3918
                                                         max_average_window=12500)
3919 3920

            exe.run(startup_program)
3921 3922 3923 3924 3925
            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])
3926 3927

            # apply ModelAverage
3928
            with model_average.apply(exe):
3929 3930 3931 3932
                x = numpy.random.random(size=(10, 1)).astype('float32')
                exe.run(program=train_program,
                        feed={'X': x},
                        fetch_list=[loss.name])
3933 3934
    """

3935 3936 3937 3938 3939 3940 3941 3942
    def __init__(
        self,
        average_window_rate,
        min_average_window=10000,
        max_average_window=10000,
        regularization=None,
        name=None,
    ):
J
Jiabin Yang 已提交
3943
        if framework._non_static_mode():
Z
zhongpu 已提交
3944
            raise Exception("In dygraph, don't support ModelAverage.")
3945
        super().__init__(0.0, regularization=regularization, name=name)
3946 3947 3948
        self.average_window = average_window_rate
        self.min_average_window = min_average_window
        self.max_average_window = max_average_window
3949

3950
        self.params_grads = []
3951 3952 3953
        for param in (
            framework.default_main_program().global_block().all_parameters()
        ):
3954
            if param.do_model_average != False:
3955
                grad = param.block.create_var(
3956 3957 3958
                    name=unique_name.generate_with_ignorable_key(
                        ".".join([param.name, 'tmp'])
                    ),
3959 3960
                    dtype=param.dtype,
                    persistable=False,
3961 3962
                    stop_gradient=True,
                )
3963
                self.params_grads.append((param, grad))
3964

3965
        for param, grad in self.params_grads:
3966 3967
            if grad is None:
                continue
X
Xin Pan 已提交
3968
            with param.block.program._optimized_guard(
3969 3970
                [param, grad]
            ), name_scope('move_average'):
3971
                self._append_average_accumulate_op(param)
3972

3973 3974 3975 3976
        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:
3977
                self._add_average_apply_op(block, param_grad)
3978 3979 3980 3981 3982

        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:
3983
                self._add_average_restore_op(block, param_grad)
3984

3985
    def _add_average_apply_op(self, block, param_grad):
L
Luo Tao 已提交
3986 3987 3988 3989 3990 3991
        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(
3992 3993
            self._get_accumulator('num_accumulates', param)
        )
L
Luo Tao 已提交
3994
        old_num_accumulates = block._clone_variable(
3995 3996
            self._get_accumulator('old_num_accumulates', param)
        )
L
Luo Tao 已提交
3997
        num_updates = block._clone_variable(
3998 3999
            self._get_accumulator('num_updates', param)
        )
4000 4001 4002
        # backup param value to grad
        layers.assign(input=param, output=grad)
        # param = (sum_1 + sum_2 + sum_3) / (num_accumulates + old_num_accumulates)
4003 4004
        tmp = paddle.add_n([num_accumulates, old_num_accumulates])
        sum = paddle.add_n([sum_1, sum_2, sum_3])
D
dzhwinter 已提交
4005
        tmp = layers.cast(
4006
            x=tmp, dtype='float32' if self._dtype is None else self._dtype
4007
        )
D
dzhwinter 已提交
4008
        sum = layers.cast(
4009
            x=sum, dtype='float32' if self._dtype is None else self._dtype
4010
        )
4011
        paddle.assign(paddle.divide(sum, tmp), output=param)
4012 4013

    def _add_average_restore_op(self, block, param_grad):
L
Luo Tao 已提交
4014 4015
        param = block._clone_variable(param_grad[0])
        grad = block._clone_variable(param_grad[1])
4016 4017 4018 4019 4020 4021 4022
        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)
4023 4024 4025 4026 4027 4028 4029 4030 4031 4032 4033 4034 4035 4036 4037 4038 4039 4040 4041 4042 4043 4044 4045 4046 4047 4048 4049 4050 4051 4052 4053 4054 4055 4056 4057 4058
        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,
        )
4059

S
rename  
sneaxiy 已提交
4060
    @signature_safe_contextmanager
4061
    def apply(self, executor, need_restore=True):
4062 4063
        """
        Apply the average of the cumulative ``Parameter`` to the parameters of the current model.
4064 4065

        Args:
4066 4067 4068 4069 4070 4071 4072 4073 4074 4075 4076
            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
2
201716010711 已提交
4077 4078
            import paddle
            paddle.enable_static()
4079 4080 4081 4082 4083 4084 4085 4086 4087 4088 4089

            # 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)
2
201716010711 已提交
4090
                loss = paddle.mean(hidden)
4091 4092 4093 4094 4095 4096 4097 4098 4099 4100 4101 4102 4103 4104 4105 4106 4107 4108 4109 4110 4111
                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])
4112
        """
4113 4114 4115 4116 4117 4118
        executor.run(self.apply_program)
        try:
            yield
        finally:
            if need_restore:
                self.restore(executor)
4119 4120

    def restore(self, executor):
4121 4122
        """
        Restore ``Parameter`` values of current model.
4123

4124
        Args:
4125 4126 4127 4128 4129 4130 4131 4132
            executor(fluid.Executor): The current network executor.

        Examples:

          .. code-block:: python

            import paddle.fluid as fluid
            import numpy
2
201716010711 已提交
4133 4134
            import paddle
            paddle.enable_static()
4135 4136 4137 4138 4139 4140 4141 4142 4143 4144 4145

            # 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)
2
201716010711 已提交
4146
                loss = paddle.mean(hidden)
4147 4148 4149 4150 4151 4152 4153 4154 4155 4156 4157 4158 4159 4160 4161 4162 4163 4164 4165 4166 4167 4168 4169 4170
                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)
4171
        """
4172
        executor.run(self.restore_program)
4173 4174


4175
class ExponentialMovingAverage:
4176
    r"""
4177
        :api_attr: Static Graph
S
swtkiwi 已提交
4178

4179 4180 4181 4182 4183 4184
    Compute the moving average of parameters with exponential decay.
    Given a parameter :math:`\\theta`, its exponential moving average (EMA)
    will be

    ..  math::

4185
        \\text{EMA}_0 & = 0
4186

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

4189 4190 4191
    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
Y
Yibing Liu 已提交
4192
    the **restore()** method is used to restore the parameters.
4193

4194 4195 4196 4197
    **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
4198 4199

    ..  math::
4200

4201 4202
        \\widehat{\\text{EMA}}_t = \\frac{\\text{EMA}_t}{1 - \\text{decay}^t}

4203 4204
    **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
4205
    relative smaller decay rate in the very beginning. The argument **thres_steps**
4206
    allows users to pass a Variable to schedule the decay rate, in this case,
4207
    the actual decay rate becomes
4208

4209
    ..  math::
4210

4211 4212 4213
        \\min(\\text{decay}, \\frac{1 + \\text{thres_steps}}{10 + \\text{thres_steps}})

    Usually **thres_steps** can be the global training steps.
4214 4215 4216


    Args:
4217 4218 4219
        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.
4220 4221 4222 4223


    Examples:

4224 4225 4226 4227 4228 4229 4230 4231 4232 4233 4234 4235 4236 4237 4238 4239 4240 4241 4242 4243 4244 4245 4246 4247 4248 4249 4250 4251
        .. 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(),
4252
                    feed={'x': data},
4253 4254 4255 4256 4257 4258
                    fetch_list=[cost.name])

                # usage 1
                with ema.apply(exe):
                    data = numpy.random.random(size=(10, 5)).astype('float32')
                    exe.run(program=test_program,
4259
                        feed={'x': data},
4260 4261 4262 4263 4264 4265
                        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,
4266
                        feed={'x': data},
4267 4268 4269
                        fetch_list=[hidden.name])
                ema.restore(exe)

4270 4271
    """

4272
    def __init__(self, decay=0.999, thres_steps=None, name=None):
J
Jiabin Yang 已提交
4273
        if framework._non_static_mode():
Z
zhongpu 已提交
4274
            raise Exception(
4275 4276
                "In dygraph, don't support ExponentialMovingAverage."
            )
4277
        self._decay = decay
4278
        self._thres_steps = thres_steps
4279
        self._name = name if name is not None else ''
4280 4281
        self._decay_var = self._get_ema_decay()

4282
        self._step_counter_name = "@EMA_STEP_COUNTER@"
Y
Yibing Liu 已提交
4283
        self._params_tmps = []
4284
        for param in default_main_program().global_block().all_parameters():
4285
            if param.do_model_average != False:
4286 4287 4288 4289 4290 4291 4292 4293
                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 已提交
4294
                self._params_tmps.append((param, tmp))
4295

Y
Yibing Liu 已提交
4296 4297
        self._ema_vars = {}
        for param, tmp in self._params_tmps:
4298 4299 4300
            with param.block.program._optimized_guard([param, tmp]), name_scope(
                'moving_average'
            ):
Y
Yibing Liu 已提交
4301
                self._ema_vars[param.name] = self._create_ema_vars(param)
4302 4303 4304 4305

        self.apply_program = Program()
        block = self.apply_program.global_block()
        with program_guard(main_program=self.apply_program):
4306
            decay_pow, global_step = self._get_decay_pow(block)
Y
Yibing Liu 已提交
4307
            for param, tmp in self._params_tmps:
4308 4309
                param = block._clone_variable(param)
                tmp = block._clone_variable(tmp)
Y
Yibing Liu 已提交
4310
                ema = block._clone_variable(self._ema_vars[param.name])
4311
                layers.assign(input=param, output=tmp)
4312
                # bias correction
4313 4314
                with layers.control_flow.Switch() as switch:
                    with switch.case(global_step > 0):
4315 4316 4317
                        layers.assign(
                            output=param, input=ema / (1.0 - decay_pow)
                        )
4318 4319
                    with switch.default():
                        layers.assign(output=param, input=ema)
4320 4321 4322 4323

        self.restore_program = Program()
        block = self.restore_program.global_block()
        with program_guard(main_program=self.restore_program):
Y
Yibing Liu 已提交
4324
            for param, tmp in self._params_tmps:
4325 4326 4327 4328
                tmp = block._clone_variable(tmp)
                param = block._clone_variable(param)
                layers.assign(input=tmp, output=param)

4329 4330
    def _get_ema_decay(self):
        with default_main_program()._lr_schedule_guard():
4331
            decay_var = paddle.static.create_global_var(
4332 4333 4334 4335
                shape=[1],
                value=self._decay,
                dtype='float32',
                persistable=True,
4336 4337
                name="scheduled_ema_decay_rate",
            )
4338 4339 4340 4341 4342 4343 4344 4345

            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(
4346 4347
                            np.array([self._decay], dtype=np.float32), decay_var
                        )
4348 4349 4350
        return decay_var

    def _get_decay_pow(self, block):
4351
        global_step = paddle.static.create_global_var(
4352 4353 4354 4355 4356 4357
            name=self._step_counter_name,
            shape=[1],
            value=0,
            dtype='int64',
            persistable=True,
        )
4358
        global_step = layers.cast(global_step, "float32")
4359
        decay_var = block._clone_variable(self._decay_var)
4360
        decay_pow_acc = paddle.pow(decay_var, global_step)
4361
        return decay_pow_acc, global_step
4362

Y
Yibing Liu 已提交
4363
    def _create_ema_vars(self, param):
4364
        param_ema = paddle.static.create_global_var(
4365 4366 4367 4368
            name=unique_name.generate(self._name + param.name + '_ema'),
            shape=param.shape,
            value=0.0,
            dtype=param.dtype,
4369 4370
            persistable=True,
        )
4371 4372 4373

        return param_ema

Y
Yibing Liu 已提交
4374
    def update(self):
4375 4376
        """
        Update Exponential Moving Average. Should only call this method in
Y
Yibing Liu 已提交
4377 4378
        train program.
        """
4379
        global_step = layers.autoincreased_step_counter(
4380 4381
            counter_name=self._step_counter_name
        )
4382
        param_master_emas = []
Y
Yibing Liu 已提交
4383
        for param, tmp in self._params_tmps:
4384 4385 4386
            with param.block.program._optimized_guard([param, tmp]), name_scope(
                'moving_average'
            ):
Y
Yibing Liu 已提交
4387
                param_ema = self._ema_vars[param.name]
4388
                if param.name + '.master' in self._ema_vars:
4389 4390 4391 4392
                    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 * (
4393 4394
                        1 - self._decay_var
                    )
4395 4396 4397 4398 4399 4400 4401 4402 4403 4404
                    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,
4405 4406 4407
                    "out_dtype": param_ema.dtype,
                },
            )
Y
Yibing Liu 已提交
4408

4409 4410 4411 4412
    @signature_safe_contextmanager
    def apply(self, executor, need_restore=True):
        """
        Apply moving average to parameters for evaluation.
4413

4414 4415
        Args:
            executor (Executor): The Executor to execute applying.
4416
            need_restore (bool, optional): Whether to restore parameters after
Y
Yibing Liu 已提交
4417
                applying. Default True.
4418 4419 4420 4421 4422 4423 4424 4425 4426 4427
        """
        executor.run(self.apply_program)
        try:
            yield
        finally:
            if need_restore:
                self.restore(executor)

    def restore(self, executor):
        """Restore parameters.
4428

4429 4430 4431 4432
        Args:
            executor (Executor): The Executor to execute restoring.
        """
        executor.run(self.restore_program)
H
hutuxian 已提交
4433 4434


4435
class PipelineOptimizer:
4436
    """
4437
        :api_attr: Static Graph
S
swtkiwi 已提交
4438

4439 4440 4441 4442
    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 已提交
4443

4444
    Args:
4445 4446 4447
        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].
4448

4449 4450
    Examples:
        .. code-block:: python
H
hutuxian 已提交
4451

4452
            import paddle.fluid as fluid
H
hutuxian 已提交
4453 4454
            import paddle.fluid.layers as layers

4455 4456 4457 4458 4459 4460 4461 4462 4463 4464 4465 4466 4467 4468 4469 4470
            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 已提交
4471
            optimizer = fluid.optimizer.SGD(learning_rate=0.5)
4472
            optimizer = fluid.optimizer.PipelineOptimizer(optimizer)
H
hutuxian 已提交
4473
            optimizer.minimize(loss)
4474 4475 4476 4477 4478 4479 4480 4481 4482

            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 已提交
4483 4484
            exe = fluid.Executor(place)
            exe.run(fluid.default_startup_program())
4485 4486
            batch_size = 1
            data_loader.start()
H
hutuxian 已提交
4487
            exe.train_from_dataset(
4488
                    fluid.default_main_program())
4489
            data_loader.reset()
4490 4491
    """

4492
    def __init__(self, optimizer, num_microbatches=1, start_cpu_core_id=0):
4493 4494 4495 4496 4497
        self._device = 'cpu'
        if core.is_compiled_with_npu():
            self._device = "npu"
        elif core.is_compiled_with_cuda():
            self._device = "gpu"
J
Jiabin Yang 已提交
4498
        if framework._non_static_mode():
Z
zhongpu 已提交
4499
            raise Exception("In dygraph, don't support PipelineOptimizer.")
4500 4501 4502 4503 4504
        valid_optimizers = (
            Optimizer,
            paddle.optimizer.Optimizer,
            paddle.fluid.contrib.mixed_precision.decorator.OptimizerWithMixedPrecision,
        )
4505
        if not isinstance(optimizer, valid_optimizers):
4506 4507 4508 4509 4510 4511 4512
            raise ValueError(
                "The 'optimizer' parameter for "
                "PipelineOptimizer must be an instance of "
                "{}, but the given type is {}.".format(
                    valid_optimizers, type(optimizer)
                )
            )
H
hutuxian 已提交
4513
        self._optimizer = optimizer
4514 4515 4516 4517 4518 4519

        # 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

4520 4521 4522
        assert (
            num_microbatches >= 1
        ), "num_microbatches must be a positive value."
4523
        self._num_microbatches = num_microbatches
4524 4525 4526
        assert (
            start_cpu_core_id >= 0
        ), "start_cpu_core_id must be a non-negative integer."
H
hutuxian 已提交
4527
        self._start_cpu_core_id = start_cpu_core_id
4528 4529 4530 4531 4532 4533
        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()
4534
        self._param_device_map = None
4535 4536
        self._pipeline_pair = []
        self._pp_ring_map = dict()
4537 4538
        self.output_var_to_op = None
        self.input_var_to_op = None
4539 4540 4541 4542 4543 4544 4545 4546 4547 4548 4549 4550 4551 4552 4553

    # 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")
4554 4555 4556 4557 4558 4559 4560 4561 4562 4563 4564 4565 4566 4567
            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,
                },
            )
4568 4569 4570 4571
            offset += 1
        block._insert_op(
            op_idx + 1 + offset,
            type='c_allreduce_max'
4572 4573
            if op.type == "reduce_any"
            else 'c_allreduce_sum',
4574 4575 4576
            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={
4577
                'ring_id': self.global_ring_id,
4578
                self._op_role_key: self._op_role.Optimize,
4579 4580 4581
                'use_calc_stream': True,
            },
        )
4582 4583
        offset += 1
        if op.type == "reduce_any":
4584 4585 4586 4587 4588 4589 4590 4591 4592 4593 4594
            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,
                },
            )
4595
            offset += 1
4596
        return offset
H
hutuxian 已提交
4597

4598
    def _create_vars(self, block, ori_block):
4599
        # Create vars for block, copied from ori_block
H
hutuxian 已提交
4600
        used_var_set = set()
4601 4602 4603 4604 4605 4606 4607 4608 4609
        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]
4610
            # For op process vars on all devices, remove its input
4611 4612 4613 4614 4615 4616 4617 4618 4619 4620 4621 4622 4623 4624 4625
            # 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)
4626 4627 4628 4629 4630 4631 4632 4633 4634 4635
            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
4636 4637 4638 4639 4640 4641 4642 4643
            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 已提交
4644
            for var in vars:
4645 4646
                # a var whose name contains "blocking_queue"
                # only exists in startup program
4647
                if var in used_var_set or "_blocking_queue" in var:
H
hutuxian 已提交
4648 4649
                    continue
                used_var_set.add(var)
4650 4651
                if block._find_var_recursive(str(var)):
                    continue
4652
                source_var = ori_block._var_recursive(str(var))
4653
                if source_var.type == core.VarDesc.VarType.READER:
4654
                    dest_var = block.create_var(
4655 4656
                        name=var,
                        type=core.VarDesc.VarType.READER,
4657 4658
                        persistable=source_var.persistable,
                    )
4659 4660 4661 4662 4663 4664 4665 4666 4667 4668 4669
                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,
4670 4671
                        error_clip=source_var.error_clip,
                    )
4672
                else:
4673
                    dest_var = block._clone_variable(source_var, False)
4674
                self._clone_var_attr(dest_var, source_var)
4675 4676 4677
            # When use with sharding, allreduce_sum and allreduce_max
            # used for global gradient clip and amp will be added by sharding.
            op_idx += 1
4678 4679
            if self.use_sharding or not should_insert:
                continue
4680 4681 4682 4683
            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 已提交
4684

4685
    def _is_loss_grad_op(self, op):
4686 4687
        assert self._op_role_key in op.attr_names
        op_role = int(op.attr(self._op_role_key))
4688
        return op_role & int(self._op_role.Backward) and op_role & int(
4689 4690
            self._op_role.Loss
        )
4691

4692
    def _is_forward_op(self, op):
4693 4694 4695
        return self._op_role_key in op.attr_names and (
            int(op.attr(self._op_role_key)) == int(self._op_role.Forward)
        )
4696

4697
    def _is_backward_op(self, op):
4698
        return self._op_role_key in op.attr_names and (
4699 4700
            int(op.attr(self._op_role_key)) & int(self._op_role.Backward)
        )
4701 4702 4703 4704

    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)
4705 4706

    def _is_optimize_op(self, op):
4707
        return self._op_role_key in op.attr_names and (
4708 4709
            int(op.attr(self._op_role_key)) & int(self._op_role.Optimize)
        )
4710 4711

    def _is_update_op(self, op):
4712 4713 4714 4715 4716
        return (
            'Param' in op.input_names
            and 'Grad' in op.input_names
            and ("LearningRate" in op.input_names)
        )
4717

4718
    def _split_program(self, main_program, devices):
H
hutuxian 已提交
4719
        """
4720
        Split a program into sections according to devices that ops run on.
4721
        The op whose op_device attr is "gpu:all" is copied to all sections.
4722 4723 4724

        Args:
            main_program (Program): the main program
4725
            devices: all used devices
H
hutuxian 已提交
4726
        """
4727
        # Map from device to its corresponding section program info
4728
        device_program_map = defaultdict(Program)
4729

4730
        block = main_program.block(0)
4731 4732
        for op in block.ops:
            device = op.attr(self._op_device_key)
4733
            # Copy ops whose op_device set to "gpu:all" to all sections.
4734
            if device == f"{self._device}:all":
4735
                for device in devices:
4736 4737
                    program = device_program_map[device]
                    op_desc = op.desc
4738
                    ap_op = program.global_block().desc.append_op()
4739
                    ap_op.copy_from(op_desc)
4740
                    ap_op._set_attr(self._op_device_key, "")
4741 4742 4743
            else:
                program = device_program_map[device]
                op_desc = op.desc
4744
                ap_op = program.global_block().desc.append_op()
4745
                ap_op.copy_from(op_desc)
4746
                ap_op._set_attr(self._op_device_key, "")
4747

4748
        program_list = []
4749
        for key in devices:
4750
            program = device_program_map[key]
4751 4752
            program._sync_with_cpp()
            program_list.append(program)
H
hutuxian 已提交
4753

4754
        return program_list
H
hutuxian 已提交
4755

4756 4757 4758 4759 4760 4761 4762
    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.
        """
4763 4764
        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 '
4765
            'or beta2_pow_acc.'
4766 4767
        )
        param_name = var_name[0 : var_name.index('_beta')]
4768 4769 4770
        device = self._param_device_map[param_name]
        return device

4771 4772
    def _split_startup_program(self, startup_program, device_id):
        block = startup_program.global_block()
4773 4774 4775
        new_startup_program = Program()
        for op in block.ops:
            device = op.attr(self._op_device_key)
4776 4777
            if device == "cpu":
                assert op.type == "fill_constant", (
4778
                    "For ops in startup program with the op_device attribute "
4779 4780
                    "of cpu, they must be of type fill_constant."
                )
4781 4782 4783
                output_var = op.output_arg_names[0]
                device = self._get_op_device_for_startup_program(output_var)

4784
            if device:
4785
                device_index = int(device.split(':')[1])
4786
            else:
4787 4788
                # LR related ops
                device = None
4789 4790
            if device and device_index != device_id:
                continue
4791
            op_desc = op.desc
4792
            ap_op = new_startup_program.global_block().desc.append_op()
4793 4794 4795
            ap_op.copy_from(op_desc)
            ap_op._set_attr(self._op_device_key, "")
        new_startup_program._sync_with_cpp()
4796
        self._create_vars(new_startup_program.global_block(), block)
4797 4798
        return new_startup_program

4799
    def _find_post_op(self, index, var_name):
H
hutuxian 已提交
4800
        """
4801
        Find the post op that has variable named var_name as input.
H
hutuxian 已提交
4802
        """
4803 4804 4805 4806 4807 4808
        # 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', '')

4809
        post_ops = self.input_var_to_op[var_name]
4810
        if post_ops is None:
4811
            return None
4812 4813 4814 4815 4816 4817
        result_op = None
        for post_op, post_idx in reversed(post_ops):
            if post_idx > index:
                result_op = post_op
                break
        return result_op
4818

4819
    def _find_prev_op(self, index, var_name):
H
hutuxian 已提交
4820
        """
4821 4822
        Find the previous op of op with index that outputs
        variable named var_name.
H
hutuxian 已提交
4823
        """
4824
        prev_ops = self.output_var_to_op[var_name]
4825
        if prev_ops is None:
4826
            return None
4827 4828 4829 4830
        result_op = None
        for prev_op, prev_idx in reversed(prev_ops):
            if prev_idx < index:
                result_op = prev_op
4831
                break
4832
        return result_op
4833 4834

    def _rename_arg(self, op, old_name, new_name):
4835 4836
        op._rename_input(old_name, new_name)
        op._rename_output(old_name, new_name)
4837

4838
    def _create_var(self, block, ref_var, name, dtype=None):
4839 4840 4841 4842 4843 4844 4845 4846
        """
        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,
4847
            dtype=ref_var.dtype if dtype is None else dtype,
4848 4849
            type=ref_var.type,
            lod_level=ref_var.lod_level,
4850 4851
            persistable=ref_var.persistable,
            is_data=ref_var.is_data,
4852 4853
            need_check_feed=ref_var.desc.need_check_feed(),
        )
4854
        self._clone_var_attr(new_var, ref_var)
4855 4856
        return new_var

4857 4858 4859 4860 4861
    def _clone_var_attr(self, dest, src):
        dest.stop_gradient = src.stop_gradient
        if hasattr(src, 'is_distributed'):
            dest.is_distributed = src.is_distributed

4862 4863 4864 4865 4866 4867
    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 已提交
4868

4869 4870 4871 4872 4873 4874
    def _append_grad_suffix(self, name):
        """
        Append grad suffix to the given variable name
        """
        return name + core.grad_var_suffix()

4875
    def _get_op_device_attr(self, op):
H
hutuxian 已提交
4876
        """
4877
        Get the op_device attribute of a op.
H
hutuxian 已提交
4878
        """
4879 4880 4881 4882 4883
        device = (
            op.attr(self._op_device_key)
            if op.has_attr(self._op_device_key)
            else None
        )
4884
        if device:
4885 4886
            assert device[0:3] == 'gpu' or device[0:3] == 'npu', (
                "Now, only gpu and npu devices are "
4887
                "supported in pipeline parallemism."
4888
            )
4889 4890 4891 4892 4893 4894 4895 4896 4897 4898 4899 4900 4901
        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
4902
            op._set_attr(self._op_device_key, f"{self._device}:all")
4903 4904 4905 4906
        # 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():
4907 4908 4909
                assert (
                    '@RENAME@' in name
                ), "The op must be sum used to accumulate renamed vars."
4910 4911 4912 4913
            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(
4914 4915 4916 4917
                'op_device'
            ), "{} has no op_device attr for var {}".format(
                post_op.type, out_name
            )
4918 4919 4920
            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)
4921 4922 4923
        elif (op.type == "cast" or op.type == "scale") and self._is_backward_op(
            op
        ):
4924
            prev_op = self._find_prev_op(idx, op.desc.input("X")[0])
4925 4926
            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):
4927
            # for checkpoint offloading
4928 4929 4930
            assert (
                len(op.input_arg_names) == 1 and len(op.output_arg_names) == 1
            )
4931 4932 4933
            input_name = op.input_arg_names[0]
            output_name = op.output_arg_names[0]
            if '@Fetch' in output_name:
4934
                post_op = self._find_post_op(idx, output_name)
4935 4936 4937
                op._set_attr(
                    self._op_device_key, post_op.attr(self._op_device_key)
                )
4938
            else:
4939
                prev_op = self._find_prev_op(idx, op.desc.input("X")[0])
4940 4941 4942
                op._set_attr(
                    self._op_device_key, prev_op.attr(self._op_device_key)
                )
4943 4944 4945
        elif self._is_loss_op(op):
            # For loss * loss_scaling op added by AMP
            offset = 1
4946 4947 4948
            while not block.ops[idx + offset].has_attr(
                self._op_device_key
            ) or not block.ops[idx + offset].attr(self._op_device_key):
4949 4950 4951 4952 4953 4954 4955 4956 4957
                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
4958
            param_name = self._strip_grad_suffix(grad_name[0])
4959 4960 4961 4962 4963
            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.
4964 4965
            assert self._op_role_var_key in op.attr_names, (
                "gradient_clip "
4966
                "and regularization ops must have op_role_var attribute."
4967
            )
4968
            op_role_var = op.attr(self._op_role_var_key)
4969 4970
            assert len(op_role_var) == 2, (
                "op_role_var for gradient_clip "
4971
                "regularization ops must have two elements."
4972
            )
4973 4974
            param_name = op_role_var[0]
            device = self._param_device_map[param_name]
4975
            # For sum op added by global gradient clip, it must be
4976
            # put on all devices
4977 4978 4979 4980 4981 4982 4983
            if (
                op.type == 'sum'
                or op.type == 'sqrt'
                or op.type == 'fill_constant'
                or op.type == 'elementwise_max'
                or op.type == 'elementwise_div'
            ):
4984
                device = f"{self._device}:all"
4985
            op._set_attr(self._op_device_key, device)
R
Roc 已提交
4986
        elif op.type == "alloc_float_status" or op.type == "clear_float_status":
4987
            op._set_attr(self._op_device_key, f"{self._device}:all")
4988 4989 4990 4991 4992 4993 4994 4995 4996 4997
            # 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
4998 4999
        else:
            other_known_ops = [
5000 5001 5002 5003 5004 5005
                'update_loss_scaling',
                'reduce_any',
                'concat',
                'sum',
                'check_finite_and_unscale',
                'memcpy',
5006
            ]
5007 5008 5009
            assert op.type in other_known_ops, (
                "For other ops without "
                "op_device set, they must be one of {}, but it "
5010
                "is {}".format(other_known_ops, op.type)
5011
            )
5012
            assert self._is_optimize_op(op)
5013
            op._set_attr(self._op_device_key, f"{self._device}:all")
5014 5015

    def _add_op_device_attr(self, block):
5016
        """
5017
        Add op_device attrribute for ops in block that have
5018
        not that attribute set.
5019
        """
5020
        for idx, op in enumerate(list(block.ops)):
5021 5022 5023 5024 5025
            if (
                op.type == "create_py_reader"
                or op.type == "read"
                or op.type == "create_double_buffer_reader"
            ):
5026
                # Copy read related ops to all section to make them exit
5027 5028 5029 5030
                # 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.
5031
                op._set_attr(self._op_device_key, f"{self._device}:all")
5032 5033
                continue
            # op_device attribute has been set
5034 5035
            if self._get_op_device_attr(op):
                continue
5036
            self._add_op_device_attr_for_op(op, idx, block)
H
hutuxian 已提交
5037

5038 5039
    def _check_validation(self, block):
        """
5040
        Check whether ops in a block have both the op_device and the
5041 5042
        op_role attributes set.
        Then, return all devices in order.
5043
        """
5044 5045 5046 5047 5048 5049 5050 5051 5052 5053
        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),
        ]
5054
        for op in block.ops:
5055
            if not op._has_kernel(op.type):
5056 5057 5058 5059 5060 5061
                assert op.type == "conditional_block" and (
                    op.attr(self._op_role_key) == int(self._op_role.LRSched)
                ), (
                    "Now, the only supported op without kernel is "
                    "conditional_block, and its op role must be LRSched."
                )
5062
            assert op.has_attr(
5063 5064
                self._op_role_key
            ), "op ({}) has no {} attribute.".format(op.type, self._op_role_key)
5065
            op_role = op.attr(self._op_role_key)
5066 5067 5068 5069 5070
            assert (
                int(op_role) in valid_op_role_value
            ), "op_role {} for op {} must be one of {}".format(
                op_role, op.type, valid_op_role_value
            )
5071

5072
            assert op.has_attr(
5073 5074 5075 5076
                self._op_device_key
            ), "op ({}) has no {} attribute.".format(
                op.type, self._op_device_key
            )
5077 5078

            device = op.attr(self._op_device_key)
5079 5080 5081 5082 5083 5084 5085
            assert (
                device
            ), "op_device attribute for op " "{} has not been set.".format(
                op.type
            )
            if device == f"{self._device}:all":
                continue
5086

5087
            dev_type = device.split(':')[0]
B
Baibaifan 已提交
5088 5089
            assert dev_type == "gpu" or dev_type == 'npu', (
                "Now only gpu and npu devices are supported "
5090 5091
                "for pipeline parallelism."
            )
5092 5093

            if device not in device_list:
5094
                device_list.append(device)
5095

5096
        return device_list
5097

5098
    def _insert_sendrecv_ops_for_boundaries(self, block):
5099
        """
5100
        Insert a pair of send and recv ops for every two
5101 5102
        consecutive ops on different devices.
        """
5103
        # A map from var to device where op takes it as input,
5104
        # avoiding multiple send and recv ops.
5105
        input_var_to_device = dict()
5106 5107 5108 5109 5110 5111 5112 5113
        # 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,
5114
            'first_optimize_index': first_optimize_index,
5115
        }
5116

5117
        for index, op in enumerate(list(block.ops)):
5118
            cur_device = op.attr(self._op_device_key)
5119 5120
            if cur_device == f"{self._device}:all":
                continue
5121 5122
            for var_name in op.input_arg_names:
                var = block.var(var_name)
5123
                # skip data var
5124 5125
                if var.is_data:
                    continue
5126
                prev_device = None
5127 5128 5129

                prev_op = self._find_prev_op(index, var_name)
                if prev_op is None:
5130 5131
                    if var_name not in self._param_device_map:
                        continue
5132
                    prev_device = self._param_device_map[var_name]
5133

5134
                if not prev_device:
5135 5136 5137
                    prev_device = (
                        prev_op.attr(self._op_device_key) if prev_op else None
                    )
5138

5139 5140
                if prev_device is None or prev_device == f"{self._device}:all":
                    continue
5141

5142 5143
                if prev_device == cur_device:
                    continue
5144

5145 5146 5147 5148 5149 5150 5151
                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] + ':'

5152 5153 5154 5155
                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)
5156 5157
                    assert is_forward or is_backward, (
                        'send/recv in pipeline should only be inserted in forward or backward,'
5158
                        'please check the op_role of op={}'.format(op)
5159
                    )
5160 5161

                    if is_forward:
5162 5163
                        assert prev_id < cur_id, (
                            "In forward, send/recv can only be passed forward, but now "
5164
                            "prev_stage={} great than cur_stage={}, please check op_device of op={}".format(
5165 5166 5167
                                prev_id, cur_id, op
                            )
                        )
5168
                    elif is_backward:
5169 5170
                        assert prev_id > cur_id, (
                            "In backward, send/recv can only be passed backward, but now "
5171
                            "prev_stage={} less than cur_stage={}, please check op_device of op={}".format(
5172 5173 5174
                                prev_id, cur_id, op
                            )
                        )
5175

5176 5177 5178 5179 5180 5181 5182 5183 5184 5185
                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(
5186 5187
                            (cur_dev, prev_dev)
                        )
5188 5189 5190 5191 5192
                        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(
5193 5194
                            (cur_dev, prev_dev)
                        )
5195 5196 5197 5198 5199 5200
                        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)
5201
                    var = block.vars[var_name]
5202 5203 5204
                    pair = (prev_id, cur_id)
                    # 1000 is just a magic number
                    pair_key = prev_id * 1000 + cur_id
5205 5206 5207 5208 5209 5210 5211
                    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]
5212

5213
                    if self.schedule_mode == 'F-then-B':  # F-then-B
F
fangshuixun007 已提交
5214
                        block._insert_op_without_sync(
5215
                            index=index + extra_index_info['index'],
5216 5217 5218
                            type='send_v2',
                            inputs={'X': var},
                            attrs={
5219
                                self._op_device_key: prev_dev,
5220 5221 5222
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 1,
5223 5224 5225
                                'ring_id': ring_id,
                            },
                        )
5226
                        extra_index_info['index'] += 1
5227
                        var_shape = list(var.shape)
5228 5229 5230 5231 5232
                        var_shape[0] = (
                            self.micro_batch_size
                            if var_shape[0] < 0
                            else var_shape[0]
                        )
F
fangshuixun007 已提交
5233
                        block._insert_op_without_sync(
5234
                            index=index + extra_index_info['index'],
5235 5236 5237
                            type='recv_v2',
                            outputs={'Out': [var]},
                            attrs={
5238
                                'out_shape': var_shape,
5239
                                'dtype': var.dtype,
5240
                                self._op_device_key: cur_dev,
5241 5242 5243
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 0,
5244 5245 5246
                                'ring_id': ring_id,
                            },
                        )
5247
                        extra_index_info['index'] += 1
5248
                    elif self.schedule_mode == '1F1B':  # 1F1B
5249
                        var_shape = list(var.shape)
5250 5251 5252 5253 5254
                        var_shape[0] = (
                            self.micro_batch_size
                            if var_shape[0] < 0
                            else var_shape[0]
                        )
5255

5256
                        numel = np.prod(var_shape)
5257 5258 5259
                        use_mp = (self.mp_degree > 1) and (
                            numel % self.mp_degree == 0
                        )
5260 5261 5262 5263 5264 5265 5266 5267 5268 5269 5270 5271 5272 5273 5274 5275 5276 5277 5278 5279 5280 5281

                        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,
5282 5283
                                },
                            )
5284 5285 5286
                            extra_index_info['index'] += 1
                            return

5287 5288
                        _check_stage(cur_id, prev_id)

5289 5290 5291 5292 5293 5294 5295 5296 5297 5298
                        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,
                            },
                        )
5299
                        extra_index_info['index'] += 1
5300 5301
                        prefix_name = var.name.split('@')[0]
                        prefix_var = block.var(prefix_name)
5302 5303 5304
                        is_param = (
                            True if isinstance(prefix_var, Parameter) else False
                        )
F
fangshuixun007 已提交
5305
                        block._insert_op_without_sync(
5306
                            index=index + extra_index_info['index'],
5307
                            type='send_v2'
5308 5309
                            if not use_mp or is_param
                            else 'partial_send',
5310 5311
                            inputs={'X': var},
                            attrs={
5312
                                self._op_device_key: prev_dev,
5313 5314 5315 5316
                                self._op_role_key: op_role,
                                'use_calc_stream': False,
                                'ring_id': ring_id,
                                'peer': 1,
5317 5318 5319
                                # if send_v2, num&id attr is not in op_attrs, will not insert
                                'num': self.mp_degree,
                                'id': self.mp_rank,
5320 5321
                            },
                        )
5322
                        extra_index_info['index'] += 1
5323 5324 5325
                        insert_index = None
                        if int(op_role) == int(self._op_role.Backward):
                            insert_index = extra_index_info[
5326 5327
                                'first_optimize_index'
                            ]
5328 5329 5330 5331
                            new_op_role = self._op_role.Optimize
                        else:
                            insert_index = index
                            new_op_role = self._op_role.Backward
5332
                        sync_comm_op = block._insert_op_without_sync(
5333
                            index=insert_index + extra_index_info['index'],
5334 5335 5336 5337
                            type='c_sync_comm_stream',
                            inputs={'X': [var]},
                            outputs={'Out': [var]},
                            attrs={
5338
                                self._op_device_key: prev_dev,
5339
                                self._op_role_key: new_op_role,
5340
                                'ring_id': ring_id,
5341 5342
                            },
                        )
5343
                        if int(op_role) == int(self._op_role.Forward):
5344
                            sync_comm_op._set_attr('pipeline_flag', '')
5345
                            extra_index_info['index'] += 1
F
fangshuixun007 已提交
5346
                        block._insert_op_without_sync(
5347
                            index=index + extra_index_info['index'],
5348
                            type='recv_v2'
5349 5350
                            if not use_mp or is_param
                            else 'partial_recv',
5351 5352 5353 5354
                            outputs={'Out': [var]},
                            attrs={
                                'out_shape': var_shape,
                                'dtype': var.dtype,
5355
                                self._op_device_key: cur_dev,
5356 5357 5358
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 0,
5359 5360 5361 5362
                                '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,
5363 5364
                            },
                        )
5365
                        extra_index_info['index'] += 1
5366
                        if use_mp and not is_param:
5367 5368 5369 5370 5371 5372 5373 5374 5375 5376 5377 5378 5379
                            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,
5380 5381
                                },
                            )
5382
                            extra_index_info['index'] += 1
5383 5384 5385
                    else:
                        raise ValueError(
                            "Now only 'F-then-B' and '1F1B' are supported."
5386 5387
                            "The given value is {}.".format(self.schedule_mode)
                        )
5388

5389 5390 5391 5392
                _insert_send_recv(
                    int(cur_device.split(':')[1]),
                    int(prev_device.split(':')[1]),
                )
5393 5394
        block._sync_with_cpp()

5395
    def _insert_loss_scale(self, block):
5396
        """
5397
        Scale the loss corresponding to number of micro-batches.
5398
        """
5399 5400
        if self._num_microbatches == 1:
            return
5401
        for index, op in reversed(tuple(enumerate(list(block.ops)))):
5402
            if self._is_loss_grad_op(op):
5403 5404
                assert op.type == 'fill_constant', (
                    "loss_grad_op must be fill_constant op, "
5405
                    "but this op is {}".format(op.type)
5406
                )
5407 5408 5409 5410
                assert op.has_attr('value')
                loss_scale = float(op.attr('value'))
                loss_scale = loss_scale / self._num_microbatches
                op._set_attr('value', loss_scale)
5411 5412
                break

5413 5414
    def _rename_gradient_var_name(self, block):
        for index, op in enumerate(block.ops):
5415 5416
            if not self._is_optimize_op(op):
                continue
5417 5418 5419
            input_names = op.input_arg_names
            output_names = op.output_arg_names
            in_out_names = input_names + output_names
5420 5421
            if op.type == 'cast' or op.type == "c_sync_comm_stream":
                continue
5422 5423 5424
            # 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:
5425 5426
                if not core.grad_var_suffix() in name:
                    continue
5427 5428 5429 5430
                param_name = name.strip(core.grad_var_suffix())
                new_grad_name = name + "@MERGED"
                self._rename_arg(op, name, new_grad_name)

5431 5432 5433
    def _accumulate_gradients(
        self, block, pp_allreduce_in_optimize=False, strategy=None, shard=None
    ):
5434 5435 5436 5437
        """
        Create a new merged gradient for each parameter and accumulate the
        corresponding gradient to it.
        """
5438 5439
        fp16_allreduce = strategy.fp16_allreduce if strategy else False
        if strategy and strategy.fuse_grad_merge:
5440
            fused_gradient_names = self._accumulate_gradients_with_fuse(
5441 5442
                block, fp16_allreduce, strategy.fuse_grad_size_in_MB, shard
            )
5443 5444
            return fused_gradient_names

5445 5446 5447
        merged_gradient_names = []
        first_opt_op_idx = None

5448 5449 5450
        merged_suffix = '@MERGED@FP16' if fp16_allreduce else '@MERGED'
        dtype = paddle.float16 if fp16_allreduce else None

5451 5452 5453 5454 5455 5456 5457 5458
        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)
5459
                    continue
5460

5461
            if self._is_backward_op(op) and first_opt_op_idx is None:
5462
                first_opt_op_idx = index + 1
5463 5464
                # maybe have no optimize
                # if first_opt_op_idx == len(block.ops): return
5465

5466 5467 5468
            if self._is_backward_op(op) and (
                self._op_role_var_key in op.attr_names
            ):
5469
                op_role_var = op.attr(self._op_role_var_key)
5470 5471
                if len(op_role_var) == 0:
                    continue
5472 5473
                assert len(op_role_var) % 2 == 0
                for i in range(0, len(op_role_var), 2):
5474 5475
                    offset = 0
                    param_name = op_role_var[i]
5476 5477 5478 5479
                    if not block.has_var(param_name):
                        continue
                    if '@BroadCast' in param_name:
                        continue
5480

5481
                    param_grad_name = param_name + core.grad_var_suffix()
5482
                    merged_param_grad_name = param_grad_name + merged_suffix
5483
                    if not block.has_var(merged_param_grad_name):
5484 5485 5486 5487 5488 5489
                        self._create_var(
                            block,
                            block.vars[param_name],
                            merged_param_grad_name,
                            dtype,
                        )
5490
                    assert block.has_var(merged_param_grad_name)
5491

5492 5493 5494
                    param_grad_var = block.var(param_grad_name)
                    merged_param_grad_var = block.var(merged_param_grad_name)
                    merged_param_grad_var.persistable = True
5495
                    block._insert_op(
5496 5497 5498 5499
                        index=first_opt_op_idx + offset,
                        type='fill_constant',
                        inputs={},
                        outputs={'Out': [merged_param_grad_var]},
5500
                        attrs={
5501 5502 5503
                            'shape': merged_param_grad_var.shape,
                            'dtype': merged_param_grad_var.dtype,
                            'value': float(0),
5504
                            # a trick to run this op once per mini-batch
5505 5506 5507
                            self._op_role_key: self._op_role.Optimize.LRSched,
                        },
                    )
5508
                    offset += 1
5509 5510
                    grad_name = op_role_var[i + 1]
                    grad_var = block.vars[grad_name]
5511 5512

                    is_fp16_grad = 'cast_fp16' in grad_name
5513
                    need_cast = is_fp16_grad is not fp16_allreduce
5514 5515 5516 5517 5518 5519

                    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
5520
                        cast_grad_var_name = param_grad_name + '@TMP'
5521
                        cast_grad_var = self._create_var(
5522 5523
                            block, param_grad_var, cast_grad_var_name, dtype
                        )
5524
                        cast_grad_var.persistable = False
5525 5526 5527 5528 5529 5530 5531 5532 5533 5534 5535
                        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,
                            },
                        )
5536
                        offset += 1
5537 5538 5539 5540 5541 5542 5543
                        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},
5544 5545
                        attrs={
                            self._op_role_key: self._op_role.Backward,
5546 5547
                        },
                    )
5548 5549 5550
                    offset += 1
                    merged_gradient_names.append(merged_param_grad_name)

5551 5552
        if not fp16_allreduce:
            return merged_gradient_names
5553 5554 5555 5556 5557 5558 5559 5560 5561 5562 5563 5564 5565 5566 5567 5568 5569 5570 5571 5572 5573 5574 5575

        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

5576 5577 5578 5579 5580 5581 5582 5583 5584 5585 5586
            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,
                },
            )
5587

5588
        return merged_gradient_names
5589

5590 5591 5592
    def _insert_accumulate_gradients_with_fuse(
        self, main_block, fp16, fused_size, grad_param_pairs, first_opt_op_idx
    ):
5593
        grad_param_pairs = self._sort_grad_param_by_dtype(
5594 5595
            main_block, grad_param_pairs
        )
5596

5597 5598 5599
        grad_param_segments = []
        merged_suffix = '@MERGED@FP16' if fp16 else '@MERGED'
        dtype = paddle.float16 if fp16 else paddle.float32
5600
        cur_size = 0.0
5601 5602 5603 5604 5605 5606 5607 5608 5609 5610
        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,
5611 5612
                stop_gradient=False,
            )
5613
            real_param = main_block.var(param)
5614 5615
            if hasattr(real_param, 'is_distributed'):
                merged_grad_var.is_distributed = real_param.is_distributed
5616 5617 5618 5619
            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
5620 5621 5622 5623 5624
            if (
                len(grad_param_segments) == 0
                or cur_size + tmp_size > fused_size
                or real_grad.dtype != last_dtype
            ):
5625
                grad_param_segments.append(
5626 5627
                    ([real_grad], [real_param], [merged_grad_var])
                )
5628
                last_dtype = real_grad.dtype
5629
                cur_size = 0.0
5630 5631 5632 5633 5634 5635 5636 5637 5638 5639 5640 5641
            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]
5642 5643 5644 5645 5646 5647
            fused_grad = main_block.create_var(
                name='FusedGrad_{}'.format(grad_segment[0].name),
                dtype=grad_segment[0].dtype,
                persistable=False,
                stop_gradient=False,
            )
5648
            # keep the '.cast_fp16' info in the fuse var name
5649 5650 5651 5652 5653 5654 5655 5656 5657
            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)
            )
5658 5659 5660 5661
            fused_merged_grad = main_block.create_var(
                name=fused_merged_grad_name,
                dtype=merged_grad_segment[0].dtype,
                persistable=True,
5662 5663
                stop_gradient=False,
            )
5664 5665 5666 5667 5668 5669 5670 5671 5672 5673 5674 5675 5676 5677 5678 5679 5680 5681 5682 5683 5684 5685 5686 5687 5688
            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},
5689
                outputs={"Output": grads, "FusedOutput": fused_grad},
5690 5691 5692 5693 5694 5695 5696 5697 5698 5699 5700 5701 5702 5703 5704 5705
                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,
5706 5707 5708 5709 5710 5711 5712
                    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),
5713 5714
                },
            )
5715 5716 5717 5718 5719 5720 5721 5722 5723 5724
            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,
5725
                    "FusedOutput": fused_merged_grad,
5726 5727 5728 5729 5730 5731 5732 5733
                },
                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,
5734 5735 5736
                    self._op_role_key: self._op_role.Optimize.LRSched,
                },
            )
5737 5738 5739 5740 5741 5742 5743 5744 5745
            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
5746
            need_cast = is_fp16_grad is not fp16
5747 5748 5749 5750
            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'
5751 5752 5753 5754 5755 5756 5757 5758 5759 5760 5761 5762 5763 5764 5765 5766 5767
                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,
                    },
                )
5768 5769 5770 5771 5772 5773 5774
                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},
5775 5776
                attrs={self._op_role_key: self._op_role.Backward},
            )
5777 5778 5779 5780 5781 5782 5783 5784 5785 5786
            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'
5787 5788 5789 5790 5791 5792 5793 5794 5795 5796 5797 5798 5799 5800 5801 5802 5803 5804
                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,
                    },
                )
5805 5806 5807 5808 5809 5810
                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

5811
        return fused_merged_gradients, first_opt_op_idx
5812

5813 5814 5815
    def _accumulate_gradients_with_fuse(
        self, main_block, fp16, fused_size, shard=None
    ):
5816 5817 5818 5819 5820 5821 5822 5823 5824 5825 5826 5827 5828 5829 5830 5831 5832 5833 5834
        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

5835 5836 5837
            if self._is_backward_op(op) and (
                self._op_role_var_key in op.attr_names
            ):
5838 5839 5840 5841 5842 5843 5844 5845 5846 5847 5848
                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(
5849 5850
                        (op_role_var[i + 1], op_role_var[i])
                    )
5851 5852 5853 5854 5855 5856 5857 5858 5859 5860 5861 5862 5863

        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:
5864 5865 5866 5867 5868 5869
            (
                fused_merged_gradients,
                first_opt_op_idx,
            ) = self._insert_accumulate_gradients_with_fuse(
                main_block, fp16, fused_size, pairs, first_opt_op_idx
            )
5870 5871 5872 5873
            all_fused_merged_gradients += fused_merged_gradients

        main_block._sync_with_cpp()
        return all_fused_merged_gradients
5874

5875 5876 5877 5878 5879 5880 5881 5882 5883 5884 5885 5886 5887 5888 5889 5890 5891 5892
    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

5893 5894 5895 5896 5897 5898 5899 5900 5901 5902 5903 5904
    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
5905 5906 5907 5908 5909 5910
        return (
            reduce(lambda x, y: x * y, var.shape)
            * dtype_to_size[var.dtype]
            / 1024.0
            / 1024.0
        )
5911

5912 5913
    def _add_sub_blocks(self, main_block, program_list):
        main_program = main_block.program
5914
        for prog in program_list:
5915 5916 5917 5918 5919 5920
            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)
5921 5922
                for sub_op in origin_sub_block.ops:
                    op_desc = sub_op.desc
5923 5924 5925
                    ap_op = new_sub_block.desc.append_op()
                    ap_op.copy_from(op_desc)
                new_sub_block._sync_with_cpp()
5926
                self._create_vars(new_sub_block, origin_sub_block)
5927
                op._set_attr('sub_block', new_sub_block)
5928 5929 5930

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

5936 5937 5938
    def _process_persistable_vars_in_multi_sections(
        self, main_program, startup_prog, program_list
    ):
5939 5940 5941 5942 5943 5944 5945
        """
        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()
5946
        for prog in program_list:
5947 5948
            block = prog.block(0)
            for var_name in block.vars:
5949 5950
                if var_name == "double_buffer_0":
                    continue
5951
                var = block.var(var_name)
5952 5953
                if not var.persistable:
                    continue
5954 5955 5956 5957 5958 5959 5960 5961 5962 5963 5964 5965 5966 5967 5968
                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:
5969 5970 5971 5972 5973 5974
                    if (
                        op.type == "recv_v2"
                        or op.type == "create_py_reader"
                        or op.type == "read"
                        or op.type == "update_loss_scaling"
                    ):
5975
                        continue
5976 5977
                    # We have processed lr related vars
                    if op.attr(self._op_role_key) == int(
5978 5979
                        self._op_role.Optimize.LRSched
                    ):
5980 5981 5982 5983
                        continue
                    if var_name in op.desc.output_arg_names():
                        assert var_name not in write_info, (
                            "two sections write the same var({}): second "
5984 5985
                            "op {}.".format(var_name, op)
                        )
5986 5987 5988 5989 5990
                        write_info[var_name] = prog
                        break

        for var_name in var_info.keys():
            # Case 1: read only variables, no special process
5991 5992
            if not var_name in write_info:
                continue
5993 5994 5995 5996 5997

            # 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)
5998
            write_dev_index = int(write_device.split(':')[1])
5999 6000
            all_progs = var_info[var_name]
            for prog in all_progs:
6001 6002
                if prog == write_prog:
                    continue
6003 6004 6005
                read_block = prog.block(0)
                read_device = self._get_device_info(read_block)
                read_dev_index = int(read_device.split(':')[1])
6006 6007 6008 6009 6010 6011 6012 6013 6014
                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]
6015 6016 6017

                write_block._insert_op(
                    index=0,
6018
                    type='send_v2',
6019 6020 6021
                    inputs={
                        'X': write_block.var(var_name),
                    },
6022
                    attrs={
6023 6024
                        self._op_device_key: write_device,
                        'use_calc_stream': False,
6025 6026
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
6027 6028 6029 6030 6031
                        self._op_role_key: self._op_role.LRSched,
                        'peer': read_dev_index,
                        'ring_id': ring_id,
                    },
                )
6032 6033
                read_block._insert_op(
                    index=0,
6034
                    type='recv_v2',
6035 6036
                    outputs={'Out': [read_block.var(var_name)]},
                    attrs={
6037 6038 6039 6040
                        '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,
6041 6042
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
6043 6044 6045 6046 6047
                        self._op_role_key: self._op_role.LRSched,
                        'peer': write_dev_index,
                        'ring_id': ring_id,
                    },
                )
6048 6049 6050 6051 6052 6053
                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={
6054
                        self._op_device_key: read_device,
6055 6056
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
6057 6058 6059 6060
                        self._op_role_key: self._op_role.LRSched,
                        'ring_id': ring_id,
                    },
                )
6061 6062

    def _is_gradient_clip_op(self, op):
6063 6064 6065
        return op.desc.has_attr("op_namescope") and op.desc.attr(
            "op_namescope"
        ).startswith("/gradient_clip")
6066 6067

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

6072 6073
    def _is_weight_decay_op(self, op):
        # in AdamW namescope is /optimizer_*/weight decay/
6074 6075 6076
        return op.desc.has_attr(
            "op_namescope"
        ) and 'weight decay' in op.desc.attr("op_namescope")
6077

6078 6079 6080 6081 6082
    def _get_input_output_info(self, block):
        '''
        Get info of op input and output.
        '''
        # A map from output var to op which generate it.
6083
        output_var_to_op = defaultdict(list)
6084
        # A map from var to op which takes it as input.
6085
        input_var_to_op = defaultdict(list)
6086

6087
        for index, op in enumerate(block.ops):
6088
            for var_name in op.input_arg_names:
6089
                input_var_to_op[var_name].append([op, index])
6090
            for var_name in op.output_arg_names:
6091 6092 6093 6094 6095 6096 6097 6098
                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
        """
6099 6100
        if self.schedule_mode != '1F1B':
            return
6101 6102 6103

        block = program.block(0)

6104
        recv_type = 'recv_v2' if self.mp_degree == 1 else 'partial_recv'
6105 6106
        backward_recv_index = None
        for index, op in enumerate(block.ops):
6107
            if op.type == recv_type and self._is_backward_op(op):
6108 6109 6110
                backward_recv_index = index
                break

6111
        # last pipeline stage
6112 6113
        if backward_recv_index is None:
            return
6114 6115 6116

        offset = 0
        for index, op in enumerate(list(block.ops)):
6117 6118
            if index >= backward_recv_index:
                break
6119 6120 6121 6122 6123 6124 6125 6126 6127 6128 6129 6130 6131 6132 6133 6134
            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]},
6135 6136
                    attrs={self._op_role_key: self._op_role.Backward},
                )
6137
        block._sync_with_cpp()
6138

6139 6140 6141 6142 6143 6144 6145 6146 6147 6148 6149 6150 6151
    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))
6152 6153 6154 6155
            if (
                op_role == int(self._op_role.Backward)
                and backward_insert_index is None
            ):
6156
                backward_insert_index = i
6157 6158 6159 6160 6161 6162
            if (
                op.type != "partial_recv"
                and op.type != "partial_allgather"
                and op.type != "nop"
                and op.type != "recv_v2"
            ):
6163 6164 6165 6166 6167 6168 6169 6170 6171 6172 6173 6174 6175 6176 6177 6178 6179 6180 6181
                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)
6182 6183 6184 6185 6186 6187 6188
            block._insert_op_without_sync(
                index=insert_index,
                type=op.type,
                inputs=op_inputs,
                outputs=op_outputs,
                attrs=op.all_attrs(),
            )
6189 6190 6191 6192 6193 6194 6195
            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()

6196 6197 6198 6199 6200 6201 6202 6203 6204 6205 6206 6207 6208 6209 6210 6211 6212 6213 6214 6215 6216 6217 6218 6219 6220 6221 6222
    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 "
6223 6224
            "forward and used in backward:\n{}".format(used_in_backward)
        )
6225

6226 6227 6228
    def minimize(
        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
    ):
6229
        main_block = loss.block
6230
        self.origin_main_block = main_block
6231
        main_program = main_block.program
6232 6233
        if startup_program is None:
            startup_program = default_startup_program()
6234

6235 6236
        pipeline_opt = main_program._pipeline_opt
        assert pipeline_opt, 'Please use pipeline with fleet.'
6237 6238 6239 6240 6241 6242 6243
        required_keys = [
            'local_rank',
            'schedule_mode',
            'micro_batch_size',
            'ring_id',
            'global_ring_id',
            'use_sharding',
6244 6245
            'mp_degree',
            'mp_rank',
6246 6247
        ]
        for key in required_keys:
6248 6249 6250
            assert (
                key in pipeline_opt
            ), 'Please use pipeline with fleet to use {}.'.format(key)
6251 6252 6253 6254 6255 6256 6257 6258
        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']
6259
        self.scale_gradient = pipeline_opt.get('scale_gradient', False)
6260 6261
        assert self.mp_degree >= 1
        assert 0 <= self.mp_rank < self.mp_degree
6262 6263

        optimize_ops, params_grads = self._optimizer.minimize(
6264 6265
            loss, startup_program, parameter_list, no_grad_set
        )
6266
        self._param_device_map = self._origin_optimizer._param_device_map
6267

6268 6269 6270 6271
        (
            self.output_var_to_op,
            self.input_var_to_op,
        ) = self._get_input_output_info(main_block)
6272 6273 6274
        # Step1: add default op_device attribute for ops.
        self._add_op_device_attr(main_block)
        device_list = self._check_validation(main_block)
6275 6276 6277 6278 6279 6280 6281 6282 6283 6284 6285

        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

6286 6287 6288
        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 "
6289 6290
            "another in the order of their ids."
        )
6291
        # Step2: add send and recv ops between section boundaries
6292
        self._insert_sendrecv_ops_for_boundaries(main_block)
6293

6294
        # Step3: split program into sections and add pairs of
6295 6296
        # send and recv ops for data var.
        main_program = main_block.program
6297
        program_list = self._split_program(main_program, device_list)
6298
        for p in program_list:
6299
            self._create_vars(p.global_block(), main_block)
6300

L
lilong12 已提交
6301 6302 6303 6304 6305
        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 "
6306 6307
                "stages."
            )
L
lilong12 已提交
6308 6309
        else:
            self.local_rank %= len(device_list)
6310 6311 6312
        # Step3.5: optimize forward send sync_comm to overlap send and recv
        self._optimize_forward_send_sync(program_list[self.local_rank])

6313
        # Step4: Special Case: process persistable vars that exist in
6314
        # multiple sections
6315
        # FIXME
6316 6317
        # self._process_persistable_vars_in_multi_sections(
        #     main_program, startup_program, program_list)
6318

6319
        # Step5: Add sub blocks for section programs
6320 6321
        self._add_sub_blocks(main_block, program_list)

6322
        place_list = []
6323 6324
        for dev in device_list:
            dev_index = int(dev.split(":")[1])
6325 6326 6327 6328
            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))
6329

6330
        # Step6: Split startup program
6331
        new_startup_program = self._split_startup_program(
6332 6333
            startup_program, self.local_rank
        )
6334 6335 6336 6337

        startup_program._pipeline_opt = {
            "startup_program": new_startup_program,
        }
6338
        real_block = program_list[self.local_rank].global_block()
6339 6340
        if not self.scale_gradient:
            self._insert_loss_scale(real_block)
6341
        if not self.use_sharding:
6342
            # Step7: clear gradients before each mini-batch and
6343 6344 6345 6346 6347
            # 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()
6348

6349 6350 6351 6352
        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"))
6353 6354 6355
        # A pass to move the recv op to the beginning of
        # the forward/backward phase
        self._mv_head_recv(program_list[self.local_rank])
6356 6357 6358 6359 6360

        # 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])

6361
        main_program._pipeline_opt = {
H
hutuxian 已提交
6362 6363
            "trainer": "PipelineTrainer",
            "device_worker": "Section",
6364
            "pipeline_stage": self.local_rank,
6365
            "num_pipeline_stages": len(device_list),
6366
            "schedule_mode": self.schedule_mode,
6367
            "inner_parallelism": len(device_list),
6368 6369
            "section_program": program_list[self.local_rank],
            "place": place_list[self.local_rank],
6370
            "place_id": place_id,
6371
            "sync_steps": -1,
L
lilong12 已提交
6372
            "num_microbatches": self._num_microbatches,
H
hutuxian 已提交
6373 6374
            "start_cpu_core_id": self._start_cpu_core_id,
        }
6375 6376 6377 6378 6379 6380 6381
        return (
            optimize_ops,
            params_grads,
            program_list,
            self._pipeline_pair,
            self._pp_ring_map,
        )
M
mapingshuo 已提交
6382 6383


M
mapingshuo 已提交
6384 6385
class RecomputeOptimizer(Optimizer):
    """
6386
        :api_attr: Static Graph
S
swtkiwi 已提交
6387

M
mapingshuo 已提交
6388 6389 6390
    Recompute Optimizer Wrapper

    Normally, a training step contains three sub-steps: first, run forward
6391
    Operators to calculate the loss; second, run backward Operators to
M
mapingshuo 已提交
6392 6393 6394
    calculate gradient of the parameters; third, apply optimization method
    to update the value of the parameters.

6395
    In the forward computation process, all variables that are needed by
M
mapingshuo 已提交
6396 6397 6398
    backward computation process will be kept in memory, which occupy a great
    amount of memory when the network becomes very deep.

6399
    Recompute split the network to k segments. In each segment, It will
M
mapingshuo 已提交
6400 6401
    recompute the forward Operators, before running backward operators. It is
    very helpful for saving memory.
6402

M
mapingshuo 已提交
6403 6404 6405 6406 6407 6408 6409 6410 6411 6412 6413 6414 6415 6416 6417 6418 6419 6420 6421 6422 6423 6424 6425 6426 6427 6428 6429 6430 6431 6432 6433 6434 6435 6436 6437 6438 6439 6440 6441 6442 6443 6444 6445 6446 6447
    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 已提交
6448
        if framework._non_static_mode():
Z
zhongpu 已提交
6449
            raise Exception("In dygraph, don't support RecomputeOptimizer.")
M
mapingshuo 已提交
6450 6451
        self._optimizer = optimizer
        self._checkpoints = None
M
mapingshuo 已提交
6452 6453
        self._learning_rate = self._optimizer._learning_rate
        self._learning_rate_map = self._optimizer._learning_rate_map
J
JZ-LIANG 已提交
6454
        self.enable_offload = False
M
mapingshuo 已提交
6455 6456

    def _set_checkpoints(self, checkpoints):
6457 6458
        """
        Args:
6459
            checkpoints (list): List of Variable or string
6460 6461 6462 6463 6464
        """
        assert isinstance(
            checkpoints, list
        ), "_checkpoints should be a list of Variable or a list of String"
        for ckpt in checkpoints:
6465 6466
            assert isinstance(ckpt, str) or isinstance(
                ckpt, Variable
6467
            ), "_checkpoints should be a list of Variable or a list of String"
M
mapingshuo 已提交
6468 6469
        self._checkpoints = checkpoints

6470
    # should enable offload before calling backward
J
JZ-LIANG 已提交
6471 6472 6473
    def _enable_offload(self):
        self.enable_offload = True

6474 6475
    @framework.deprecate_stat_dict
    def load(self, state_dict):
M
mapingshuo 已提交
6476
        """
6477
            :api_attr: Static Graph
S
swtkiwi 已提交
6478

M
mapingshuo 已提交
6479 6480 6481 6482
        load function is not supported by Recompute Optimizer for now.
        :return: None

        Args:
6483
            state_dict: the dict load by load_persistable method
M
mapingshuo 已提交
6484 6485 6486 6487

        Examples:
            .. code-block:: python

6488
                import paddle
M
mapingshuo 已提交
6489
                import paddle.fluid as fluid
6490

6491
                paddle.enable_static()
M
mapingshuo 已提交
6492 6493 6494 6495 6496 6497
                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
6498

M
mapingshuo 已提交
6499 6500 6501 6502
                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")
6503

M
mapingshuo 已提交
6504 6505 6506 6507
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
                sgd._set_checkpoints([fc_1, pred])
                try:
6508 6509
                    state_dict = {}
                    sgd.load(state_dict)
M
mapingshuo 已提交
6510
                except NotImplementedError as e:
6511
                    print(e)
M
mapingshuo 已提交
6512 6513
        """
        raise NotImplementedError(
6514 6515
            "load function is not supported by Recompute Optimizer for now"
        )
M
mapingshuo 已提交
6516 6517 6518 6519 6520 6521 6522 6523 6524 6525 6526 6527 6528 6529 6530 6531 6532 6533 6534 6535 6536 6537 6538 6539 6540 6541 6542 6543 6544 6545 6546 6547

    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)
6548
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
6549 6550 6551 6552
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
6553
                    no_grad_set=None)
M
mapingshuo 已提交
6554 6555 6556 6557 6558 6559 6560 6561 6562 6563

                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 已提交
6564 6565 6566 6567 6568 6569 6570 6571 6572
    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,
6573 6574
            stop_gradient=True,
        )
J
JZ-LIANG 已提交
6575 6576 6577 6578 6579 6580

        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,
6581 6582
            stop_gradient=False,
        )
J
JZ-LIANG 已提交
6583 6584 6585 6586 6587 6588 6589 6590

        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
6591 6592 6593
        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.
J
JZ-LIANG 已提交
6594 6595 6596 6597 6598 6599 6600 6601 6602 6603 6604 6605 6606
        """
        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,
6607 6608 6609 6610 6611 6612 6613 6614 6615 6616 6617 6618 6619
                stop_gradient=True,
            )
            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 已提交
6620 6621 6622

        return

6623 6624 6625
    def _insert_async_memcpy_op(
        self, insert_idx, src_varname, dst_varname, op_role, dst_place_type
    ):
J
JZ-LIANG 已提交
6626 6627 6628 6629 6630 6631 6632 6633
        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)]
            },
6634 6635
            attrs={"dst_place_type": int(dst_place_type), OP_ROLE_KEY: op_role},
        )
J
JZ-LIANG 已提交
6636 6637

    def _insert_fetch_op(self, idx, varname):
6638 6639 6640 6641 6642
        assert (
            varname in self.checkpoint_name2pinned_name
        ), "Try to fetch {} from Pinned Memory, but it is NOT a checkpoint".format(
            varname
        )
J
JZ-LIANG 已提交
6643 6644 6645

        pinned_varname = self.checkpoint_name2pinned_name[varname]
        fetch_varname = self.checkpoint_name2fetch_name[varname]
6646
        self._insert_async_memcpy_op(idx, pinned_varname, fetch_varname, 1, 1)
J
JZ-LIANG 已提交
6647 6648

    def _insert_offload_op(self, idx, varname):
6649 6650 6651 6652 6653
        assert (
            varname in self.checkpoint_name2pinned_name
        ), "Try to offload {} to Pinned Memory, but it is NOT a checkpoint".format(
            varname
        )
J
JZ-LIANG 已提交
6654
        pinned_varname = self.checkpoint_name2pinned_name[varname]
6655
        self._insert_async_memcpy_op(idx, varname, pinned_varname, 0, 2)
J
JZ-LIANG 已提交
6656 6657

    def _insert_sync_op(self, op_idx, checkpoint_name):
6658
        # single stream offload no need sync
J
JZ-LIANG 已提交
6659 6660 6661
        pass

    def _record_fetch_op(self, idx):
6662 6663 6664
        assert (
            len(self.un_fetch_checkpoint_names) > 0
        ), "Could NOT found checkpoint to fetch"
J
JZ-LIANG 已提交
6665 6666 6667 6668 6669 6670 6671 6672
        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)
6673 6674 6675 6676 6677
        assert (
            checkpoint_name == expected_checkpoint_name
        ), "expected to offload [{}] but got [{}]".format(
            expected_checkpoint_name, checkpoint_name
        )
J
JZ-LIANG 已提交
6678 6679 6680 6681
        logging.debug("Record offload [{}]".format(checkpoint_name))
        self.idx2insertions[idx] = ("offload", checkpoint_name)

    def _record_sync_op(self, idx, checkpoint_name):
6682 6683 6684
        assert (
            checkpoint_name not in self.synced_checkpoints
        ), "Try to sync the checkpoint [{}] twice".format(checkpoint_name)
J
JZ-LIANG 已提交
6685 6686 6687 6688 6689 6690 6691
        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 = {}
6692
        # don't offload the last checkpoints, to favor throughput
J
JZ-LIANG 已提交
6693 6694 6695 6696 6697 6698 6699 6700 6701 6702 6703 6704 6705 6706
        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(
6707 6708
            self.block.ops
        ), "Could NOT found backword op in prog"
J
JZ-LIANG 已提交
6709 6710 6711

        # fetch second to last checkpoint at the beginning of BW
        fetched_checkpoint_varname = self._record_fetch_op(
6712 6713
            self.bw_strart_op_idx
        )
J
JZ-LIANG 已提交
6714 6715
        last_last_fetch_checkpoint = None

6716
        for i, op in enumerate(self.block.ops[self.bw_strart_op_idx :]):
J
JZ-LIANG 已提交
6717 6718 6719 6720 6721 6722 6723 6724 6725
            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
6726 6727 6728
                            second_to_last_fetch_checkpoint = (
                                fetched_checkpoint_varname
                            )
6729
                            # there is NO fetch ahead the first checkpoint
J
JZ-LIANG 已提交
6730
                            if input_var != self.sorted_checkpoint_names[0]:
6731 6732 6733
                                fetched_checkpoint_varname = (
                                    self._record_fetch_op(idx)
                                )
J
JZ-LIANG 已提交
6734

6735
                        # should check the current used checkpoint is ths last fetch one
6736 6737 6738 6739 6740
                        assert (
                            second_to_last_fetch_checkpoint == input_var
                        ), "Current recompute segment should use [{}] BUT got [{}]".format(
                            second_to_last_fetch_checkpoint, input_var
                        )
J
JZ-LIANG 已提交
6741 6742 6743
                        # rename
                        self.block.ops[idx]._rename_input(
                            input_var,
6744 6745
                            self.checkpoint_name2fetch_name[input_var],
                        )
J
JZ-LIANG 已提交
6746 6747 6748 6749
                        self.checkpoint_usage_count[input_var] += 1
                    else:
                        raise ValueError(
                            "use checkpoint [{}] before fetch in BW".format(
6750 6751 6752
                                input_var
                            )
                        )
J
JZ-LIANG 已提交
6753

6754 6755 6756 6757 6758
        assert (
            len(self.un_fetch_checkpoint_names) == 0
        ), "{} checkpoints have NOT been Recorded".format(
            self.un_fetch_checkpoint_names
        )
J
JZ-LIANG 已提交
6759 6760 6761 6762 6763 6764 6765 6766 6767 6768

    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)
6769
                    logging.debug(
6770 6771
                        "Insert [{}] fetch op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
6772 6773 6774 6775 6776
                    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()
6777 6778 6779 6780 6781
        assert (
            len(self.idx2insertions) == 0
        ), "{} checkpoints left un-Fecthed".format(
            [ele[1] for ele in self.idx2insertions.values()]
        )
J
JZ-LIANG 已提交
6782 6783 6784 6785

    def _parse_forward(self):

        self.idx2insertions = {}
6786
        # don't offload the last checkpoints, faster, less memory saving
J
JZ-LIANG 已提交
6787 6788 6789 6790 6791 6792 6793
        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,
6794
                'idx': -1,
J
JZ-LIANG 已提交
6795 6796 6797 6798 6799 6800 6801 6802 6803
            }
        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(
6804 6805
            self.block.ops
        ), "Could NOT found Forward op in prog"
J
JZ-LIANG 已提交
6806 6807
        last_offload_checkpoint = None

6808
        for i, op in enumerate(
6809 6810
            self.block.ops[self.fw_strart_op_idx : self.bw_strart_op_idx]
        ):
J
JZ-LIANG 已提交
6811 6812 6813 6814 6815 6816 6817

            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:
6818 6819 6820 6821 6822
                    assert (
                        len(output_vars) == 1
                    ), "chekpoint should be the only Output of a certain op, but [{}] is from [{}]".format(
                        output_var, op
                    )
J
JZ-LIANG 已提交
6823 6824 6825

                    if output_var in self.un_offload_checkpoint_names:
                        # insert sync op if last checkpoint has not been sync
6826
                        if last_offload_checkpoint is not None:
6827 6828 6829 6830 6831 6832 6833 6834 6835
                            if (
                                self.checkpoint_usage_count_and_idx[
                                    last_offload_checkpoint
                                ]['count']
                                == 0
                            ):
                                self._record_sync_op(
                                    idx, last_offload_checkpoint
                                )
J
JZ-LIANG 已提交
6836
                            else:
6837 6838 6839 6840 6841 6842 6843 6844 6845 6846 6847 6848 6849
                                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
                                )
J
JZ-LIANG 已提交
6850 6851 6852 6853 6854
                        # 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(
6855 6856 6857 6858
                            "There should be just ONE op that output checkpoint [{}]".format(
                                output_var
                            )
                        )
J
JZ-LIANG 已提交
6859 6860
                # need to sync the last need to offload checkpoint before the last checkpoint as output op
                if output_var == last_checkpoint:
6861 6862 6863 6864 6865 6866 6867 6868 6869 6870 6871 6872 6873
                    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,
                    )
J
JZ-LIANG 已提交
6874
                    # sync if last checkpoint has not been sync
6875 6876 6877 6878 6879 6880
                    if (
                        self.checkpoint_usage_count_and_idx[
                            last_offload_checkpoint
                        ]['idx']
                        == 0
                    ):
J
JZ-LIANG 已提交
6881 6882 6883
                        self._record_sync_op(idx, last_offload_checkpoint)
                    else:
                        last_usage_idx = self.checkpoint_usage_count_and_idx[
6884 6885 6886 6887 6888 6889 6890 6891 6892 6893
                            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
                        )
6894
            # record checkpoint usage
J
JZ-LIANG 已提交
6895 6896
            for input_var in input_vars:
                if input_var in need_offload_checkpoint_names:
6897 6898 6899
                    assert (
                        input_var not in self.synced_checkpoints
                    ), "checkpoint [{}] used after sync".format(input_var)
J
JZ-LIANG 已提交
6900 6901 6902
                    self.checkpoint_usage_count_and_idx[input_var]['count'] += 1
                    self.checkpoint_usage_count_and_idx[input_var]['idx'] = idx

6903 6904 6905 6906 6907
        assert (
            len(self.un_offload_checkpoint_names) == 0
        ), "{} checkpoints have NOT been Recorded".format(
            self.un_fetch_checkpoint_names
        )
J
JZ-LIANG 已提交
6908 6909 6910
        assert len(self.synced_checkpoints) == len(
            need_offload_checkpoint_names
        ), "{} checkpoints have NOT been Recorded".format(
6911 6912
            set(need_offload_checkpoint_names) - set(self.synced_checkpoints)
        )
J
JZ-LIANG 已提交
6913 6914 6915 6916 6917

    def _update_forward(self):
        if len(self.idx2insertions) == 0:
            return
        for op_idx in reversed(
6918 6919
            range(self.fw_strart_op_idx, self.bw_strart_op_idx)
        ):
J
JZ-LIANG 已提交
6920 6921 6922 6923
            if op_idx in self.idx2insertions:
                operation, checkpoint_name = self.idx2insertions[op_idx]
                if operation == "offload":
                    self._insert_offload_op(op_idx, checkpoint_name)
6924
                    logging.debug(
6925 6926
                        "Insert [{}] offload op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
6927 6928 6929
                    del self.idx2insertions[op_idx]
                elif operation == "sync":
                    self._insert_sync_op(op_idx, checkpoint_name)
6930
                    logging.debug(
6931 6932
                        "Insert [{}] offload_sync op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
6933 6934 6935
                    del self.idx2insertions[op_idx]

        self.block._sync_with_cpp()
6936 6937 6938 6939 6940
        assert (
            len(self.idx2insertions) == 0
        ), "{} checkpoints left un-Offloaded".format(
            [ele[1] for ele in self.idx2insertions.values()]
        )
J
JZ-LIANG 已提交
6941 6942 6943 6944 6945 6946 6947 6948

    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
6949
        1. create pinned vars and temp vars
J
JZ-LIANG 已提交
6950 6951 6952 6953 6954 6955
        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
6956
        if startup_program is None:
J
JZ-LIANG 已提交
6957
            startup_program = paddle.static.default_startup_program()
J
JZ-LIANG 已提交
6958 6959

        with program_guard(self._main_program, startup_program):
6960 6961 6962 6963 6964 6965 6966 6967 6968 6969
            assert (
                len(self.checkpoint_shape) > 0
            ), "checkpoints shape {} should be an non empty list like: [12, 512, 1024]".format(
                self.checkpoint_shape
            )
            assert all(
                [ele > 0 for ele in self.checkpoint_shape]
            ), "all ele in checkpoints shape {} should be a determined integer larger than 0".format(
                self.checkpoint_shape
            )
J
JZ-LIANG 已提交
6970 6971 6972 6973
            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(
6974 6975
                    checkpoint_varname
                )
J
JZ-LIANG 已提交
6976
                self.checkpoint_name2pinned_name[
6977 6978
                    checkpoint_varname
                ] = pinned_var_name
J
JZ-LIANG 已提交
6979
                self.checkpoint_name2fetch_name[
6980 6981
                    checkpoint_varname
                ] = fetch_var_name
J
JZ-LIANG 已提交
6982 6983 6984 6985 6986 6987 6988 6989 6990 6991 6992 6993 6994
            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

6995 6996 6997 6998 6999 7000 7001 7002
    def backward(
        self,
        loss,
        startup_program=None,
        parameter_list=None,
        no_grad_set=None,
        callbacks=None,
    ):
M
mapingshuo 已提交
7003 7004 7005 7006 7007 7008 7009
        """
        call append_backward with checkpoints.

        Args:
            loss (Variable): loss variable to run optimizations.
            startup_program (Program): startup_program for initializing parameters
                in `parameter_list`.
7010 7011
            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 已提交
7012 7013 7014 7015 7016 7017 7018 7019
            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
7020

M
mapingshuo 已提交
7021 7022 7023 7024 7025 7026
                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
7027 7028


M
mapingshuo 已提交
7029 7030 7031 7032
                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")
7033

M
mapingshuo 已提交
7034 7035
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
7036
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
7037 7038 7039 7040
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
7041
                    no_grad_set=None)
M
mapingshuo 已提交
7042 7043
                print("Finished backward")
        """
7044 7045 7046
        assert (
            self._checkpoints is not None
        ), "You should call _set_checkpoints first"
M
mapingshuo 已提交
7047

J
Jiabin Yang 已提交
7048
        if framework._non_static_mode():
M
mapingshuo 已提交
7049
            raise NotImplementedError(
7050 7051
                "DyGraph current does not support recompute"
            )
M
mapingshuo 已提交
7052 7053 7054 7055

        self._dtype = loss.dtype
        program = loss.block.program
        with program_guard(program, startup_program):
7056 7057 7058 7059 7060 7061 7062
            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 已提交
7063 7064 7065 7066 7067 7068
            # 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,
7069 7070
                    checkpoints=checkpoint_vars,
                )
J
JZ-LIANG 已提交
7071
            else:
7072 7073 7074 7075 7076 7077
                params_grads = append_backward(
                    loss,
                    parameter_list,
                    no_grad_set,
                    checkpoints=checkpoint_vars,
                )
J
JZ-LIANG 已提交
7078 7079 7080 7081 7082

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

M
mapingshuo 已提交
7083 7084 7085 7086 7087 7088 7089 7090 7091 7092 7093 7094 7095
        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
7096

M
mapingshuo 已提交
7097 7098 7099 7100 7101
                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)
7102 7103
                    return sum_cost, fc_1, prediction

M
mapingshuo 已提交
7104 7105 7106 7107
                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")
7108

M
mapingshuo 已提交
7109 7110
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
7111
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
7112 7113 7114 7115
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
7116
                    no_grad_set=None)
7117

M
mapingshuo 已提交
7118 7119
                optimize_ops = sgd.apply_optimize(
                    cost, startup_program=None, params_grads=params_grads)
7120

M
mapingshuo 已提交
7121 7122 7123
                print("Finished apply_optimize")
        """

7124 7125 7126 7127 7128 7129 7130 7131 7132 7133 7134 7135
        func = (
            self._optimizer.apply_optimize
            if hasattr(self._optimizer, 'apply_optimize')
            else self._optimizer._apply_optimize
        )
        return func(
            loss, startup_program=startup_program, params_grads=params_grads
        )

    def minimize(
        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
    ):
7136
        assert isinstance(loss, Variable), "The loss should be an Variable."
7137 7138 7139
        assert (
            self._checkpoints is not None
        ), "You should call _set_checkpoints first"
J
Jiabin Yang 已提交
7140
        if framework._non_static_mode():
M
mapingshuo 已提交
7141
            raise NotImplementedError(
7142 7143 7144 7145 7146 7147 7148 7149
                "DyGraph current does not support recompute"
            )
        params_grads = self.backward(
            loss,
            startup_program=startup_program,
            parameter_list=parameter_list,
            no_grad_set=no_grad_set,
        )
M
mapingshuo 已提交
7150

7151 7152 7153
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
M
mapingshuo 已提交
7154 7155 7156 7157

        return optimize_ops, params_grads


7158
class LookaheadOptimizer:
7159
    r"""
7160
        :api_attr: Static Graph
S
swtkiwi 已提交
7161

M
mapingshuo 已提交
7162 7163 7164 7165
    This implements the Lookahead optimizer of the
    paper : https://arxiv.org/abs/1907.08610.

    Lookahead keeps two sets of params: the fast_params and
7166 7167
    the slow_params. inner_optimizer update fast_params every
    training step. Lookahead updates the slow_params and fast_params
M
mapingshuo 已提交
7168 7169 7170
    every k training steps as follows:

    .. math::
7171

M
mapingshuo 已提交
7172
        slow\_param_t &= slow\_param_{t-1} + \\alpha * (fast\_param_{t-1} - slow\_param_{t-1})
7173

7174
        fast\_param_t &=  slow\_param_t
M
mapingshuo 已提交
7175 7176

    Args:
7177
        inner_optimizer (Optimizer): The optimizer that update fast params step by step.
M
mapingshuo 已提交
7178 7179 7180 7181 7182 7183 7184 7185 7186
        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
7187
            import numpy.random as random
M
mapingshuo 已提交
7188

7189
            paddle.enable_static()
7190

7191 7192 7193 7194
            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)
7195
            loss = paddle.mean(x=loss)
7196 7197 7198 7199 7200 7201 7202 7203 7204
            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 已提交
7205

7206 7207 7208
            def train_reader(limit=5):
                for i in range(limit):
                    yield random.random([2]).astype('float32'), random.random([1]).astype('int64')
7209

7210 7211
            feeder = fluid.DataFeeder(feed_list=[x, label], place=place)
            reader = paddle.batch(paddle.reader.shuffle(train_reader, buf_size=50000),batch_size=1)
7212

7213 7214 7215
            for batch_data in reader():
                exe.run(fluid.default_main_program(),
                feed=feeder.feed(batch_data))
M
mapingshuo 已提交
7216 7217 7218 7219 7220

    """

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

J
Jiabin Yang 已提交
7221
        if framework._non_static_mode():
Z
zhongpu 已提交
7222
            raise Exception("In dygraph, don't support LookaheadOptimizer.")
7223
        assert inner_optimizer is not None, "inner optimizer can not be None"
M
mapingshuo 已提交
7224 7225 7226
        assert (
            0.0 <= alpha <= 1.0
        ), "alpha should be larger or equal to 0.0, and less or equal than 1.0"
7227
        assert isinstance(k, int) and k > 0, "k should be a positive integer"
M
mapingshuo 已提交
7228 7229 7230 7231 7232 7233 7234 7235 7236 7237

        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(
7238 7239
            loss, startup_program=startup_program
        )
M
mapingshuo 已提交
7240 7241 7242 7243 7244 7245 7246 7247 7248 7249 7250

        # 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)
7251 7252 7253 7254 7255 7256 7257
            assert fast_var is not None
            slow_var = main_block.create_var(
                name=param + "@SLOW",
                shape=fast_var.shape,
                dtype=fast_var.dtype,
                persistable=True,
            )
M
mapingshuo 已提交
7258 7259 7260 7261 7262 7263
            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)
7264 7265 7266 7267 7268 7269 7270
            assert fast_var is not None
            slow_var = startup_block.create_var(
                name=param + "@SLOW",
                shape=fast_var.shape,
                dtype=fast_var.dtype,
                persistable=True,
            )
M
mapingshuo 已提交
7271

7272 7273 7274
            startup_block.append_op(
                type="assign", inputs={"X": fast_var}, outputs={"Out": slow_var}
            )
M
mapingshuo 已提交
7275

7276 7277
        with framework.program_guard(main_block.program, startup_program):
            # Add Var k to main prog and startup prog
7278
            k = paddle.static.create_global_var(
7279 7280 7281 7282 7283 7284
                name="lookahead_k",
                shape=[1],
                value=int(self.k),
                dtype='int32',
                persistable=True,
            )
M
mapingshuo 已提交
7285

7286
            # Add Var alpha to main prog and startup prog
7287
            alpha = paddle.static.create_global_var(
7288 7289 7290 7291 7292 7293
                name="lookahead_alpha",
                shape=[1],
                value=float(self.alpha),
                dtype='float32',
                persistable=True,
            )
M
mapingshuo 已提交
7294

7295
            # Add Var step
7296
            step = paddle.static.create_global_var(
7297 7298 7299 7300 7301 7302
                name="lookahead_step",
                shape=[1],
                value=int(0),
                dtype='int32',
                persistable=True,
            )
7303
            paddle.increment(x=step, value=1.0)
7304 7305

            # lookahead
7306 7307 7308
            zero_var = layers.fill_constant(
                shape=[1], dtype='float32', value=0.0
            )
7309

7310 7311 7312
            one_var = layers.fill_constant(
                shape=[1], dtype='float32', value=1.0
            )
7313

7314
            mod = paddle.remainder(step, k)
7315
            with layers.control_flow.Switch() as switch:
7316 7317 7318 7319 7320
                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)
7321 7322 7323 7324
                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]
7325 7326 7327 7328
                        tmp_var = paddle.add(
                            paddle.multiply(fast_var, alpha),
                            paddle.multiply(
                                slow_var, paddle.subtract(one_var, alpha)
7329 7330
                            ),
                        )
7331 7332 7333 7334
                        layers.assign(input=tmp_var, output=slow_var)
                        layers.assign(input=tmp_var, output=fast_var)
                with switch.default():
                    pass
M
mapingshuo 已提交
7335
        return mini_out
7336 7337


7338
class GradientMergeOptimizer:
7339 7340 7341 7342 7343 7344 7345 7346 7347 7348 7349 7350 7351 7352 7353 7354 7355 7356 7357 7358 7359 7360 7361 7362 7363 7364 7365 7366 7367 7368 7369 7370 7371 7372 7373 7374 7375 7376 7377 7378 7379 7380 7381 7382 7383 7384 7385 7386 7387 7388 7389 7390 7391 7392
    """
    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]))
    """

7393 7394
    GRAD_MERGE_COND_NAME = "grad_merge_cond_name"

7395
    def __init__(self, inner_optimizer, k_steps=1, avg=True):
J
Jiabin Yang 已提交
7396
        if framework._non_static_mode():
7397 7398 7399
            raise Exception(
                "In dygraph, we don't support GradientMergeOptimizer."
                "You can do Gradient merge by yourself with k-times forward + backward, "
7400 7401
                "and one-time optimizer.minimize()"
            )
7402

7403 7404 7405 7406
        assert inner_optimizer is not None, "inner optimizer can not be None"
        assert (
            isinstance(k_steps, int) and k_steps > 0
        ), "k_steps should be a positive integer"
7407 7408 7409 7410 7411

        self.inner_optimizer = inner_optimizer
        self.k_steps = k_steps
        self.type = "gradient_merge"
        self.avg = avg
7412
        self._optimize_ops = None
7413

7414 7415 7416 7417 7418 7419
    def _set_k_steps(self, k_steps):
        self.k_steps = k_steps

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

7420 7421 7422 7423 7424 7425 7426 7427
    def backward(
        self,
        loss,
        startup_program=None,
        parameter_list=None,
        no_grad_set=None,
        callbacks=None,
    ):
7428 7429 7430 7431 7432 7433 7434 7435 7436
        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(
7437 7438
            loss, startup_program=startup_program
        )
7439 7440 7441 7442 7443 7444 7445 7446 7447 7448 7449
        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
7450 7451 7452
        if op_maker.kOpRoleVarAttrName() in op.attr_names and int(
            op.all_attrs()[op_maker.kOpRoleAttrName()]
        ) == int(backward):
7453 7454 7455 7456 7457 7458
            return True
        return False

    def _remove_op_role_var(self, param, grad):
        op_maker = core.op_proto_and_checker_maker
        op = grad.op
7459 7460 7461 7462 7463
        assert self._is_the_backward_op(
            op
        ), 'grad.op={} is not the backward op which produces the grad={}'.format(
            op, grad.name
        )
7464 7465 7466

        block = grad.block
        var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
7467 7468 7469 7470 7471 7472 7473 7474 7475 7476
        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
        )
7477 7478 7479 7480 7481 7482 7483 7484 7485 7486 7487 7488 7489 7490 7491 7492 7493 7494 7495 7496 7497 7498 7499 7500 7501 7502

        # 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
7503
        k_step_var = paddle.static.create_global_var(
7504 7505 7506 7507 7508 7509 7510 7511
            name="gradient_merge_k",
            shape=[1],
            value=int(self.k_steps),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )

7512
        zero_var = paddle.static.create_global_var(
7513 7514 7515 7516 7517 7518 7519
            name="gradient_merge_zero",
            shape=[1],
            value=int(0),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )
7520 7521

        # Add step var & cond var
7522
        step_var = paddle.static.create_global_var(
7523 7524 7525 7526 7527 7528 7529
            name="gradient_merge_step",
            shape=[1],
            value=int(0),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )
7530

7531 7532 7533
        cond_var = main_block.create_var(
            name="gradient_merge_cond", shape=[1], dtype='bool'
        )
7534 7535 7536

        with device_guard("cpu"):
            # step_var = (step_var + 1) % k_step
7537
            paddle.increment(x=step_var, value=1.0)
7538 7539 7540 7541 7542 7543
            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},
            )
7544 7545

            # cond_var = (step_var == 0)
7546 7547 7548 7549 7550
            main_block.append_op(
                type='equal',
                inputs={'X': step_var, 'Y': zero_var},
                outputs={'Out': cond_var},
            )
7551 7552 7553 7554 7555 7556 7557 7558 7559 7560

        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)
7561

7562
        # TODO(mapingshuo) support sparse embedding
7563 7564
        # step1: remove grad.op's op_role_var
        for param, grad in params_grads:
7565
            assert (
7566
                param.type != core.VarDesc.VarType.SELECTED_ROWS
7567 7568
            ), "SELECTED_ROWS is not supported in GradientMergeOptimizer for now"

7569
            self._remove_op_role_var(param, grad)
7570

7571
        param_to_grad = {k.name: v for (k, v) in params_grads}
7572 7573 7574
        param_names = param_to_grad.keys()
        param_to_gradient_merge = {}

7575 7576 7577 7578 7579
        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
7580
            param_var = main_block.var(param_name)
7581 7582 7583 7584 7585 7586 7587
            assert param_var is not None
            gradient_merge_var = main_block.create_var(
                name=param_name + "@GRAD@GradientMerge",
                shape=param_var.shape,
                dtype=param_var.dtype,
                persistable=True,
            )
7588
            param_to_gradient_merge[param_name] = gradient_merge_var
7589

7590 7591 7592 7593
            startup_gradient_merge_var = startup_block.create_var(
                name=param_name + "@GRAD@GradientMerge",
                shape=param_var.shape,
                dtype=param_var.dtype,
7594 7595 7596 7597 7598 7599 7600 7601 7602 7603 7604
                persistable=True,
            )
            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),
                },
            )
7605

7606 7607 7608
            # grad_merge += grad
            new_grad_op = main_block.append_op(
                type="elementwise_add",
7609
                inputs={'X': grad, 'Y': gradient_merge_var},
7610
                outputs={'Out': gradient_merge_var},
7611 7612 7613 7614 7615
                attrs={'axis': -1, 'use_mkldnn': False},
            )
            self._add_gm_op_role_var(
                new_grad_op, param, gradient_merge_var, cond
            )
7616 7617 7618 7619 7620 7621 7622 7623
            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)
7624
            op_maker = core.op_proto_and_checker_maker
7625 7626 7627 7628

            if self.avg:
                for param, new_grad in new_params_grads:
                    # grad /= k_steps
7629 7630 7631 7632 7633 7634 7635 7636 7637 7638 7639 7640 7641
                    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,
                        },
                    )
                    new_grad.op._set_attr(
                        op_maker.kOpRoleAttrName(), op_maker.OpRole.Backward
                    )
7642

7643 7644 7645 7646 7647 7648
            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
7649

7650
            self._optimize_ops = self.inner_optimizer.apply_gradients(
7651 7652
                new_params_grads
            )
7653

7654 7655
            # clear gradient_merge_vars
            for param, new_grad in new_params_grads:
7656 7657 7658 7659 7660 7661 7662 7663 7664
                layers.fill_constant(
                    shape=new_grad.shape,
                    dtype=new_grad.dtype,
                    value=0.0,
                    out=new_grad,
                )
                new_grad.op._set_attr(
                    op_maker.kOpRoleAttrName(), op_maker.OpRole.Optimize
                )
7665 7666

        # step3. apply gradient
7667
        paddle.static.nn.cond(cond, true_fn=true_apply_gradient, false_fn=None)
7668 7669 7670

        return self._optimize_ops

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

7676 7677 7678 7679 7680 7681
        params_grads = self.backward(
            loss,
            startup_program=startup_program,
            parameter_list=parameter_list,
            no_grad_set=no_grad_set,
        )
7682

7683 7684 7685
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
7686 7687

        return optimize_ops, params_grads