optimizer.py 318.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
S
sneaxiy 已提交
51
from .layers import ops
52
from .dygraph import base as imperative_base
53
from .dygraph import no_grad
54 55 56 57
from .dygraph.learning_rate_scheduler import (
    LearningRateDecay,
    _LearningRateEpochDecay,
)
58 59 60
from paddle.fluid import core
from paddle.fluid.layers import tensor
from functools import reduce
61
from functools import cmp_to_key
62
from .wrapped_decorator import signature_safe_contextmanager
63
import warnings
64
from paddle import _C_ops, _legacy_C_ops
65 66 67 68 69
from ..fluid.framework import (
    _in_legacy_dygraph,
    in_dygraph_mode,
    _current_expected_place,
)
70

71
__all__ = [
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 98
    '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',
99
]
Q
Qiao Longfei 已提交
100 101 102 103 104 105


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

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

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

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

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

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

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

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

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

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

                import paddle.fluid as fluid
218 219 220 221 222 223

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

                    adam = fluid.optimizer.Adam(0.001, parameter_list=emb.parameters())
                    state_dict = adam.state_dict()
H
hong 已提交
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 422
                    self._learning_rate_map[
                        framework.default_main_program()
                    ] = layers.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 453 454 455 456 457 458 459
                framework.default_main_program()
            ] = layers.create_global_var(
                name=unique_name.generate("learning_rate"),
                shape=[1],
                value=float(self._learning_rate),
                dtype='float32' if self._dtype is None else self._dtype,
                persistable=True,
            )
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 478
        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
479

480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513
                with fluid.dygraph.guard():
                    linear = fluid.dygraph.nn.Linear(10, 10)

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

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


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



        """
        if not isinstance(value, (framework.Variable, float)):
            raise TypeError(
                "The type of 'value' in optimizer.set_lr must be (float, Variable), but received %s."
514 515
                % (type(value))
            )
516 517 518 519 520 521 522 523
        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:
524 525
                if in_dygraph_mode():
                    place = _current_expected_place()
526 527 528 529 530 531 532
                    _C_ops.full_(
                        current_lr,
                        list(current_lr.shape),
                        float(value),
                        current_lr.dtype,
                        place,
                    )
533 534

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

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

569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594
        Get current step learning rate. The return value is all the same When LearningRateDecay is not used,
        otherwise return the step learning rate.

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

        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
                import numpy as np

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

                # example2: PiecewiseDecay is used, return the step learning rate
                with fluid.dygraph.guard():
                    inp = np.random.uniform(-0.1, 0.1, [10, 10]).astype("float32")
                    linear = fluid.dygraph.nn.Linear(10, 10)
                    inp = fluid.dygraph.to_variable(inp)
                    out = linear(inp)
                    loss = fluid.layers.reduce_mean(out)
595

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

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

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

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

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

    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 已提交
662
        """
663 664
        pass

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

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

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

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

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

Q
Qiao Longfei 已提交
744
        self._accumulators[name][param.name] = var
745
        return var
746

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

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

        self._global_accumulators[name] = var
        return var

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

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

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

    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

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

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

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

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

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

        start = len(target_block.ops)
904

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

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

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

937 938
        end = len(target_block.ops)
        return target_block._slice_ops(start, end)
939 940

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

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

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

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

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

L
Leo Chen 已提交
1023 1024 1025 1026
        # Infer dtype by loss if None
        if self._dtype is None:
            self._dtype = loss.dtype

J
Jiabin Yang 已提交
1027
        if framework._non_static_mode():
1028 1029 1030
            parameter_list = (
                parameter_list if parameter_list else self._parameter_list
            )
1031

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

1061
    def _create_regularization_of_grad(self, param, grad, regularization=None):
1062
        """Create and add backward regularization Operators
1063

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

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

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

        return new_grad

1107 1108 1109
    def append_regularization_ops(
        self, parameters_and_grads, regularization=None
    ):
1110
        r"""Create and add backward regularization Operators
1111

1112 1113 1114 1115
        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.
1116

1117 1118 1119 1120 1121
        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.
1122

1123 1124 1125
        Returns:
            list[(Variable, Variable)]: list of (parameters, gradients) \
            pair with the regularized gradient
1126

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

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

        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)],
1200 1201
            belong_to_optimizer=True,
        )
1202 1203

        with program_guard(default_main_program()):
1204 1205 1206 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
            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,
                },
            )
1233

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

        return [(flatten_param, flatten_grad)]

1245 1246 1247 1248 1249 1250 1251
    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 已提交
1252

1253 1254
        Returns:
            list: A list of operators appended to the current program.
M
minqiyang 已提交
1255

1256 1257 1258
        Examples:
            .. code-block:: python

1259
                import paddle.fluid as fluid
1260 1261 1262 1263 1264 1265 1266 1267 1268
                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)

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

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

        # Add regularization if any
1283 1284 1285
        params_grads = self.append_regularization_ops(
            params_grads, self.regularization
        )
1286 1287 1288 1289

        optimize_ops = self._create_optimization_pass(params_grads)
        return optimize_ops

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

        return no_grad_set

1330 1331 1332 1333
    @framework.dygraph_only
    def clear_gradients(self):
        """
        Clear the gradients of all optimized parameters for model.
1334 1335

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

1337 1338
        Returns:
            None
1339

1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350
        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
                import numpy as np

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

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

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

1381
        Returns:
1382 1383 1384
            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.
1385 1386
            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
1387
            ``fetch_list`` before run, see details in ``Executor``.
1388 1389 1390

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

1394 1395 1396
        parameter_list = (
            parameter_list if parameter_list else self._parameter_list
        )
1397

1398 1399 1400 1401 1402 1403
        params_grads = self.backward(
            loss,
            startup_program=startup_program,
            parameter_list=parameter_list,
            no_grad_set=no_grad_set,
        )
1404

1405 1406 1407
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
M
minqiyang 已提交
1408

Q
Qiao Longfei 已提交
1409
        return optimize_ops, params_grads
Q
Qiao Longfei 已提交
1410 1411 1412


class SGDOptimizer(Optimizer):
1413
    r"""
Q
qiaolongfei 已提交
1414 1415 1416 1417 1418 1419
    Optimizer of the stochastic gradient descent algorithm.

    .. math::

        param\_out = param - learning\_rate * grad

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

    Examples:
        .. code-block:: python

1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465
            import paddle
            import paddle.fluid as fluid
            import numpy as np

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

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

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

Q
Qiao Longfei 已提交
1466 1467
    """

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

1537
    @no_grad
1538
    def _append_optimize_op(self, block, param_and_grad):
1539

1540 1541 1542 1543 1544 1545 1546 1547 1548
        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
        )
1549

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

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

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

        attrs = {"multi_precision": find_master}

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

1587 1588 1589 1590 1591 1592 1593
        sgd_op = block.append_op(
            type=self.type,
            inputs=inputs,
            outputs=outputs,
            attrs=attrs,
            stop_gradient=True,
        )
Q
Qiao Longfei 已提交
1594 1595

        return sgd_op
1596 1597 1598


class MomentumOptimizer(Optimizer):
1599
    r"""
Q
qiaolongfei 已提交
1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612

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

1613
        &\quad   param = param - (gradient + mu * velocity) * learning\_rate
Q
qiaolongfei 已提交
1614 1615 1616

        & else:

Q
qiaolongfei 已提交
1617
        &\quad   param = param - learning\_rate * velocity
Q
qiaolongfei 已提交
1618

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

    Examples:
        .. code-block:: python

1642 1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666
            import paddle
            import paddle.fluid as fluid
            import numpy as np

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

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

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

1667 1668 1669
    """
    _velocity_acc_str = "velocity"

1670 1671 1672 1673 1674 1675 1676 1677 1678 1679
    def __init__(
        self,
        learning_rate,
        momentum,
        parameter_list=None,
        use_nesterov=False,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
1680 1681
        assert learning_rate is not None
        assert momentum is not None
1682 1683 1684 1685 1686 1687 1688
        super(MomentumOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
1689 1690
        self.type = "momentum"
        self._momentum = momentum
1691
        self._use_nesterov = bool(use_nesterov)
1692 1693 1694 1695 1696

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

        for p in parameters:
Q
Qiao Longfei 已提交
1697
            self._add_accumulator(self._velocity_acc_str, p)
1698 1699 1700 1701

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

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

1724
        attrs = {"mu": self._momentum, "use_nesterov": self._use_nesterov}
1725 1726 1727 1728
        inputs = {
            "Param": [param_and_grad[0]],
            "Grad": [param_and_grad[1]],
            "Velocity": [velocity_acc],
1729
            "LearningRate": [lr],
1730 1731 1732 1733
        }

        outputs = {
            "ParamOut": [param_and_grad[0]],
1734
            "VelocityOut": [velocity_acc],
1735
        }
1736
        # create the momentum optimize op
1737 1738 1739 1740 1741 1742 1743
        momentum_op = block.append_op(
            type=self.type,
            inputs=inputs,
            outputs=outputs,
            attrs=attrs,
            stop_gradient=True,
        )
1744 1745

        return momentum_op
1746 1747


1748
class DGCMomentumOptimizer(Optimizer):
1749
    r"""
1750
	:api_attr: Static Graph
S
swtkiwi 已提交
1751

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

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

G
gongweibao 已提交
1757
    To avoid losing information, DGC accumulates the rest of the gradients locally.
1758 1759 1760

    Eventually, these gradients become large enough to be transmitted.

1761
    Thus, DGC sends the large gradients immediately but eventually sends all of the gradients over time.
1762

G
gongweibao 已提交
1763
    To ensure no loss of accuracy, DGC employs momentum correction and local gradient clipping on top of the gradient sparsification to maintain model performance.
1764 1765 1766 1767

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

    This optimizer will do two things:
1768

1769 1770
        1. Compress the gradient by get TopK import value from tensor \
            and use it for allreduce to reduce network bandwidth.
1771

1772
        2. Call momentum to optimize the cost.
1773 1774

    Args:
1775 1776
        learning_rate (float|Variable): The learning rate used to update parameters. \
            It can be a float value or a Variable with one float value as a data element.
1777
        momentum (float): Momentum factor.
G
gongweibao 已提交
1778
        rampup_begin_step (int): The beginning step from which gradient compression is implemented.
1779 1780 1781 1782 1783 1784 1785
        rampup_step (int): Time steps used in sparsity warm-up periods. Default is 1.
            For example, if the sparsity is [0.75, 0.9375, 0.984375, 0.996, 0.999], and the rampup_step is 100, \
                it will use 0.75 at 0~19 steps, and 0.9375 at 20~39 steps, and so on. \
                And when reach sparsity array ends, it will use 0.999 then and after.
        sparsity (list[float]): Get top important element from gradient tensor, the ratio is (1 - current sparsity). \
            Default is [0.999]. For example, if the sparsity is [0.99, 0.999], \
                the top [1%, 0.1%] important element will be transmitted.
H
hong 已提交
1786
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1787 1788
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
1789
        use_nesterov (bool): Enables Nesterov momentum. True means use Nesterov. Default is False.
1790 1791 1792 1793 1794
        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.
1795 1796
        grad_clip (GradientClipByNorm, optional): Gradient cliping strategy. ``DGCMomentumOptimizer`` only support
            :ref:`api_fluid_clip_GradientClipByNorm` , and if not, it will raise TypeError. Default None,
1797
            meaning there is no gradient clipping.
1798 1799
        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.
1800 1801 1802 1803

    Examples:
        .. code-block:: python

1804
            import paddle.fluid as fluid
1805
            optimizer = fluid.optimizer.DGCMomentumOptimizer(
G
gongweibao 已提交
1806 1807 1808 1809 1810
                        learning_rate=0.0001,
                        momentum=0.9,
                        rampup_step=1000,
                        rampup_begin_step=1252,
                        sparsity=[0.999, 0.999])
1811 1812

    """
1813 1814
    _u_velocity_acc_str = "_dgc_u_"
    _v_velocity_acc_str = "_dgc_v_"
1815

1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829
    def __init__(
        self,
        learning_rate,
        momentum,
        rampup_begin_step,
        rampup_step=1,
        sparsity=[0.999],
        parameter_list=None,
        use_nesterov=False,
        num_trainers=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
J
Jiabin Yang 已提交
1830
        if framework._non_static_mode():
Z
zhongpu 已提交
1831
            raise Exception("In dygraph, don't support DGCMomentumOptimizer.")
1832

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

1837 1838
        assert learning_rate is not None
        assert momentum is not None
1839 1840 1841 1842 1843 1844 1845
        super(DGCMomentumOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
1846 1847 1848
        self.type = "dgc_momentum"
        self._momentum = momentum
        self._use_nesterov = bool(use_nesterov)
1849

1850
        assert rampup_begin_step >= 0, "rampup_begin_step must >= 0"
1851
        self._rampup_begin_step = rampup_begin_step
1852 1853
        self._rampup_step = rampup_step
        self._sparsity = sparsity
1854

1855
        self._rampup_begin_step_var = None
1856
        self._global_step_var = None
1857

1858 1859 1860 1861 1862 1863
        self._dgc_clip_norm = None
        if grad_clip is not None:
            if not isinstance(grad_clip, GradientClipByNorm):
                raise TypeError(
                    "The type of grad_clip should be 'GradientClipByNorm', because DGCMomentumOptimizer only support GradientClipByNorm"
                )
1864 1865 1866 1867 1868 1869 1870
            assert isinstance(num_trainers, int), (
                "The type of num_trainers should be 'int', but received %s"
                % type(num_trainers)
            )
            assert (
                num_trainers > 0
            ), "The value of num_trainers should be greater than 0!"
1871 1872

            self._num_trainers = num_trainers
1873
            self._dgc_clip_norm = grad_clip.clip_norm * (num_trainers**-0.5)
1874

1875
        self.regular_type, self.regular_coeff = self._get_regularization_param(
1876 1877
            self.regularization
        )
1878

1879 1880 1881
    def _get_regularization_param(self, regularization):
        regular_type = 0
        regular_coeff = 0.0
1882

1883 1884
        if regularization is not None:
            regular_coeff = regularization._regularization_coeff
1885
            from .regularizer import L1Decay, L2Decay
1886

1887 1888 1889 1890
            if isinstance(regularization, L1Decay):
                regular_type = 1
            elif isinstance(regularization, L2Decay):
                regular_type = 2
1891 1892
            else:
                assert False, 'regularization must be None|L1Decay|L2Deacy'
1893
        return regular_type, regular_coeff
1894

1895 1896
    def _is_use_dgc(self, param_var, grad_var):
        var_numel = abs(reduce(lambda x, y: x * y, param_var.shape))
1897 1898 1899 1900 1901 1902
        if (
            var_numel < 16384
            or param_var.type == core.VarDesc.VarType.SELECTED_ROWS
            or grad_var.type == core.VarDesc.VarType.SELECTED_ROWS
            or param_var.dtype != core.VarDesc.VarType.FP32
        ):
1903 1904 1905 1906 1907
            return False
        return True

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)
1908 1909 1910
        velocity_acc = self._get_accumulator(
            self._u_velocity_acc_str, param_and_grad[0]
        )
1911 1912 1913 1914 1915 1916 1917 1918 1919 1920 1921 1922 1923
        assert velocity_acc is not None

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

        if not self._is_use_dgc(param_and_grad[0], param_and_grad[1]):
1926 1927 1928
            type = "momentum"
        else:
            type = "dgc_momentum"
1929 1930 1931 1932 1933 1934
            inputs.update(
                {
                    "current_step": self._global_step_var,
                    "nranks": self._nranks_var,
                }
            )
1935
            outputs.update({'Grad_out': param_and_grad[1]})
1936
            attrs.update({"rampup_begin_step": float(self._rampup_begin_step)})
1937 1938

        # create the dgc momentum optimize op
1939 1940 1941 1942 1943 1944 1945
        dgc_momentum_op = block.append_op(
            type=type,
            inputs=inputs,
            outputs=outputs,
            attrs=attrs,
            stop_gradient=True,
        )
1946 1947
        return dgc_momentum_op

1948 1949 1950
    def _add_auto_increment_var(self, counter_name, begin, step=1):
        helper = LayerHelper('global_step_counter')
        counter, is_new_var = helper.create_or_get_global_variable(
1951 1952
            name=counter_name, dtype='float32', shape=[1], persistable=True
        )
1953
        if is_new_var:
1954 1955 1956 1957
            helper.set_variable_initializer(
                counter,
                initializer=Constant(value=float(begin - 1), force_cpu=True),
            )
1958 1959 1960 1961 1962
            helper.main_program.global_block()._prepend_op(
                type='increment',
                inputs={'X': [counter]},
                outputs={'Out': [counter]},
                attrs={'step': float(step)},
1963 1964
                stop_gradient=True,
            )
1965 1966 1967 1968
            counter.stop_gradient = True

        return counter

1969 1970 1971
    def _add_nranks_var(self, name, value=-1):
        helper = LayerHelper('global_step_counter')
        counter, is_new_var = helper.create_or_get_global_variable(
1972 1973
            name=name, dtype='float32', shape=[1], persistable=True
        )
1974
        if is_new_var:
1975 1976 1977 1978
            helper.set_variable_initializer(
                counter,
                initializer=Constant(value=float(value), force_cpu=True),
            )
1979 1980 1981 1982
            counter.stop_gradient = True

        return counter

1983 1984 1985 1986 1987 1988
    def _append_dgc_ops(self, param_and_grads):
        main_program = default_main_program()
        main_program._enable_dgc = True

        # step counter
        self._global_step_var = self._add_auto_increment_var(
1989 1990
            counter_name=core.dgc.kDGCCounterName(), begin=0
        )
1991

1992 1993 1994
        self._nranks_var = self._add_nranks_var(
            name=core.dgc.kDGCNRanksName(), value=-1
        )
1995

1996 1997 1998 1999 2000
        # rampup begin step var for all_reduce_op_handle
        self._rampup_begin_step_var = tensor.create_global_var(
            shape=[1],
            dtype=core.VarDesc.VarType.FP32,
            persistable=True,
G
gongweibao 已提交
2001
            name=core.dgc.kDGCRampUpBeginStepName(),
2002
            value=self._rampup_begin_step * 1.0,
2003 2004
            force_cpu=True,
        )
2005

2006 2007
        self.helper = LayerHelper(self.__class__.__name__)

2008
        for param_var, grad_var in param_and_grads:
2009 2010 2011
            # reuse velocity in dgc_op and dgc_momentum_op
            u_var = self._add_accumulator(self._u_velocity_acc_str, param_var)

2012
            if not self._is_use_dgc(param_var, grad_var):
2013 2014
                continue

2015
            v_var = self._add_accumulator(self._v_velocity_acc_str, param_var)
2016

2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027 2028 2029 2030 2031 2032 2033 2034 2035 2036 2037 2038 2039 2040 2041 2042
            k_var = tensor.create_global_var(
                shape=[1],
                dtype=param_var.dtype,
                persistable=True,
                name=param_var.name + core.dgc.kDGCKName(),
                value=0.0,
                force_cpu=True,
            )

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

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

2044 2045 2046 2047 2048 2049 2050 2051 2052 2053 2054 2055 2056 2057 2058 2059 2060 2061 2062
            # del back oprolevarname
            op_maker = core.op_proto_and_checker_maker
            backward = core.op_proto_and_checker_maker.OpRole.Backward
            for op in main_program.global_block().ops:
                if not self._is_the_backward_op(op):
                    continue

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

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

            clip_var = grad_var
2063 2064
            if self._dgc_clip_norm is not None:
                clip_var = self._append_clip_norm(grad_var, self._dgc_clip_norm)
2065 2066 2067 2068 2069 2070 2071 2072 2073 2074
            self._dgc_op(
                param_var,
                clip_var,
                grad_var,
                u_var,
                v_var,
                k_var,
                encoded_var,
                gather_var,
            )
2075 2076 2077 2078

    def _is_the_backward_op(self, op):
        op_maker = core.op_proto_and_checker_maker
        backward = core.op_proto_and_checker_maker.OpRole.Backward
2079 2080 2081
        if op_maker.kOpRoleVarAttrName() in op.attr_names and int(
            op.all_attrs()[op_maker.kOpRoleAttrName()]
        ) == int(backward):
2082 2083 2084 2085 2086 2087 2088 2089 2090
            return True
        return False

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

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

        if name is None:
2091 2092 2093 2094 2095 2096 2097 2098 2099 2100 2101 2102 2103 2104 2105 2106 2107
            name = unique_name.generate_with_ignorable_key(
                ".".join([helper.name, 'tmp'])
            )

        out = helper.create_variable(
            type=x.type, name=name, dtype=x.dtype, persistable=False
        )

        helper.append_op(
            type="dgc_clip_by_norm",
            inputs={"X": x, "current_step": self._global_step_var},
            attrs={
                "max_norm": max_norm,
                "rampup_begin_step": float(self._rampup_begin_step),
            },
            outputs={"Out": out},
        )
2108 2109 2110 2111
        return out

    def _append_clip_norm(self, grad_var, clip_norm):
        with grad_var.block.program._backward_role_guard():
2112 2113 2114
            return self._clip_by_norm(
                x=grad_var, max_norm=clip_norm, name=grad_var.name
            )
2115

2116 2117 2118 2119 2120 2121 2122 2123 2124 2125 2126
    def _dgc_op(
        self,
        param_var,
        clip_var,
        grad_var,
        u_var,
        v_var,
        k_var,
        encoded_var,
        gather_var,
    ):
2127 2128
        block = framework.default_main_program().global_block()
        op_maker = core.op_proto_and_checker_maker
2129

2130 2131 2132 2133 2134
        regular_type = self.regular_type
        regular_coeff = self.regular_coeff
        # The regularizer of the Parameters have higher priority
        if param_var.regularizer is not None:
            regular_type, regular_coeff = self._get_regularization_param(
2135 2136 2137 2138 2139 2140 2141 2142 2143 2144 2145 2146 2147 2148 2149 2150 2151 2152 2153 2154 2155 2156 2157 2158 2159 2160 2161 2162 2163 2164 2165 2166
                param_var.regularizer
            )

        dgc_op = block.append_op(
            type="dgc",
            inputs={
                "U": u_var,
                "V": v_var,
                "Grad": clip_var,
                "Param": param_var,
                "current_step": self._global_step_var,
                "nranks": self._nranks_var,
            },
            outputs={
                "U_out": u_var,
                "V_out": v_var,
                "EncodeGrad": encoded_var,
                "k": k_var,
                "Grad_out": grad_var,
                "GatherBuff": gather_var,
            },
            attrs={
                "m": self._momentum,
                "sparsity": self._sparsity,
                "use_nesterov": self._use_nesterov,
                "rampup_begin_step": float(self._rampup_begin_step),
                "rampup_step": float(self._rampup_step),
                "regular_coeff": float(regular_coeff),
                "regular_type": int(regular_type),
            },
            stop_gradient=True,
        )
2167 2168 2169

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

2174
    @imperative_base.no_grad
2175
    def apply_gradients(self, params_grads):
2176 2177 2178 2179 2180
        # Note: since we can't use all_reduce_op now,
        # dgc_op should be the last op of one grad.
        # Maybe need a grad allreduce pass.
        self._append_dgc_ops(params_grads)

2181
        params_grads = sorted(params_grads, key=lambda x: x[0].name)
2182 2183 2184 2185 2186
        (
            params_grads,
            table_param_and_grad,
            table_optimize_op,
        ) = self._process_distribute_lookuptable(params_grads)
2187 2188 2189

        not_dgc_params_grads = []
        dgc_params_grads = []
2190
        # DGC clip and regularization in optimizer.backward
2191 2192 2193 2194 2195 2196
        for param, grad in params_grads:
            if not self._is_use_dgc(param, grad):
                not_dgc_params_grads.append((param, grad))
            else:
                dgc_params_grads.append((param, grad))

2197
        # 'optimizer(grad_clip)' or 'set_gradient_clip'
2198 2199 2200 2201
        if self._grad_clip is not None:
            not_dgc_params_grads = self._grad_clip(not_dgc_params_grads)
        else:
            not_dgc_params_grads = append_gradient_clip_ops(
2202 2203
                not_dgc_params_grads
            )
2204

2205
        not_dgc_params_grads = self.append_regularization_ops(
2206 2207
            not_dgc_params_grads, self.regularization
        )
2208 2209 2210 2211 2212 2213 2214 2215 2216 2217 2218

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

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

        return optimize_ops

2219

2220
class LarsMomentumOptimizer(Optimizer):
2221
    r"""
2222 2223 2224 2225 2226 2227 2228 2229 2230
    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||}

2231
        & velocity = mu * velocity + local\_learning\_rate * (gradient + lars\_weight\_decay * param + epsilon)
2232 2233 2234

        & param = param - velocity

2235 2236 2237 2238 2239 2240
    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 已提交
2241
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2242 2243
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2244 2245 2246 2247 2248
        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.
2249 2250 2251
        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` ,
2252
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2253 2254
        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.
2255 2256
        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.
2257 2258 2259
        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`.
2260

2261 2262 2263
    Examples:
        .. code-block:: python

2264 2265 2266 2267 2268 2269 2270 2271 2272 2273 2274 2275 2276 2277 2278 2279
            import paddle.fluid as fluid
            import numpy as np

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

            exe = fluid.Executor(fluid.CPUPlace())
            exe.run(fluid.default_startup_program())
            exe.run(
                feed={"inp": np_inp},
                fetch_list=[out.name])
2280 2281 2282
    """
    _velocity_acc_str = "velocity"

2283 2284 2285 2286 2287 2288 2289 2290 2291 2292 2293 2294 2295 2296 2297
    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,
    ):
2298 2299
        assert learning_rate is not None
        assert momentum is not None
2300 2301 2302 2303 2304 2305 2306
        super(LarsMomentumOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
2307 2308 2309 2310
        self.type = "lars_momentum"
        self._momentum = momentum
        self._lars_coeff = float(lars_coeff)
        self._lars_weight_decay = float(lars_weight_decay)
2311 2312 2313 2314 2315
        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
2316 2317 2318 2319 2320
        self._multi_precision = multi_precision
        self._rescale_grad = float(rescale_grad)
        self._master_weights = {}

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

2326 2327
            var_name = param.name + '_fp32_master'
            var_name = unique_name.generate(var_name)
2328 2329 2330 2331 2332 2333 2334
            var = layers.create_global_var(
                name=var_name,
                shape=param.shape,
                value=0,
                dtype='float32',
                persistable=True,
            )
2335
            block = self.helper.startup_program.global_block()
2336 2337 2338 2339 2340 2341 2342 2343 2344
            block.append_op(
                type="cast",
                inputs={"X": [param]},
                outputs={"Out": [var]},
                attrs={
                    "in_dtype": param.dtype,
                    "out_dtype": core.VarDesc.VarType.FP32,
                },
            )
2345
            self._master_weights[param.name] = var
2346 2347 2348 2349 2350 2351 2352 2353 2354 2355 2356 2357
        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
2358 2359 2360 2361 2362 2363
        find_master = (
            self._multi_precision and param.dtype == core.VarDesc.VarType.FP16
        )
        target_param = (
            self._master_weights[param.name] if find_master else param
        )
2364
        target_name = target_param.name
2365 2366 2367 2368
        if (
            name not in self._accumulators
            or target_name not in self._accumulators[name]
        ):
2369 2370
            raise Exception(
                "Accumulator {} does not exist for parameter {}".format(
2371 2372 2373
                    name, target_name
                )
            )
2374
        return self._accumulators[name][target_name]
2375 2376 2377 2378 2379

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

        for p in parameters:
2380 2381 2382 2383
            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
2384 2385 2386 2387
            if (
                p.dtype == core.VarDesc.VarType.FP16
                and not self._multi_precision
            ):
2388 2389 2390 2391
                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."
                )
2392 2393 2394 2395
            self._add_accumulator(self._velocity_acc_str, p)

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)
2396 2397 2398 2399 2400 2401 2402 2403
        _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

2404 2405 2406
        velocity_acc = self._get_accumulator(
            self._velocity_acc_str, param_and_grad[0]
        )
2407 2408
        lr = self._create_param_lr(param_and_grad)

2409 2410 2411 2412 2413 2414 2415 2416 2417
        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
        )
2418 2419 2420

        attrs = {
            "mu": self._momentum,
2421
            "lars_coeff": self._lars_coeff,
L
limingshu 已提交
2422
            "lars_weight_decay": [_lars_weight_decay],
2423
            "multi_precision": find_master,
L
limingshu 已提交
2424
            "epsilon": self._epsilon,
2425
            "rescale_grad": self._rescale_grad,
2426 2427 2428 2429 2430 2431
        }

        inputs = {
            "Param": param_and_grad[0],
            "Grad": param_and_grad[1],
            "Velocity": velocity_acc,
2432
            "LearningRate": lr,
2433 2434 2435 2436 2437 2438 2439 2440
        }

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

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

J
Jiabin Yang 已提交
2441
        if framework._non_static_mode():
2442
            tmp, tmp2 = _legacy_C_ops.lars_momentum(
2443 2444 2445 2446 2447 2448 2449 2450 2451 2452 2453 2454 2455 2456 2457 2458 2459 2460 2461
                [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,
            )
2462 2463
        else:
            # create the momentum optimize op
2464 2465 2466 2467 2468 2469 2470
            momentum_op = block.append_op(
                type=self.type,
                inputs=inputs,
                outputs=outputs,
                attrs=attrs,
                stop_gradient=True,
            )
2471

2472
            return momentum_op
2473 2474


2475
class AdagradOptimizer(Optimizer):
2476
    r"""
2477 2478
    The Adaptive Gradient optimizer (Adagrad for short) can adaptively assign
    different learning rates to individual parameters.
Q
qiaolongfei 已提交
2479

2480
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2481 2482 2483 2484 2485 2486 2487

    .. math::

        moment\_out &= moment + grad * grad

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

2488 2489 2490 2491 2492 2493
    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 已提交
2494 2495 2496
    for numerical stability to avoid the division by zero error.

    Args:
2497 2498 2499 2500
        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 已提交
2501
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2502 2503
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2504 2505 2506 2507 2508
        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.
2509 2510 2511
        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` ,
2512
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2513 2514 2515 2516 2517
        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 已提交
2518 2519 2520 2521

    Examples:
        .. code-block:: python

2522
            import numpy as np
2523
            import paddle.fluid as fluid
2524 2525

            np_inp = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
2526
            inp = fluid.data(name="inp", shape=[2, 2])
2527 2528
            out = fluid.layers.fc(inp, size=3)
            out = fluid.layers.reduce_sum(out)
2529
            optimizer = fluid.optimizer.AdagradOptimizer(learning_rate=0.2)
2530 2531 2532 2533 2534 2535 2536
            optimizer.minimize(out)

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

2540 2541 2542 2543 2544 2545 2546 2547 2548 2549
    def __init__(
        self,
        learning_rate,
        epsilon=1.0e-6,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
        initial_accumulator_value=0.0,
    ):
2550 2551
        assert learning_rate is not None
        assert epsilon is not None
2552 2553 2554 2555 2556 2557 2558
        super(AdagradOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
2559 2560
        self.type = "adagrad"
        self._epsilon = epsilon
2561
        self.initial_accumulator_value = initial_accumulator_value
2562 2563 2564 2565 2566

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

        for p in parameters:
2567 2568 2569 2570 2571
            self._add_accumulator(
                self._moment_acc_str,
                p,
                fill_value=self.initial_accumulator_value,
            )
2572 2573 2574 2575

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

2576 2577 2578
        moment_acc = self._get_accumulator(
            self._moment_acc_str, param_and_grad[0]
        )
C
caozhou 已提交
2579
        if in_dygraph_mode():
2580 2581 2582 2583 2584 2585 2586
            _C_ops.adagrad_(
                param_and_grad[0],
                param_and_grad[1],
                moment_acc,
                self._create_param_lr(param_and_grad),
                self._epsilon,
            )
C
caozhou 已提交
2587 2588
            return None
        elif _in_legacy_dygraph():
2589 2590 2591 2592 2593 2594 2595 2596 2597 2598
            _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 已提交
2599
            return None
2600 2601 2602 2603 2604 2605 2606 2607
        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,
2608
                    "LearningRate": self._create_param_lr(param_and_grad),
2609 2610 2611
                },
                outputs={
                    "ParamOut": param_and_grad[0],
2612
                    "MomentOut": moment_acc,
2613 2614
                },
                attrs={"epsilon": self._epsilon},
2615 2616
                stop_gradient=True,
            )
2617

2618
            return adagrad_op
2619 2620 2621


class AdamOptimizer(Optimizer):
2622
    r"""
T
tianshuo78520a 已提交
2623
    The Adam optimizer uses an optimization described at the end
2624 2625 2626
    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.
2627

2628
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2629 2630 2631 2632 2633 2634 2635 2636 2637 2638 2639 2640 2641 2642

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

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

Q
qiaolongfei 已提交
2645
    Args:
2646 2647
        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.
2648 2649
        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.
2650
            The default value is 0.9.
2651 2652
        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.
2653
            The default value is 0.999.
2654 2655
        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.
2656
            The default value is 1e-08.
H
hong 已提交
2657
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2658 2659
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
2660 2661 2662 2663 2664
        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.
2665 2666 2667
        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` ,
2668
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2669 2670 2671 2672 2673 2674 2675 2676 2677 2678
        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.
2679
        use_global_beta_pow (bool, optional): Whether to use global beta_pow. If true, Adam will use global beta_pow
2680
            for whole model instead of creating beta_pow for each parameter. Default is false.
2681 2682
        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
2683
            use same align_size as allocator.
Q
qiaolongfei 已提交
2684 2685 2686 2687

    Examples:
        .. code-block:: python

2688 2689 2690 2691 2692 2693
            import paddle
            import paddle.fluid as fluid

            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
2694 2695
                x = fluid.data(name='x', shape=[None, 13], dtype='float32')
                y = fluid.data(name='y', shape=[None, 1], dtype='float32')
2696 2697 2698 2699 2700 2701 2702 2703 2704 2705 2706 2707 2708 2709 2710
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
                cost = fluid.layers.square_error_cost(input=y_predict, label=y)
                avg_cost = fluid.layers.mean(cost)

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

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

2712 2713 2714 2715 2716 2717 2718 2719 2720 2721 2722 2723 2724 2725 2726 2727 2728
        .. code-block:: python

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

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

                # define beta decay variable
2729
                def get_decayed_betas(beta1_init, beta2_init, decay_steps, decay_rate, epsilon_init):
2730 2731 2732 2733 2734 2735 2736 2737 2738 2739 2740 2741 2742 2743 2744 2745
                    global_step = lr_scheduler._decay_step_counter()

                    beta1 = fluid.layers.create_global_var(
                        shape=[1],
                        value=float(beta1_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="beta1")
                    beta2 = fluid.layers.create_global_var(
                        shape=[1],
                        value=float(beta2_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="beta2")
2746 2747 2748 2749 2750 2751 2752
                    epsilon = fluid.layers.create_global_var(
                        shape=[1],
                        value=float(epsilon_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="epsilon")
2753 2754 2755 2756 2757 2758 2759

                    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)

2760
                    return beta1, beta2, epsilon
2761

2762
                beta1, beta2, epsilon = get_decayed_betas(0.9, 0.99, 1e5, 0.9, 1e-8)
2763 2764
                adam_optimizer = fluid.optimizer.AdamOptimizer(
                                                    learning_rate=0.01,
2765
                                                    beta1=beta1,
2766 2767
                                                    beta2=beta2,
                                                    epsilon=epsilon)
2768 2769 2770 2771 2772 2773 2774 2775 2776 2777
                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)
2778 2779 2780
    """
    _moment1_acc_str = "moment1"
    _moment2_acc_str = "moment2"
Q
qiaolongfei 已提交
2781 2782
    _beta1_pow_acc_str = "beta1_pow_acc"
    _beta2_pow_acc_str = "beta2_pow_acc"
2783

2784 2785 2786 2787 2788 2789 2790 2791 2792 2793 2794 2795 2796 2797 2798
    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,
    ):
2799 2800 2801 2802
        assert learning_rate is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
2803 2804 2805 2806 2807 2808 2809 2810 2811
        super(AdamOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            flatten_param_grads=flatten_param_grads,
            align_size=align_size,
            name=name,
        )
2812 2813 2814 2815
        self.type = "adam"
        self._beta1 = beta1
        self._beta2 = beta2
        self._epsilon = epsilon
Q
Qiao Longfei 已提交
2816
        self._lazy_mode = lazy_mode
2817
        self._use_global_beta_pow = use_global_beta_pow
2818 2819 2820 2821 2822 2823

    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 已提交
2824 2825
            self._add_accumulator(self._moment1_acc_str, p)
            self._add_accumulator(self._moment2_acc_str, p)
2826 2827 2828 2829
            if not self._use_global_beta_pow:
                self._add_accumulator(
                    name=self._beta1_pow_acc_str,
                    param=p,
2830 2831 2832
                    fill_value=0.9
                    if isinstance(self._beta1, Variable)
                    else self._beta1,
2833
                    shape=[1],
2834 2835 2836
                    type=core.VarDesc.VarType.LOD_TENSOR,
                    device='cpu',
                )
2837 2838 2839
                self._add_accumulator(
                    name=self._beta2_pow_acc_str,
                    param=p,
2840 2841 2842
                    fill_value=0.999
                    if isinstance(self._beta2, Variable)
                    else self._beta2,
2843
                    shape=[1],
2844 2845 2846
                    type=core.VarDesc.VarType.LOD_TENSOR,
                    device='cpu',
                )
2847 2848
        if self._use_global_beta_pow:
            self._add_global_accumulator(
Q
qiaolongfei 已提交
2849
                name=self._beta1_pow_acc_str,
2850 2851 2852
                fill_value=0.9
                if isinstance(self._beta1, Variable)
                else self._beta1,
2853
                shape=[1],
2854 2855 2856
                type=core.VarDesc.VarType.LOD_TENSOR,
                device='cpu',
            )
2857
            self._add_global_accumulator(
Q
qiaolongfei 已提交
2858
                name=self._beta2_pow_acc_str,
2859 2860 2861
                fill_value=0.999
                if isinstance(self._beta2, Variable)
                else self._beta2,
2862
                shape=[1],
2863 2864 2865
                type=core.VarDesc.VarType.LOD_TENSOR,
                device='cpu',
            )
2866 2867 2868 2869

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

2870 2871 2872 2873 2874 2875
        moment1 = self._get_accumulator(
            self._moment1_acc_str, param_and_grad[0]
        )
        moment2 = self._get_accumulator(
            self._moment2_acc_str, param_and_grad[0]
        )
2876 2877
        if self._use_global_beta_pow:
            beta1_pow_acc = self._get_global_accumulator(
2878 2879
                self._beta1_pow_acc_str
            )
2880
            beta2_pow_acc = self._get_global_accumulator(
2881 2882
                self._beta2_pow_acc_str
            )
2883
        else:
2884 2885 2886 2887 2888 2889
            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]
            )
2890
        lr = self._create_param_lr(param_and_grad)
2891
        # create the adam optimize op
2892

J
Jiabin Yang 已提交
2893
        if framework._non_static_mode():
2894 2895 2896 2897 2898 2899 2900 2901 2902 2903
            _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)
            )
2904
            master_weight = None
2905
            _, _, _, _, _, _ = _legacy_C_ops.adam(
2906 2907 2908 2909 2910 2911 2912 2913 2914 2915 2916 2917 2918 2919 2920 2921 2922 2923 2924 2925 2926 2927 2928 2929 2930 2931 2932
                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,
            )
2933 2934 2935

            return None

2936
        inputs = {
2937 2938
            "Param": [param_and_grad[0]],
            "Grad": [param_and_grad[1]],
2939
            "LearningRate": [lr],
2940 2941 2942
            "Moment1": [moment1],
            "Moment2": [moment2],
            "Beta1Pow": [beta1_pow_acc],
2943
            "Beta2Pow": [beta2_pow_acc],
2944
        }
2945 2946 2947 2948 2949 2950 2951

        # 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

2952
        outputs = {
2953 2954 2955 2956 2957
            "ParamOut": [param_and_grad[0]],
            "Moment1Out": [moment1],
            "Moment2Out": [moment2],
            "Beta1PowOut": [beta1_pow_acc],
            "Beta2PowOut": [beta2_pow_acc],
2958 2959 2960
        }
        attrs = {
            "lazy_mode": self._lazy_mode,
2961
            "min_row_size_to_use_multithread": 1000,
2962
            'use_global_beta_pow': self._use_global_beta_pow,
2963 2964 2965 2966 2967 2968 2969 2970 2971 2972
        }

        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
2973 2974 2975 2976
        if isinstance(self._epsilon, Variable):
            inputs['EpsilonTensor'] = self._epsilon
        else:
            attrs['epsilon'] = self._epsilon
2977

2978 2979 2980 2981 2982 2983 2984
        adam_op = block.append_op(
            type=self.type,
            inputs=inputs,
            outputs=outputs,
            attrs=attrs,
            stop_gradient=True,
        )
2985 2986 2987

        return adam_op

2988
    def _finish_update(self, block, parameters_and_grads):
2989
        r"""Update beta1_pow and beta2_pow accumulator"""
2990 2991 2992
        assert isinstance(block, framework.Block)
        if self._use_global_beta_pow:
            beta1_pow_acc = self._get_global_accumulator(
2993 2994
                self._beta1_pow_acc_str
            )
2995
            beta2_pow_acc = self._get_global_accumulator(
2996 2997
                self._beta2_pow_acc_str
            )
2998 2999 3000

            with block.program._optimized_guard([]):
                inputs = {"X": beta1_pow_acc}
3001
                outputs = {"Out": beta1_pow_acc}
3002 3003
                attrs = {}
                if isinstance(self._beta1, Variable):
3004 3005
                    inputs["Y"] = self._beta1
                    # use elementwise_mul for better performance
3006 3007 3008 3009 3010 3011 3012
                    block.append_op(
                        type="elementwise_mul",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
3013 3014
                else:
                    attrs['scale'] = self._beta1
3015 3016 3017 3018 3019 3020 3021
                    block.append_op(
                        type="scale",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
3022 3023

                inputs = {"X": beta2_pow_acc}
3024
                outputs = {"Out": beta2_pow_acc}
3025 3026
                attrs = {}
                if isinstance(self._beta2, Variable):
3027 3028
                    inputs["Y"] = self._beta2
                    # use elementwise_mul for better performance
3029 3030 3031 3032 3033 3034 3035
                    block.append_op(
                        type="elementwise_mul",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
3036 3037
                else:
                    attrs['scale'] = self._beta2
3038 3039 3040 3041 3042 3043 3044
                    block.append_op(
                        type="scale",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
3045

3046 3047

class AdamaxOptimizer(Optimizer):
3048
    r"""
3049
    The Adamax optimizer is implemented based on the Adamax Optimization
3050 3051 3052
    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 已提交
3053

3054
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
3055 3056 3057 3058 3059 3060 3061 3062 3063 3064 3065 3066 3067

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

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

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

3101 3102 3103 3104 3105 3106 3107 3108 3109 3110 3111 3112 3113
    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          import numpy

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

          train_program = fluid.Program()
          startup_program = fluid.Program()
          with fluid.program_guard(train_program, startup_program):
3114
              data = fluid.data(name='X', shape=[None, 1], dtype='float32')
3115 3116
              hidden = fluid.layers.fc(input=data, size=10)
              loss = fluid.layers.mean(hidden)
3117
              adam = fluid.optimizer.AdamaxOptimizer(learning_rate=0.2)
3118 3119 3120 3121 3122 3123 3124 3125 3126
              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])
3127 3128 3129
    """
    _moment_acc_str = "moment"
    _inf_norm_acc_str = "inf_norm"
Q
qiaolongfei 已提交
3130
    _beta1_pow_acc_str = "beta1_pow_acc"
3131

3132 3133 3134 3135 3136 3137 3138 3139 3140 3141 3142
    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,
    ):
3143 3144 3145 3146
        assert learning_rate is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
3147 3148 3149 3150 3151 3152 3153
        super(AdamaxOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
3154 3155 3156 3157 3158 3159 3160 3161
        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 已提交
3162 3163
            self._add_accumulator(self._moment_acc_str, p)
            self._add_accumulator(self._inf_norm_acc_str, p)
3164 3165 3166 3167 3168 3169
            self._add_accumulator(
                name=self._beta1_pow_acc_str,
                param=p,
                fill_value=self._beta1,
                shape=[1],
            )
3170 3171 3172 3173 3174

    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])
3175 3176 3177 3178 3179 3180
        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]
        )
3181 3182

        if framework.in_dygraph_mode():
3183 3184 3185 3186 3187 3188 3189 3190 3191 3192 3193
            _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,
            )
3194
        elif framework._in_legacy_dygraph():
3195 3196 3197 3198 3199 3200 3201 3202 3203 3204 3205 3206 3207 3208 3209 3210 3211
            _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,
            )
3212 3213 3214 3215 3216 3217 3218 3219 3220 3221
        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,
3222
                    "Beta1Pow": beta1_pow_acc,
3223 3224 3225 3226
                },
                outputs={
                    "ParamOut": param_and_grad[0],
                    "MomentOut": moment,
3227
                    "InfNormOut": inf_norm,
3228 3229 3230 3231
                },
                attrs={
                    "beta1": self._beta1,
                    "beta2": self._beta2,
3232
                    "epsilon": self._epsilon,
3233
                },
3234 3235
                stop_gradient=True,
            )
3236

3237
            return adamax_op
3238

3239
    def _finish_update(self, block, parameters_and_grads):
3240
        """Update Beta1 Power accumulator"""
3241
        assert isinstance(block, framework.Block)
3242
        for param, grad in parameters_and_grads:
C
chengduo 已提交
3243
            if grad is None or param.trainable is False:
3244
                continue
3245 3246 3247 3248 3249 3250
            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 已提交
3251
                if framework._non_static_mode():
3252
                    if framework.in_dygraph_mode():
3253 3254 3255
                        tmp = _C_ops.scale(
                            beta1_pow_acc, self._beta1, 0.0, True
                        )
3256
                    else:
3257 3258 3259
                        tmp = _legacy_C_ops.scale(
                            beta1_pow_acc, "scale", self._beta1
                        )
3260 3261
                    beta1_pow_acc.copy_(tmp, False)
                else:
3262 3263 3264 3265 3266 3267 3268
                    block.append_op(
                        type="scale",
                        inputs={"X": beta1_pow_acc},
                        outputs={"Out": beta1_pow_acc},
                        attrs={"scale": self._beta1},
                        stop_gradient=True,
                    )
3269 3270


3271
class DpsgdOptimizer(Optimizer):
3272
    r"""
3273 3274 3275 3276 3277 3278 3279 3280 3281 3282 3283 3284 3285 3286 3287 3288 3289 3290 3291 3292 3293 3294 3295 3296 3297 3298 3299 3300 3301 3302 3303 3304 3305 3306 3307 3308
    We implement the Dpsgd optimizer according to CCS16 paper -
    Deep Learning with Differential Privacy.

    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          import numpy

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

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

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

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

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

3316 3317 3318 3319 3320 3321 3322 3323
    def __init__(
        self,
        learning_rate=0.001,
        clip=0.9,
        batch_size=0.999,
        sigma=1e-8,
        parameter_list=None,
    ):
3324 3325 3326 3327
        assert learning_rate is not None
        assert clip is not None
        assert batch_size is not None
        assert sigma is not None
3328 3329 3330
        super(DpsgdOptimizer, self).__init__(
            learning_rate=learning_rate, parameter_list=parameter_list
        )
3331 3332 3333 3334
        self.type = "dpsgd"
        self._clip = clip
        self._batch_size = batch_size
        self._sigma = sigma
Z
zhongpu 已提交
3335 3336 3337 3338 3339 3340 3341
        '''
        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
3342 3343 3344 3345 3346

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

        # create the dpsgd optimize op
3347
        if self._seed is None:
Z
zhongpu 已提交
3348 3349
            self._seed = 0

J
Jiabin Yang 已提交
3350
        if framework._non_static_mode():
3351 3352 3353 3354 3355 3356 3357 3358 3359 3360 3361 3362 3363 3364
            _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,
            )
3365
        else:
3366 3367 3368 3369 3370 3371 3372 3373 3374 3375 3376 3377 3378 3379 3380 3381
            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,
            )
3382

3383
            return dpsgd_op
3384 3385


3386
class DecayedAdagradOptimizer(Optimizer):
3387
    r"""
3388 3389 3390
    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.
3391

3392
    The parameter ``param_out`` update rule with gradient ``grad``:
3393 3394 3395 3396 3397 3398 3399

    .. math::

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

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

3400 3401 3402 3403
    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
3404 3405 3406
    stability to avoid the division by zero error.

    Args:
3407 3408 3409 3410 3411
        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 已提交
3412
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3413 3414
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3415 3416 3417 3418 3419
        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.
3420 3421 3422
        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` ,
3423
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3424 3425 3426 3427 3428 3429
        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.**
3430 3431 3432 3433

    Examples:
        .. code-block:: python

3434 3435
            import paddle.fluid as fluid

3436 3437 3438 3439
            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)
3440
            optimizer.minimize(cost)
3441 3442 3443
    """
    _moment_acc_str = "moment"

3444 3445 3446 3447 3448 3449 3450 3451 3452 3453
    def __init__(
        self,
        learning_rate,
        decay=0.95,
        epsilon=1.0e-6,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
3454 3455 3456 3457
        assert learning_rate is not None
        assert decay is not None
        assert epsilon is not None

3458 3459 3460 3461 3462 3463 3464
        super(DecayedAdagradOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
3465 3466 3467 3468 3469 3470 3471 3472 3473 3474 3475 3476 3477
        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)

3478 3479 3480
        moment_acc = self._get_accumulator(
            self._moment_acc_str, param_and_grad[0]
        )
3481

J
Jiabin Yang 已提交
3482
        if framework._non_static_mode():
3483 3484 3485 3486 3487 3488 3489 3490 3491 3492 3493 3494
            _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,
            )
3495 3496 3497 3498 3499 3500 3501 3502
        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,
3503
                    "LearningRate": self._create_param_lr(param_and_grad),
3504 3505 3506
                },
                outputs={
                    "ParamOut": param_and_grad[0],
3507
                    "MomentOut": moment_acc,
3508
                },
3509 3510 3511
                attrs={"epsilon": self._epsilon, "decay": self._decay},
                stop_gradient=True,
            )
3512

3513
            return decayed_adagrad_op
3514 3515


3516
class AdadeltaOptimizer(Optimizer):
3517
    r"""
Z
Zeng Jinle 已提交
3518
    **Notes: This API does not support sparse parameter optimization.**
Q
qiaolongfei 已提交
3519

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

    The update is done as follows:
3524

Z
Zeng Jinle 已提交
3525 3526
    .. math::

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

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

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

    Args:
Z
Zeng Jinle 已提交
3534 3535 3536
        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 已提交
3537
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3538 3539
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3540 3541 3542 3543 3544
        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.
3545 3546 3547
        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` ,
3548
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3549 3550 3551
        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` .
3552 3553 3554 3555

    Examples:
        .. code-block:: python

3556
            import paddle.fluid as fluid
Z
Zeng Jinle 已提交
3557

3558
            image = fluid.data(name='image', shape=[None, 28], dtype='float32')
Z
Zeng Jinle 已提交
3559 3560
            fc = fluid.layers.fc(image, size=10)
            cost = fluid.layers.reduce_mean(fc)
3561 3562
            optimizer = fluid.optimizer.Adadelta(
                learning_rate=0.0003, epsilon=1.0e-6, rho=0.95)
C
chengduo 已提交
3563

Z
Zeng Jinle 已提交
3564 3565 3566 3567
            # 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)
3568
    """
3569

3570 3571 3572
    _avg_squared_grad_acc_str = "_avg_squared_grad"
    _avg_squared_update_acc_str = "_avg_squared_update"

3573 3574 3575 3576 3577 3578 3579 3580 3581 3582
    def __init__(
        self,
        learning_rate,
        epsilon=1.0e-6,
        rho=0.95,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
3583 3584 3585 3586 3587 3588
        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.")
3589 3590 3591 3592 3593 3594 3595
        super(AdadeltaOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
3596 3597 3598 3599 3600
        self.type = "adadelta"
        self._epsilon = epsilon
        self._rho = rho

    def _create_accumulators(self, block, parameters):
3601 3602
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")
3603 3604 3605 3606 3607 3608

        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):
3609 3610
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")
3611 3612

        avg_squared_grad_acc = self._get_accumulator(
3613 3614
            self._avg_squared_grad_acc_str, param_and_grad[0]
        )
3615
        avg_squared_update_acc = self._get_accumulator(
3616 3617
            self._avg_squared_update_acc_str, param_and_grad[0]
        )
3618

3619
        if framework.in_dygraph_mode():
3620 3621 3622 3623 3624 3625 3626 3627
            _C_ops.adadelta_(
                param_and_grad[0],
                param_and_grad[1],
                avg_squared_grad_acc,
                avg_squared_update_acc,
                self._rho,
                self._epsilon,
            )
3628
        elif framework._in_legacy_dygraph():
3629 3630 3631 3632 3633 3634 3635 3636 3637 3638 3639 3640 3641
            _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,
            )
3642 3643
        else:
            # Create the adadelta optimizer op
3644 3645 3646 3647 3648 3649 3650 3651 3652 3653 3654 3655 3656 3657 3658 3659
            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,
            )
3660

3661
            return adadelta_op
3662 3663


Q
qingqing01 已提交
3664
class RMSPropOptimizer(Optimizer):
3665
    r"""
Q
qingqing01 已提交
3666 3667 3668 3669 3670 3671 3672 3673
    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 已提交
3674
        r(w, t) & = \\rho r(w, t-1) + (1 - \\rho)(\\nabla Q_{i}(w))^2
Q
qingqing01 已提交
3675 3676 3677 3678

        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 已提交
3679
    each weight. Then dividing the gradient by :math:`sqrt{v(w,t)}`.
Q
qingqing01 已提交
3680 3681 3682 3683 3684 3685

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

    ..  math::

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

3688 3689 3690 3691 3692 3693 3694 3695 3696 3697 3698 3699 3700 3701
        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 已提交
3702 3703 3704 3705
            \\epsilon}} \\nabla Q_{i}(w)

        w & = w - v(w, t)

Q
qiaolongfei 已提交
3706
    where, :math:`\\rho` is a hyperparameter and typical values are 0.9, 0.95
Q
qingqing01 已提交
3707 3708 3709 3710 3711
    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.


3712 3713 3714
    Parameters:
        learning_rate(float): Global learning rate.
        rho(float): rho is :math: `\\rho` in equation, default is 0.95.
Q
qingqing01 已提交
3715
        epsilon(float): :math: `\\epsilon` in equation is smoothing term to
3716
            avoid division by zero, default is 1e-6.
Q
qiaolongfei 已提交
3717
        momentum(float): :math:`\\beta` in equation is the momentum term,
3718
            default is 0.0.
3719 3720 3721 3722
        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 已提交
3723
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3724 3725
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3726 3727 3728 3729 3730
        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.
3731 3732 3733
        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` ,
3734
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3735 3736
        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 已提交
3737 3738 3739 3740 3741 3742 3743

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

    Examples:
          .. code-block:: python

3744 3745 3746 3747 3748 3749 3750 3751 3752 3753 3754 3755 3756 3757 3758 3759 3760 3761 3762 3763 3764 3765 3766 3767 3768
            import paddle
            import paddle.fluid as fluid
            import numpy as np

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

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

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

Q
qingqing01 已提交
3769 3770 3771 3772
    """

    _momentum_acc_str = "momentum"
    _mean_square_acc_str = "mean_square"
3773
    _mean_grad_acc_str = "mean_grad"
Q
qingqing01 已提交
3774

3775 3776 3777 3778 3779 3780 3781 3782 3783 3784 3785 3786 3787 3788 3789 3790 3791 3792 3793
    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,
    ):
        super(RMSPropOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
Q
qingqing01 已提交
3794 3795 3796 3797 3798 3799 3800 3801 3802 3803 3804 3805 3806
        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
3807
        self._centered = centered
Q
qingqing01 已提交
3808 3809 3810 3811 3812 3813 3814 3815

    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)
3816
            self._add_accumulator(self._mean_grad_acc_str, p)
Q
qingqing01 已提交
3817 3818 3819 3820 3821

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

3822 3823 3824 3825 3826 3827 3828 3829 3830
        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 已提交
3831
        if in_dygraph_mode():
3832 3833 3834 3835 3836 3837 3838 3839 3840 3841 3842 3843
            _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 已提交
3844 3845
            return None
        elif _in_legacy_dygraph():
3846 3847 3848 3849 3850 3851 3852 3853 3854 3855 3856 3857 3858 3859 3860 3861 3862 3863 3864
            _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 已提交
3865
            return None
3866 3867 3868 3869 3870 3871 3872 3873 3874 3875 3876 3877 3878 3879 3880
        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,
3881
                    "MeanGradOut": mean_grad_acc,
3882 3883 3884 3885 3886
                },
                attrs={
                    "epsilon": self._epsilon,
                    "decay": self._rho,
                    "momentum": self._momentum,
3887
                    "centered": self._centered,
3888
                },
3889 3890
                stop_gradient=True,
            )
Q
qingqing01 已提交
3891

3892
            return rmsprop_op
Q
qingqing01 已提交
3893 3894


Q
qiaolongfei 已提交
3895
class FtrlOptimizer(Optimizer):
3896
    r"""
Q
qiaolongfei 已提交
3897 3898 3899 3900 3901 3902 3903 3904 3905 3906 3907 3908 3909 3910 3911 3912 3913 3914 3915 3916 3917 3918 3919 3920 3921 3922 3923 3924 3925 3926 3927 3928 3929 3930 3931 3932 3933 3934
    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

3935 3936 3937 3938 3939
    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 已提交
3940
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3941 3942
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
3943 3944 3945 3946 3947
        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.
3948 3949 3950
        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` ,
3951
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3952 3953
        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 已提交
3954 3955 3956 3957 3958 3959 3960

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

    Examples:
          .. code-block:: python

3961 3962 3963 3964 3965 3966 3967 3968 3969 3970 3971 3972 3973 3974 3975 3976 3977 3978 3979 3980 3981 3982 3983 3984
            import paddle
            import paddle.fluid as fluid
            import numpy as np

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

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

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

3986
    NOTE:
C
chengduo 已提交
3987
       Currently, FtrlOptimizer doesn't support sparse parameter optimization.
Q
qiaolongfei 已提交
3988 3989 3990 3991 3992
    """

    _squared_acc_str = "squared"
    _linear_acc_str = "linear"

3993 3994 3995 3996 3997 3998 3999 4000 4001 4002 4003 4004 4005 4006 4007 4008 4009 4010
    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,
    ):
        super(FtrlOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
Q
qiaolongfei 已提交
4011 4012 4013 4014 4015 4016 4017 4018 4019 4020 4021 4022 4023 4024 4025 4026 4027 4028 4029 4030
        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.")

4031 4032 4033 4034 4035 4036
        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 已提交
4037
        if framework._non_static_mode():
4038 4039 4040 4041 4042 4043 4044 4045 4046 4047 4048 4049 4050 4051 4052 4053
            _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,
            )
4054 4055

        else:
4056 4057 4058 4059 4060 4061 4062 4063 4064 4065 4066 4067 4068 4069 4070 4071 4072 4073 4074 4075 4076
            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 已提交
4077

4078
            return ftrl_op
Q
qiaolongfei 已提交
4079 4080


Y
Yibing Liu 已提交
4081
class LambOptimizer(AdamOptimizer):
4082
    r"""
Y
Yibing Liu 已提交
4083 4084
    LAMB (Layer-wise Adaptive Moments optimizer for Batching training) Optimizer.

4085 4086 4087
    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 已提交
4088
    Deep Learning: Training BERT in 76 minutes <https://arxiv.org/abs/1904.00962>`_ .
Y
Yibing Liu 已提交
4089 4090 4091 4092 4093

    The updating of parameters follows:

    ..  math::

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

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

4098 4099 4100 4101
        m_t &= \\frac{m_t}{\\beta_1^t}

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

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

Y
Yibing Liu 已提交
4104
        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 已提交
4105 4106


4107
    where :math:`m` is the 1st moment, and :math:`v` the 2nd moment, :math:`\\eta` the
Y
Yibing Liu 已提交
4108 4109 4110
    learning rate, :math:`\\lambda` the LAMB weight decay rate.

    Args:
Y
Yibing Liu 已提交
4111 4112 4113 4114 4115 4116 4117 4118
        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 已提交
4119
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
4120 4121
            This parameter is required in dygraph mode. \
            The default value is None in static mode, at this time all parameters will be updated.
4122 4123 4124 4125 4126
        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.
4127 4128
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of
            some derived class of ``GradientClipBase`` . There are three cliping strategies
4129 4130 4131
            ( :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.
4132 4133
        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 已提交
4134
            Default None.
4135
        name(str|None): For detailed information, please refer to
Y
Yibing Liu 已提交
4136
            :ref:`api_guide_Name` . Usually name is no need to set and None by default.
Y
Yibing Liu 已提交
4137 4138 4139

    Examples:
        .. code-block:: python
4140 4141

            import paddle.fluid as fluid
Y
Yibing Liu 已提交
4142

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

Y
Yibing Liu 已提交
4147 4148 4149 4150 4151
            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 已提交
4152 4153 4154 4155 4156 4157 4158
            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"

4159 4160 4161 4162 4163 4164 4165 4166 4167 4168 4169 4170 4171
    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 已提交
4172 4173 4174 4175 4176
        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
4177 4178 4179 4180 4181 4182 4183 4184 4185 4186
        super(LambOptimizer, self).__init__(
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            beta1=beta1,
            beta2=beta2,
            epsilon=epsilon,
            name=name,
        )
Y
Yibing Liu 已提交
4187 4188
        self.type = "lamb"
        self._weight_decay = lamb_weight_decay
Y
Yibing Liu 已提交
4189
        self._exclude_from_weight_decay_fn = exclude_from_weight_decay_fn
Y
Yibing Liu 已提交
4190 4191 4192

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

4195 4196 4197 4198 4199 4200 4201 4202 4203 4204 4205 4206 4207 4208 4209 4210 4211
        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 已提交
4212 4213 4214
            weight_decay = 0.0
        else:
            weight_decay = self._weight_decay
4215
        lr = self._create_param_lr(param_and_grad)
4216
        master_weight = None
J
Jiabin Yang 已提交
4217
        if framework._non_static_mode():
4218 4219 4220 4221 4222 4223 4224 4225 4226 4227 4228 4229 4230 4231 4232 4233 4234 4235 4236 4237 4238 4239 4240 4241
            _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,
            )
4242
            return None
Y
Yibing Liu 已提交
4243

Y
Yibing Liu 已提交
4244
        # create the lamb optimize op
4245 4246 4247 4248 4249 4250 4251 4252 4253 4254 4255 4256 4257 4258 4259 4260 4261 4262 4263 4264 4265 4266 4267 4268 4269 4270
        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 已提交
4271 4272 4273 4274

        return lamb_op


4275 4276 4277 4278 4279 4280 4281 4282 4283 4284 4285 4286 4287
# 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
4288
Dpsgd = DpsgdOptimizer
4289
DecayedAdagrad = DecayedAdagradOptimizer
4290
Adadelta = AdadeltaOptimizer
Q
qingqing01 已提交
4291
RMSProp = RMSPropOptimizer
Q
qiaolongfei 已提交
4292
Ftrl = FtrlOptimizer
4293
LarsMomentum = LarsMomentumOptimizer
Y
Yibing Liu 已提交
4294
Lamb = LambOptimizer
4295 4296 4297


class ModelAverage(Optimizer):
4298
    r"""
4299
	:api_attr: Static Graph
S
swtkiwi 已提交
4300

4301 4302 4303 4304 4305 4306 4307 4308 4309 4310 4311 4312 4313 4314 4315 4316 4317 4318
    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:

    ::
4319

4320 4321 4322 4323 4324 4325 4326 4327 4328
        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.
4329 4330

    Args:
4331 4332 4333
        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.
4334 4335 4336 4337 4338
        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.
4339 4340 4341
        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.
4342

4343
    Examples:
Q
qiaolongfei 已提交
4344 4345 4346

      .. code-block:: python

4347 4348 4349 4350 4351 4352
        import paddle.fluid as fluid
        import numpy

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

4354 4355 4356 4357
        train_program = fluid.Program()
        startup_program = fluid.Program()
        with fluid.program_guard(train_program, startup_program):
            # build net
4358
            data = fluid.data(name='X', shape=[None, 1], dtype='float32')
4359 4360 4361 4362 4363 4364 4365 4366
            hidden = fluid.layers.fc(input=data, size=10)
            loss = fluid.layers.mean(hidden)
            optimizer = fluid.optimizer.Momentum(learning_rate=0.2, momentum=0.1)
            optimizer.minimize(loss)

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

            exe.run(startup_program)
4370 4371 4372 4373 4374
            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])
4375 4376

            # apply ModelAverage
4377
            with model_average.apply(exe):
4378 4379 4380 4381
                x = numpy.random.random(size=(10, 1)).astype('float32')
                exe.run(program=train_program,
                        feed={'X': x},
                        fetch_list=[loss.name])
4382 4383
    """

4384 4385 4386 4387 4388 4389 4390 4391
    def __init__(
        self,
        average_window_rate,
        min_average_window=10000,
        max_average_window=10000,
        regularization=None,
        name=None,
    ):
J
Jiabin Yang 已提交
4392
        if framework._non_static_mode():
Z
zhongpu 已提交
4393
            raise Exception("In dygraph, don't support ModelAverage.")
4394 4395 4396
        super(ModelAverage, self).__init__(
            0.0, regularization=regularization, name=name
        )
4397 4398 4399
        self.average_window = average_window_rate
        self.min_average_window = min_average_window
        self.max_average_window = max_average_window
4400

4401
        self.params_grads = []
4402 4403 4404
        for param in (
            framework.default_main_program().global_block().all_parameters()
        ):
4405
            if param.do_model_average != False:
4406
                grad = param.block.create_var(
4407 4408 4409
                    name=unique_name.generate_with_ignorable_key(
                        ".".join([param.name, 'tmp'])
                    ),
4410 4411
                    dtype=param.dtype,
                    persistable=False,
4412 4413
                    stop_gradient=True,
                )
4414
                self.params_grads.append((param, grad))
4415

4416
        for param, grad in self.params_grads:
4417 4418
            if grad is None:
                continue
X
Xin Pan 已提交
4419
            with param.block.program._optimized_guard(
4420 4421
                [param, grad]
            ), name_scope('move_average'):
4422
                self._append_average_accumulate_op(param)
4423

4424 4425 4426 4427
        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:
4428
                self._add_average_apply_op(block, param_grad)
4429 4430 4431 4432 4433

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

4436
    def _add_average_apply_op(self, block, param_grad):
L
Luo Tao 已提交
4437 4438 4439 4440 4441 4442
        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(
4443 4444
            self._get_accumulator('num_accumulates', param)
        )
L
Luo Tao 已提交
4445
        old_num_accumulates = block._clone_variable(
4446 4447
            self._get_accumulator('old_num_accumulates', param)
        )
L
Luo Tao 已提交
4448
        num_updates = block._clone_variable(
4449 4450
            self._get_accumulator('num_updates', param)
        )
4451 4452 4453 4454 4455
        # backup param value to grad
        layers.assign(input=param, output=grad)
        # param = (sum_1 + sum_2 + sum_3) / (num_accumulates + old_num_accumulates)
        tmp = layers.sum(x=[num_accumulates, old_num_accumulates])
        sum = layers.sum(x=[sum_1, sum_2, sum_3])
D
dzhwinter 已提交
4456
        tmp = layers.cast(
4457
            x=tmp, dtype='float32' if self._dtype is None else self._dtype
4458
        )
D
dzhwinter 已提交
4459
        sum = layers.cast(
4460
            x=sum, dtype='float32' if self._dtype is None else self._dtype
4461
        )
S
sneaxiy 已提交
4462
        ops._elementwise_div(x=sum, y=tmp, out=param)
4463 4464

    def _add_average_restore_op(self, block, param_grad):
L
Luo Tao 已提交
4465 4466
        param = block._clone_variable(param_grad[0])
        grad = block._clone_variable(param_grad[1])
4467 4468 4469 4470 4471 4472 4473
        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)
4474 4475 4476 4477 4478 4479 4480 4481 4482 4483 4484 4485 4486 4487 4488 4489 4490 4491 4492 4493 4494 4495 4496 4497 4498 4499 4500 4501 4502 4503 4504 4505 4506 4507 4508 4509
        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,
        )
4510

S
rename  
sneaxiy 已提交
4511
    @signature_safe_contextmanager
4512
    def apply(self, executor, need_restore=True):
4513 4514
        """
        Apply the average of the cumulative ``Parameter`` to the parameters of the current model.
4515 4516

        Args:
4517 4518 4519 4520 4521 4522 4523 4524 4525 4526 4527 4528 4529 4530 4531 4532 4533 4534 4535 4536 4537 4538 4539 4540 4541 4542 4543 4544 4545 4546 4547 4548 4549 4550 4551 4552 4553 4554 4555 4556 4557 4558 4559 4560
            executor(fluid.Executor): The current network executor.
            need_restore(bool): Restore flag variable, if set to True, the network will restore
                the parameters of the network to the default value, if set to False,
                it will not be restored. The default value is True.

        Examples:

          .. code-block:: python

            import paddle.fluid as fluid
            import numpy

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

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

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

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

                # apply ModelAverage
                with model_average.apply(exe):
                    x = numpy.random.random(size=(10, 1)).astype('float32')
                    exe.run(program=train_program,
                            feed={'X': x},
                            fetch_list=[loss.name])
4561
        """
4562 4563 4564 4565 4566 4567
        executor.run(self.apply_program)
        try:
            yield
        finally:
            if need_restore:
                self.restore(executor)
4568 4569

    def restore(self, executor):
4570 4571
        """
        Restore ``Parameter`` values of current model.
4572

4573
        Args:
4574 4575 4576 4577 4578 4579 4580 4581 4582 4583 4584 4585 4586 4587 4588 4589 4590 4591 4592 4593 4594 4595 4596 4597 4598 4599 4600 4601 4602 4603 4604 4605 4606 4607 4608 4609 4610 4611 4612 4613 4614 4615 4616 4617
            executor(fluid.Executor): The current network executor.

        Examples:

          .. code-block:: python

            import paddle.fluid as fluid
            import numpy

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

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

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

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

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

                # restore Parameters
                model_average.restore(exe)
4618
        """
4619
        executor.run(self.restore_program)
4620 4621 4622


class ExponentialMovingAverage(object):
4623
    r"""
4624
        :api_attr: Static Graph
S
swtkiwi 已提交
4625

4626 4627 4628 4629 4630 4631
    Compute the moving average of parameters with exponential decay.
    Given a parameter :math:`\\theta`, its exponential moving average (EMA)
    will be

    ..  math::

4632
        \\text{EMA}_0 & = 0
4633

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

4636 4637 4638
    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 已提交
4639
    the **restore()** method is used to restore the parameters.
4640

4641 4642 4643 4644
    **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
4645 4646

    ..  math::
4647

4648 4649
        \\widehat{\\text{EMA}}_t = \\frac{\\text{EMA}_t}{1 - \\text{decay}^t}

4650 4651
    **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
4652
    relative smaller decay rate in the very beginning. The argument **thres_steps**
4653
    allows users to pass a Variable to schedule the decay rate, in this case,
4654
    the actual decay rate becomes
4655

4656
    ..  math::
4657

4658 4659 4660
        \\min(\\text{decay}, \\frac{1 + \\text{thres_steps}}{10 + \\text{thres_steps}})

    Usually **thres_steps** can be the global training steps.
4661 4662 4663


    Args:
4664 4665 4666
        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.
4667 4668 4669 4670


    Examples:

4671 4672 4673 4674 4675 4676 4677 4678 4679 4680 4681 4682 4683 4684 4685 4686 4687 4688 4689 4690 4691 4692 4693 4694 4695 4696 4697 4698
        .. 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(),
4699
                    feed={'x': data},
4700 4701 4702 4703 4704 4705
                    fetch_list=[cost.name])

                # usage 1
                with ema.apply(exe):
                    data = numpy.random.random(size=(10, 5)).astype('float32')
                    exe.run(program=test_program,
4706
                        feed={'x': data},
4707 4708 4709 4710 4711 4712
                        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,
4713
                        feed={'x': data},
4714 4715 4716
                        fetch_list=[hidden.name])
                ema.restore(exe)

4717 4718
    """

4719
    def __init__(self, decay=0.999, thres_steps=None, name=None):
J
Jiabin Yang 已提交
4720
        if framework._non_static_mode():
Z
zhongpu 已提交
4721
            raise Exception(
4722 4723
                "In dygraph, don't support ExponentialMovingAverage."
            )
4724
        self._decay = decay
4725
        self._thres_steps = thres_steps
4726
        self._name = name if name is not None else ''
4727 4728
        self._decay_var = self._get_ema_decay()

4729
        self._step_counter_name = "@EMA_STEP_COUNTER@"
Y
Yibing Liu 已提交
4730
        self._params_tmps = []
4731
        for param in default_main_program().global_block().all_parameters():
4732
            if param.do_model_average != False:
4733 4734 4735 4736 4737 4738 4739 4740
                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 已提交
4741
                self._params_tmps.append((param, tmp))
4742

Y
Yibing Liu 已提交
4743 4744
        self._ema_vars = {}
        for param, tmp in self._params_tmps:
4745 4746 4747
            with param.block.program._optimized_guard([param, tmp]), name_scope(
                'moving_average'
            ):
Y
Yibing Liu 已提交
4748
                self._ema_vars[param.name] = self._create_ema_vars(param)
4749 4750 4751 4752

        self.apply_program = Program()
        block = self.apply_program.global_block()
        with program_guard(main_program=self.apply_program):
4753
            decay_pow, global_step = self._get_decay_pow(block)
Y
Yibing Liu 已提交
4754
            for param, tmp in self._params_tmps:
4755 4756
                param = block._clone_variable(param)
                tmp = block._clone_variable(tmp)
Y
Yibing Liu 已提交
4757
                ema = block._clone_variable(self._ema_vars[param.name])
4758
                layers.assign(input=param, output=tmp)
4759
                # bias correction
4760 4761
                with layers.control_flow.Switch() as switch:
                    with switch.case(global_step > 0):
4762 4763 4764
                        layers.assign(
                            output=param, input=ema / (1.0 - decay_pow)
                        )
4765 4766
                    with switch.default():
                        layers.assign(output=param, input=ema)
4767 4768 4769 4770

        self.restore_program = Program()
        block = self.restore_program.global_block()
        with program_guard(main_program=self.restore_program):
Y
Yibing Liu 已提交
4771
            for param, tmp in self._params_tmps:
4772 4773 4774 4775
                tmp = block._clone_variable(tmp)
                param = block._clone_variable(param)
                layers.assign(input=tmp, output=param)

4776 4777 4778 4779 4780 4781 4782
    def _get_ema_decay(self):
        with default_main_program()._lr_schedule_guard():
            decay_var = layers.tensor.create_global_var(
                shape=[1],
                value=self._decay,
                dtype='float32',
                persistable=True,
4783 4784
                name="scheduled_ema_decay_rate",
            )
4785 4786 4787 4788 4789 4790 4791 4792

            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(
4793 4794
                            np.array([self._decay], dtype=np.float32), decay_var
                        )
4795 4796 4797
        return decay_var

    def _get_decay_pow(self, block):
4798 4799 4800 4801 4802 4803 4804
        global_step = layers.create_global_var(
            name=self._step_counter_name,
            shape=[1],
            value=0,
            dtype='int64',
            persistable=True,
        )
4805
        global_step = layers.cast(global_step, "float32")
4806
        decay_var = block._clone_variable(self._decay_var)
4807 4808
        decay_pow_acc = layers.elementwise_pow(decay_var, global_step)
        return decay_pow_acc, global_step
4809

Y
Yibing Liu 已提交
4810
    def _create_ema_vars(self, param):
4811 4812 4813 4814 4815
        param_ema = layers.create_global_var(
            name=unique_name.generate(self._name + param.name + '_ema'),
            shape=param.shape,
            value=0.0,
            dtype=param.dtype,
4816 4817
            persistable=True,
        )
4818 4819 4820

        return param_ema

Y
Yibing Liu 已提交
4821
    def update(self):
4822 4823
        """
        Update Exponential Moving Average. Should only call this method in
Y
Yibing Liu 已提交
4824 4825
        train program.
        """
4826
        global_step = layers.autoincreased_step_counter(
4827 4828
            counter_name=self._step_counter_name
        )
4829
        param_master_emas = []
Y
Yibing Liu 已提交
4830
        for param, tmp in self._params_tmps:
4831 4832 4833
            with param.block.program._optimized_guard([param, tmp]), name_scope(
                'moving_average'
            ):
Y
Yibing Liu 已提交
4834
                param_ema = self._ema_vars[param.name]
4835
                if param.name + '.master' in self._ema_vars:
4836 4837 4838 4839
                    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 * (
4840 4841
                        1 - self._decay_var
                    )
4842 4843 4844 4845 4846 4847 4848 4849 4850 4851
                    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,
4852 4853 4854
                    "out_dtype": param_ema.dtype,
                },
            )
Y
Yibing Liu 已提交
4855

4856 4857 4858 4859
    @signature_safe_contextmanager
    def apply(self, executor, need_restore=True):
        """
        Apply moving average to parameters for evaluation.
4860

4861 4862
        Args:
            executor (Executor): The Executor to execute applying.
4863
            need_restore (bool, optional): Whether to restore parameters after
Y
Yibing Liu 已提交
4864
                applying. Default True.
4865 4866 4867 4868 4869 4870 4871 4872 4873 4874
        """
        executor.run(self.apply_program)
        try:
            yield
        finally:
            if need_restore:
                self.restore(executor)

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

4876 4877 4878 4879
        Args:
            executor (Executor): The Executor to execute restoring.
        """
        executor.run(self.restore_program)
H
hutuxian 已提交
4880 4881 4882


class PipelineOptimizer(object):
4883
    """
4884
        :api_attr: Static Graph
S
swtkiwi 已提交
4885

4886 4887 4888 4889
    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 已提交
4890

4891
    Args:
4892 4893 4894
        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].
4895

4896 4897
    Examples:
        .. code-block:: python
H
hutuxian 已提交
4898

4899
            import paddle.fluid as fluid
H
hutuxian 已提交
4900 4901
            import paddle.fluid.layers as layers

4902 4903 4904 4905 4906 4907 4908 4909 4910 4911 4912 4913 4914 4915 4916 4917
            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 已提交
4918
            optimizer = fluid.optimizer.SGD(learning_rate=0.5)
4919
            optimizer = fluid.optimizer.PipelineOptimizer(optimizer)
H
hutuxian 已提交
4920
            optimizer.minimize(loss)
4921 4922 4923 4924 4925 4926 4927 4928 4929

            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 已提交
4930 4931
            exe = fluid.Executor(place)
            exe.run(fluid.default_startup_program())
4932 4933
            batch_size = 1
            data_loader.start()
H
hutuxian 已提交
4934
            exe.train_from_dataset(
4935
                    fluid.default_main_program())
4936
            data_loader.reset()
4937 4938
    """

4939
    def __init__(self, optimizer, num_microbatches=1, start_cpu_core_id=0):
4940 4941 4942 4943 4944
        self._device = 'cpu'
        if core.is_compiled_with_npu():
            self._device = "npu"
        elif core.is_compiled_with_cuda():
            self._device = "gpu"
J
Jiabin Yang 已提交
4945
        if framework._non_static_mode():
Z
zhongpu 已提交
4946
            raise Exception("In dygraph, don't support PipelineOptimizer.")
4947 4948 4949 4950 4951
        valid_optimizers = (
            Optimizer,
            paddle.optimizer.Optimizer,
            paddle.fluid.contrib.mixed_precision.decorator.OptimizerWithMixedPrecision,
        )
4952
        if not isinstance(optimizer, valid_optimizers):
4953 4954 4955 4956 4957 4958 4959
            raise ValueError(
                "The 'optimizer' parameter for "
                "PipelineOptimizer must be an instance of "
                "{}, but the given type is {}.".format(
                    valid_optimizers, type(optimizer)
                )
            )
H
hutuxian 已提交
4960
        self._optimizer = optimizer
4961 4962 4963 4964 4965 4966

        # 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

4967 4968 4969
        assert (
            num_microbatches >= 1
        ), "num_microbatches must be a positive value."
4970
        self._num_microbatches = num_microbatches
4971 4972 4973
        assert (
            start_cpu_core_id >= 0
        ), "start_cpu_core_id must be a non-negative integer."
H
hutuxian 已提交
4974
        self._start_cpu_core_id = start_cpu_core_id
4975 4976 4977 4978 4979 4980
        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()
4981
        self._param_device_map = None
4982 4983
        self._pipeline_pair = []
        self._pp_ring_map = dict()
4984 4985
        self.output_var_to_op = None
        self.input_var_to_op = None
4986 4987 4988 4989 4990 4991 4992 4993 4994 4995 4996 4997 4998 4999 5000

    # 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")
5001 5002 5003 5004 5005 5006 5007 5008 5009 5010 5011 5012 5013 5014
            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,
                },
            )
5015 5016 5017 5018
            offset += 1
        block._insert_op(
            op_idx + 1 + offset,
            type='c_allreduce_max'
5019 5020
            if op.type == "reduce_any"
            else 'c_allreduce_sum',
5021 5022 5023
            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={
5024
                'ring_id': self.global_ring_id,
5025
                self._op_role_key: self._op_role.Optimize,
5026 5027 5028
                'use_calc_stream': True,
            },
        )
5029 5030
        offset += 1
        if op.type == "reduce_any":
5031 5032 5033 5034 5035 5036 5037 5038 5039 5040 5041
            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,
                },
            )
5042
            offset += 1
5043
        return offset
H
hutuxian 已提交
5044

5045
    def _create_vars(self, block, ori_block):
5046
        # Create vars for block, copied from ori_block
H
hutuxian 已提交
5047
        used_var_set = set()
5048 5049 5050 5051 5052 5053 5054 5055 5056
        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]
5057
            # For op process vars on all devices, remove its input
5058 5059 5060 5061 5062 5063 5064 5065 5066 5067 5068 5069 5070 5071 5072
            # 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)
5073 5074 5075 5076 5077 5078 5079 5080 5081 5082
            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
5083 5084 5085 5086 5087 5088 5089 5090
            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 已提交
5091
            for var in vars:
5092 5093
                # a var whose name contains "blocking_queue"
                # only exists in startup program
5094
                if var in used_var_set or "_blocking_queue" in var:
H
hutuxian 已提交
5095 5096
                    continue
                used_var_set.add(var)
5097 5098
                if block._find_var_recursive(str(var)):
                    continue
5099
                source_var = ori_block._var_recursive(str(var))
5100
                if source_var.type == core.VarDesc.VarType.READER:
5101
                    dest_var = block.create_var(
5102 5103
                        name=var,
                        type=core.VarDesc.VarType.READER,
5104 5105
                        persistable=source_var.persistable,
                    )
5106 5107 5108 5109 5110 5111 5112 5113 5114 5115 5116
                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,
5117 5118
                        error_clip=source_var.error_clip,
                    )
5119
                else:
5120
                    dest_var = block._clone_variable(source_var, False)
5121
                self._clone_var_attr(dest_var, source_var)
5122 5123 5124
            # When use with sharding, allreduce_sum and allreduce_max
            # used for global gradient clip and amp will be added by sharding.
            op_idx += 1
5125 5126
            if self.use_sharding or not should_insert:
                continue
5127 5128 5129 5130
            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 已提交
5131

5132
    def _is_loss_grad_op(self, op):
5133 5134
        assert self._op_role_key in op.attr_names
        op_role = int(op.attr(self._op_role_key))
5135
        return op_role & int(self._op_role.Backward) and op_role & int(
5136 5137
            self._op_role.Loss
        )
5138

5139
    def _is_forward_op(self, op):
5140 5141 5142
        return self._op_role_key in op.attr_names and (
            int(op.attr(self._op_role_key)) == int(self._op_role.Forward)
        )
5143

5144
    def _is_backward_op(self, op):
5145
        return self._op_role_key in op.attr_names and (
5146 5147
            int(op.attr(self._op_role_key)) & int(self._op_role.Backward)
        )
5148 5149 5150 5151

    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)
5152 5153

    def _is_optimize_op(self, op):
5154
        return self._op_role_key in op.attr_names and (
5155 5156
            int(op.attr(self._op_role_key)) & int(self._op_role.Optimize)
        )
5157 5158

    def _is_update_op(self, op):
5159 5160 5161 5162 5163
        return (
            'Param' in op.input_names
            and 'Grad' in op.input_names
            and ("LearningRate" in op.input_names)
        )
5164

5165
    def _split_program(self, main_program, devices):
H
hutuxian 已提交
5166
        """
5167
        Split a program into sections according to devices that ops run on.
5168
        The op whose op_device attr is "gpu:all" is copied to all sections.
5169 5170 5171

        Args:
            main_program (Program): the main program
5172
            devices: all used devices
H
hutuxian 已提交
5173
        """
5174
        # Map from device to its corresponding section program info
5175
        device_program_map = defaultdict(Program)
5176

5177
        block = main_program.block(0)
5178 5179
        for op in block.ops:
            device = op.attr(self._op_device_key)
5180
            # Copy ops whose op_device set to "gpu:all" to all sections.
5181
            if device == f"{self._device}:all":
5182
                for device in devices:
5183 5184
                    program = device_program_map[device]
                    op_desc = op.desc
5185
                    ap_op = program.global_block().desc.append_op()
5186
                    ap_op.copy_from(op_desc)
5187
                    ap_op._set_attr(self._op_device_key, "")
5188 5189 5190
            else:
                program = device_program_map[device]
                op_desc = op.desc
5191
                ap_op = program.global_block().desc.append_op()
5192
                ap_op.copy_from(op_desc)
5193
                ap_op._set_attr(self._op_device_key, "")
5194

5195
        program_list = []
5196
        for key in devices:
5197
            program = device_program_map[key]
5198 5199
            program._sync_with_cpp()
            program_list.append(program)
H
hutuxian 已提交
5200

5201
        return program_list
H
hutuxian 已提交
5202

5203 5204 5205 5206 5207 5208 5209
    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.
        """
5210 5211
        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 '
5212
            'or beta2_pow_acc.'
5213 5214
        )
        param_name = var_name[0 : var_name.index('_beta')]
5215 5216 5217
        device = self._param_device_map[param_name]
        return device

5218 5219
    def _split_startup_program(self, startup_program, device_id):
        block = startup_program.global_block()
5220 5221 5222
        new_startup_program = Program()
        for op in block.ops:
            device = op.attr(self._op_device_key)
5223 5224
            if device == "cpu":
                assert op.type == "fill_constant", (
5225
                    "For ops in startup program with the op_device attribute "
5226 5227
                    "of cpu, they must be of type fill_constant."
                )
5228 5229 5230
                output_var = op.output_arg_names[0]
                device = self._get_op_device_for_startup_program(output_var)

5231
            if device:
5232
                device_index = int(device.split(':')[1])
5233
            else:
5234 5235
                # LR related ops
                device = None
5236 5237
            if device and device_index != device_id:
                continue
5238
            op_desc = op.desc
5239
            ap_op = new_startup_program.global_block().desc.append_op()
5240 5241 5242
            ap_op.copy_from(op_desc)
            ap_op._set_attr(self._op_device_key, "")
        new_startup_program._sync_with_cpp()
5243
        self._create_vars(new_startup_program.global_block(), block)
5244 5245
        return new_startup_program

5246
    def _find_post_op(self, index, var_name):
H
hutuxian 已提交
5247
        """
5248
        Find the post op that has variable named var_name as input.
H
hutuxian 已提交
5249
        """
5250 5251 5252 5253 5254 5255
        # 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', '')

5256
        post_ops = self.input_var_to_op[var_name]
5257
        if post_ops is None:
5258
            return None
5259 5260 5261 5262 5263 5264
        result_op = None
        for post_op, post_idx in reversed(post_ops):
            if post_idx > index:
                result_op = post_op
                break
        return result_op
5265

5266
    def _find_prev_op(self, index, var_name):
H
hutuxian 已提交
5267
        """
5268 5269
        Find the previous op of op with index that outputs
        variable named var_name.
H
hutuxian 已提交
5270
        """
5271
        prev_ops = self.output_var_to_op[var_name]
5272
        if prev_ops is None:
5273
            return None
5274 5275 5276 5277
        result_op = None
        for prev_op, prev_idx in reversed(prev_ops):
            if prev_idx < index:
                result_op = prev_op
5278
                break
5279
        return result_op
5280 5281

    def _rename_arg(self, op, old_name, new_name):
5282 5283
        op._rename_input(old_name, new_name)
        op._rename_output(old_name, new_name)
5284

5285
    def _create_var(self, block, ref_var, name, dtype=None):
5286 5287 5288 5289 5290 5291 5292 5293
        """
        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,
5294
            dtype=ref_var.dtype if dtype is None else dtype,
5295 5296
            type=ref_var.type,
            lod_level=ref_var.lod_level,
5297 5298
            persistable=ref_var.persistable,
            is_data=ref_var.is_data,
5299 5300
            need_check_feed=ref_var.desc.need_check_feed(),
        )
5301
        self._clone_var_attr(new_var, ref_var)
5302 5303
        return new_var

5304 5305 5306 5307 5308
    def _clone_var_attr(self, dest, src):
        dest.stop_gradient = src.stop_gradient
        if hasattr(src, 'is_distributed'):
            dest.is_distributed = src.is_distributed

5309 5310 5311 5312 5313 5314
    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 已提交
5315

5316 5317 5318 5319 5320 5321
    def _append_grad_suffix(self, name):
        """
        Append grad suffix to the given variable name
        """
        return name + core.grad_var_suffix()

5322
    def _get_op_device_attr(self, op):
H
hutuxian 已提交
5323
        """
5324
        Get the op_device attribute of a op.
H
hutuxian 已提交
5325
        """
5326 5327 5328 5329 5330
        device = (
            op.attr(self._op_device_key)
            if op.has_attr(self._op_device_key)
            else None
        )
5331
        if device:
5332 5333
            assert device[0:3] == 'gpu' or device[0:3] == 'npu', (
                "Now, only gpu and npu devices are "
5334
                "supported in pipeline parallemism."
5335
            )
5336 5337 5338 5339 5340 5341 5342 5343 5344 5345 5346 5347 5348
        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
5349
            op._set_attr(self._op_device_key, f"{self._device}:all")
5350 5351 5352 5353
        # 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():
5354 5355 5356
                assert (
                    '@RENAME@' in name
                ), "The op must be sum used to accumulate renamed vars."
5357 5358 5359 5360
            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(
5361 5362 5363 5364
                'op_device'
            ), "{} has no op_device attr for var {}".format(
                post_op.type, out_name
            )
5365 5366 5367
            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)
5368 5369 5370
        elif (op.type == "cast" or op.type == "scale") and self._is_backward_op(
            op
        ):
5371
            prev_op = self._find_prev_op(idx, op.desc.input("X")[0])
5372 5373
            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):
5374
            # for checkpoint offloading
5375 5376 5377
            assert (
                len(op.input_arg_names) == 1 and len(op.output_arg_names) == 1
            )
5378 5379 5380
            input_name = op.input_arg_names[0]
            output_name = op.output_arg_names[0]
            if '@Fetch' in output_name:
5381
                post_op = self._find_post_op(idx, output_name)
5382 5383 5384
                op._set_attr(
                    self._op_device_key, post_op.attr(self._op_device_key)
                )
5385
            else:
5386
                prev_op = self._find_prev_op(idx, op.desc.input("X")[0])
5387 5388 5389
                op._set_attr(
                    self._op_device_key, prev_op.attr(self._op_device_key)
                )
5390 5391 5392
        elif self._is_loss_op(op):
            # For loss * loss_scaling op added by AMP
            offset = 1
5393 5394 5395
            while not block.ops[idx + offset].has_attr(
                self._op_device_key
            ) or not block.ops[idx + offset].attr(self._op_device_key):
5396 5397 5398 5399 5400 5401 5402 5403 5404
                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
5405
            param_name = self._strip_grad_suffix(grad_name[0])
5406 5407 5408 5409 5410
            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.
5411 5412
            assert self._op_role_var_key in op.attr_names, (
                "gradient_clip "
5413
                "and regularization ops must have op_role_var attribute."
5414
            )
5415
            op_role_var = op.attr(self._op_role_var_key)
5416 5417
            assert len(op_role_var) == 2, (
                "op_role_var for gradient_clip "
5418
                "regularization ops must have two elements."
5419
            )
5420 5421
            param_name = op_role_var[0]
            device = self._param_device_map[param_name]
5422
            # For sum op added by global gradient clip, it must be
5423
            # put on all devices
5424 5425 5426 5427 5428 5429 5430
            if (
                op.type == 'sum'
                or op.type == 'sqrt'
                or op.type == 'fill_constant'
                or op.type == 'elementwise_max'
                or op.type == 'elementwise_div'
            ):
5431
                device = f"{self._device}:all"
5432
            op._set_attr(self._op_device_key, device)
R
Roc 已提交
5433
        elif op.type == "alloc_float_status" or op.type == "clear_float_status":
5434
            op._set_attr(self._op_device_key, f"{self._device}:all")
5435 5436 5437 5438 5439 5440 5441 5442 5443 5444
            # 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
5445 5446
        else:
            other_known_ops = [
5447 5448 5449 5450 5451 5452
                'update_loss_scaling',
                'reduce_any',
                'concat',
                'sum',
                'check_finite_and_unscale',
                'memcpy',
5453
            ]
5454 5455 5456
            assert op.type in other_known_ops, (
                "For other ops without "
                "op_device set, they must be one of {}, but it "
5457
                "is {}".format(other_known_ops, op.type)
5458
            )
5459
            assert self._is_optimize_op(op)
5460
            op._set_attr(self._op_device_key, f"{self._device}:all")
5461 5462

    def _add_op_device_attr(self, block):
5463
        """
5464
        Add op_device attrribute for ops in block that have
5465
        not that attribute set.
5466
        """
5467
        for idx, op in enumerate(list(block.ops)):
5468 5469 5470 5471 5472
            if (
                op.type == "create_py_reader"
                or op.type == "read"
                or op.type == "create_double_buffer_reader"
            ):
5473
                # Copy read related ops to all section to make them exit
5474 5475 5476 5477
                # 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.
5478
                op._set_attr(self._op_device_key, f"{self._device}:all")
5479 5480
                continue
            # op_device attribute has been set
5481 5482
            if self._get_op_device_attr(op):
                continue
5483
            self._add_op_device_attr_for_op(op, idx, block)
H
hutuxian 已提交
5484

5485 5486
    def _check_validation(self, block):
        """
5487
        Check whether ops in a block have both the op_device and the
5488 5489
        op_role attributes set.
        Then, return all devices in order.
5490
        """
5491 5492 5493 5494 5495 5496 5497 5498 5499 5500
        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),
        ]
5501
        for op in block.ops:
5502
            if not op._has_kernel(op.type):
5503 5504 5505 5506 5507 5508
                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."
                )
5509
            assert op.has_attr(
5510 5511
                self._op_role_key
            ), "op ({}) has no {} attribute.".format(op.type, self._op_role_key)
5512
            op_role = op.attr(self._op_role_key)
5513 5514 5515 5516 5517
            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
            )
5518

5519
            assert op.has_attr(
5520 5521 5522 5523
                self._op_device_key
            ), "op ({}) has no {} attribute.".format(
                op.type, self._op_device_key
            )
5524 5525

            device = op.attr(self._op_device_key)
5526 5527 5528 5529 5530 5531 5532
            assert (
                device
            ), "op_device attribute for op " "{} has not been set.".format(
                op.type
            )
            if device == f"{self._device}:all":
                continue
5533

5534
            dev_type = device.split(':')[0]
B
Baibaifan 已提交
5535 5536
            assert dev_type == "gpu" or dev_type == 'npu', (
                "Now only gpu and npu devices are supported "
5537 5538
                "for pipeline parallelism."
            )
5539 5540

            if device not in device_list:
5541
                device_list.append(device)
5542

5543
        return device_list
5544

5545
    def _insert_sendrecv_ops_for_boundaries(self, block):
5546
        """
5547
        Insert a pair of send and recv ops for every two
5548 5549
        consecutive ops on different devices.
        """
5550
        # A map from var to device where op takes it as input,
5551
        # avoiding multiple send and recv ops.
5552
        input_var_to_device = dict()
5553 5554 5555 5556 5557 5558 5559 5560
        # 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,
5561
            'first_optimize_index': first_optimize_index,
5562
        }
5563

5564
        for index, op in enumerate(list(block.ops)):
5565
            cur_device = op.attr(self._op_device_key)
5566 5567
            if cur_device == f"{self._device}:all":
                continue
5568 5569
            for var_name in op.input_arg_names:
                var = block.var(var_name)
5570
                # skip data var
5571 5572
                if var.is_data:
                    continue
5573
                prev_device = None
5574 5575 5576

                prev_op = self._find_prev_op(index, var_name)
                if prev_op is None:
5577 5578
                    if var_name not in self._param_device_map:
                        continue
5579
                    prev_device = self._param_device_map[var_name]
5580

5581
                if not prev_device:
5582 5583 5584
                    prev_device = (
                        prev_op.attr(self._op_device_key) if prev_op else None
                    )
5585

5586 5587
                if prev_device is None or prev_device == f"{self._device}:all":
                    continue
5588

5589 5590
                if prev_device == cur_device:
                    continue
5591

5592 5593 5594 5595 5596 5597 5598
                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] + ':'

5599 5600 5601 5602
                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)
5603 5604
                    assert is_forward or is_backward, (
                        'send/recv in pipeline should only be inserted in forward or backward,'
5605
                        'please check the op_role of op={}'.format(op)
5606
                    )
5607 5608

                    if is_forward:
5609 5610
                        assert prev_id < cur_id, (
                            "In forward, send/recv can only be passed forward, but now "
5611
                            "prev_stage={} great than cur_stage={}, please check op_device of op={}".format(
5612 5613 5614
                                prev_id, cur_id, op
                            )
                        )
5615
                    elif is_backward:
5616 5617
                        assert prev_id > cur_id, (
                            "In backward, send/recv can only be passed backward, but now "
5618
                            "prev_stage={} less than cur_stage={}, please check op_device of op={}".format(
5619 5620 5621
                                prev_id, cur_id, op
                            )
                        )
5622

5623 5624 5625 5626 5627 5628 5629 5630 5631 5632
                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(
5633 5634
                            (cur_dev, prev_dev)
                        )
5635 5636 5637 5638 5639
                        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(
5640 5641
                            (cur_dev, prev_dev)
                        )
5642 5643 5644 5645 5646 5647
                        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)
5648
                    var = block.vars[var_name]
5649 5650 5651
                    pair = (prev_id, cur_id)
                    # 1000 is just a magic number
                    pair_key = prev_id * 1000 + cur_id
5652 5653 5654 5655 5656 5657 5658
                    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]
5659

5660
                    if self.schedule_mode == 'F-then-B':  # F-then-B
F
fangshuixun007 已提交
5661
                        block._insert_op_without_sync(
5662
                            index=index + extra_index_info['index'],
5663 5664 5665
                            type='send_v2',
                            inputs={'X': var},
                            attrs={
5666
                                self._op_device_key: prev_dev,
5667 5668 5669
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 1,
5670 5671 5672
                                'ring_id': ring_id,
                            },
                        )
5673
                        extra_index_info['index'] += 1
5674
                        var_shape = list(var.shape)
5675 5676 5677 5678 5679
                        var_shape[0] = (
                            self.micro_batch_size
                            if var_shape[0] < 0
                            else var_shape[0]
                        )
F
fangshuixun007 已提交
5680
                        block._insert_op_without_sync(
5681
                            index=index + extra_index_info['index'],
5682 5683 5684
                            type='recv_v2',
                            outputs={'Out': [var]},
                            attrs={
5685
                                'out_shape': var_shape,
5686
                                'dtype': var.dtype,
5687
                                self._op_device_key: cur_dev,
5688 5689 5690
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 0,
5691 5692 5693
                                'ring_id': ring_id,
                            },
                        )
5694
                        extra_index_info['index'] += 1
5695
                    elif self.schedule_mode == '1F1B':  # 1F1B
5696
                        var_shape = list(var.shape)
5697 5698 5699 5700 5701
                        var_shape[0] = (
                            self.micro_batch_size
                            if var_shape[0] < 0
                            else var_shape[0]
                        )
5702

5703
                        numel = np.prod(var_shape)
5704 5705 5706
                        use_mp = (self.mp_degree > 1) and (
                            numel % self.mp_degree == 0
                        )
5707 5708 5709 5710 5711 5712 5713 5714 5715 5716 5717 5718 5719 5720 5721 5722 5723 5724 5725 5726 5727 5728

                        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,
5729 5730
                                },
                            )
5731 5732 5733
                            extra_index_info['index'] += 1
                            return

5734 5735
                        _check_stage(cur_id, prev_id)

5736 5737 5738 5739 5740 5741 5742 5743 5744 5745
                        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,
                            },
                        )
5746
                        extra_index_info['index'] += 1
5747 5748
                        prefix_name = var.name.split('@')[0]
                        prefix_var = block.var(prefix_name)
5749 5750 5751
                        is_param = (
                            True if isinstance(prefix_var, Parameter) else False
                        )
F
fangshuixun007 已提交
5752
                        block._insert_op_without_sync(
5753
                            index=index + extra_index_info['index'],
5754
                            type='send_v2'
5755 5756
                            if not use_mp or is_param
                            else 'partial_send',
5757 5758
                            inputs={'X': var},
                            attrs={
5759
                                self._op_device_key: prev_dev,
5760 5761 5762 5763
                                self._op_role_key: op_role,
                                'use_calc_stream': False,
                                'ring_id': ring_id,
                                'peer': 1,
5764 5765 5766
                                # if send_v2, num&id attr is not in op_attrs, will not insert
                                'num': self.mp_degree,
                                'id': self.mp_rank,
5767 5768
                            },
                        )
5769
                        extra_index_info['index'] += 1
5770 5771 5772
                        insert_index = None
                        if int(op_role) == int(self._op_role.Backward):
                            insert_index = extra_index_info[
5773 5774
                                'first_optimize_index'
                            ]
5775 5776 5777 5778
                            new_op_role = self._op_role.Optimize
                        else:
                            insert_index = index
                            new_op_role = self._op_role.Backward
5779
                        sync_comm_op = block._insert_op_without_sync(
5780
                            index=insert_index + extra_index_info['index'],
5781 5782 5783 5784
                            type='c_sync_comm_stream',
                            inputs={'X': [var]},
                            outputs={'Out': [var]},
                            attrs={
5785
                                self._op_device_key: prev_dev,
5786
                                self._op_role_key: new_op_role,
5787
                                'ring_id': ring_id,
5788 5789
                            },
                        )
5790
                        if int(op_role) == int(self._op_role.Forward):
5791
                            sync_comm_op._set_attr('pipeline_flag', '')
5792
                            extra_index_info['index'] += 1
F
fangshuixun007 已提交
5793
                        block._insert_op_without_sync(
5794
                            index=index + extra_index_info['index'],
5795
                            type='recv_v2'
5796 5797
                            if not use_mp or is_param
                            else 'partial_recv',
5798 5799 5800 5801
                            outputs={'Out': [var]},
                            attrs={
                                'out_shape': var_shape,
                                'dtype': var.dtype,
5802
                                self._op_device_key: cur_dev,
5803 5804 5805
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 0,
5806 5807 5808 5809
                                '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,
5810 5811
                            },
                        )
5812
                        extra_index_info['index'] += 1
5813
                        if use_mp and not is_param:
5814 5815 5816 5817 5818 5819 5820 5821 5822 5823 5824 5825 5826
                            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,
5827 5828
                                },
                            )
5829
                            extra_index_info['index'] += 1
5830 5831 5832
                    else:
                        raise ValueError(
                            "Now only 'F-then-B' and '1F1B' are supported."
5833 5834
                            "The given value is {}.".format(self.schedule_mode)
                        )
5835

5836 5837 5838 5839
                _insert_send_recv(
                    int(cur_device.split(':')[1]),
                    int(prev_device.split(':')[1]),
                )
5840 5841
        block._sync_with_cpp()

5842
    def _insert_loss_scale(self, block):
5843
        """
5844
        Scale the loss corresponding to number of micro-batches.
5845
        """
5846 5847
        if self._num_microbatches == 1:
            return
5848
        for index, op in reversed(tuple(enumerate(list(block.ops)))):
5849
            if self._is_loss_grad_op(op):
5850 5851
                assert op.type == 'fill_constant', (
                    "loss_grad_op must be fill_constant op, "
5852
                    "but this op is {}".format(op.type)
5853
                )
5854 5855 5856 5857
                assert op.has_attr('value')
                loss_scale = float(op.attr('value'))
                loss_scale = loss_scale / self._num_microbatches
                op._set_attr('value', loss_scale)
5858 5859
                break

5860 5861
    def _rename_gradient_var_name(self, block):
        for index, op in enumerate(block.ops):
5862 5863
            if not self._is_optimize_op(op):
                continue
5864 5865 5866
            input_names = op.input_arg_names
            output_names = op.output_arg_names
            in_out_names = input_names + output_names
5867 5868
            if op.type == 'cast' or op.type == "c_sync_comm_stream":
                continue
5869 5870 5871
            # 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:
5872 5873
                if not core.grad_var_suffix() in name:
                    continue
5874 5875 5876 5877
                param_name = name.strip(core.grad_var_suffix())
                new_grad_name = name + "@MERGED"
                self._rename_arg(op, name, new_grad_name)

5878 5879 5880
    def _accumulate_gradients(
        self, block, pp_allreduce_in_optimize=False, strategy=None, shard=None
    ):
5881 5882 5883 5884
        """
        Create a new merged gradient for each parameter and accumulate the
        corresponding gradient to it.
        """
5885 5886
        fp16_allreduce = strategy.fp16_allreduce if strategy else False
        if strategy and strategy.fuse_grad_merge:
5887
            fused_gradient_names = self._accumulate_gradients_with_fuse(
5888 5889
                block, fp16_allreduce, strategy.fuse_grad_size_in_MB, shard
            )
5890 5891
            return fused_gradient_names

5892 5893 5894
        merged_gradient_names = []
        first_opt_op_idx = None

5895 5896 5897
        merged_suffix = '@MERGED@FP16' if fp16_allreduce else '@MERGED'
        dtype = paddle.float16 if fp16_allreduce else None

5898 5899 5900 5901 5902 5903 5904 5905
        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)
5906
                    continue
5907

5908
            if self._is_backward_op(op) and first_opt_op_idx is None:
5909
                first_opt_op_idx = index + 1
5910 5911
                # maybe have no optimize
                # if first_opt_op_idx == len(block.ops): return
5912

5913 5914 5915
            if self._is_backward_op(op) and (
                self._op_role_var_key in op.attr_names
            ):
5916
                op_role_var = op.attr(self._op_role_var_key)
5917 5918
                if len(op_role_var) == 0:
                    continue
5919 5920
                assert len(op_role_var) % 2 == 0
                for i in range(0, len(op_role_var), 2):
5921 5922
                    offset = 0
                    param_name = op_role_var[i]
5923 5924 5925 5926
                    if not block.has_var(param_name):
                        continue
                    if '@BroadCast' in param_name:
                        continue
5927

5928
                    param_grad_name = param_name + core.grad_var_suffix()
5929
                    merged_param_grad_name = param_grad_name + merged_suffix
5930
                    if not block.has_var(merged_param_grad_name):
5931 5932 5933 5934 5935 5936
                        self._create_var(
                            block,
                            block.vars[param_name],
                            merged_param_grad_name,
                            dtype,
                        )
5937
                    assert block.has_var(merged_param_grad_name)
5938

5939 5940 5941
                    param_grad_var = block.var(param_grad_name)
                    merged_param_grad_var = block.var(merged_param_grad_name)
                    merged_param_grad_var.persistable = True
5942
                    block._insert_op(
5943 5944 5945 5946
                        index=first_opt_op_idx + offset,
                        type='fill_constant',
                        inputs={},
                        outputs={'Out': [merged_param_grad_var]},
5947
                        attrs={
5948 5949 5950
                            'shape': merged_param_grad_var.shape,
                            'dtype': merged_param_grad_var.dtype,
                            'value': float(0),
5951
                            # a trick to run this op once per mini-batch
5952 5953 5954
                            self._op_role_key: self._op_role.Optimize.LRSched,
                        },
                    )
5955
                    offset += 1
5956 5957
                    grad_name = op_role_var[i + 1]
                    grad_var = block.vars[grad_name]
5958 5959

                    is_fp16_grad = 'cast_fp16' in grad_name
5960
                    need_cast = is_fp16_grad is not fp16_allreduce
5961 5962 5963 5964 5965 5966

                    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
5967
                        cast_grad_var_name = param_grad_name + '@TMP'
5968
                        cast_grad_var = self._create_var(
5969 5970
                            block, param_grad_var, cast_grad_var_name, dtype
                        )
5971
                        cast_grad_var.persistable = False
5972 5973 5974 5975 5976 5977 5978 5979 5980 5981 5982
                        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,
                            },
                        )
5983
                        offset += 1
5984 5985 5986 5987 5988 5989 5990
                        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},
5991 5992
                        attrs={
                            self._op_role_key: self._op_role.Backward,
5993 5994
                        },
                    )
5995 5996 5997
                    offset += 1
                    merged_gradient_names.append(merged_param_grad_name)

5998 5999
        if not fp16_allreduce:
            return merged_gradient_names
6000 6001 6002 6003 6004 6005 6006 6007 6008 6009 6010 6011 6012 6013 6014 6015 6016 6017 6018 6019 6020 6021 6022

        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

6023 6024 6025 6026 6027 6028 6029 6030 6031 6032 6033
            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,
                },
            )
6034

6035
        return merged_gradient_names
6036

6037 6038 6039
    def _insert_accumulate_gradients_with_fuse(
        self, main_block, fp16, fused_size, grad_param_pairs, first_opt_op_idx
    ):
6040
        grad_param_pairs = self._sort_grad_param_by_dtype(
6041 6042
            main_block, grad_param_pairs
        )
6043

6044 6045 6046
        grad_param_segments = []
        merged_suffix = '@MERGED@FP16' if fp16 else '@MERGED'
        dtype = paddle.float16 if fp16 else paddle.float32
6047
        cur_size = 0.0
6048 6049 6050 6051 6052 6053 6054 6055 6056 6057
        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,
6058 6059
                stop_gradient=False,
            )
6060
            real_param = main_block.var(param)
6061 6062
            if hasattr(real_param, 'is_distributed'):
                merged_grad_var.is_distributed = real_param.is_distributed
6063 6064 6065 6066
            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
6067 6068 6069 6070 6071
            if (
                len(grad_param_segments) == 0
                or cur_size + tmp_size > fused_size
                or real_grad.dtype != last_dtype
            ):
6072
                grad_param_segments.append(
6073 6074
                    ([real_grad], [real_param], [merged_grad_var])
                )
6075
                last_dtype = real_grad.dtype
6076
                cur_size = 0.0
6077 6078 6079 6080 6081 6082 6083 6084 6085 6086 6087 6088
            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]
6089 6090 6091 6092 6093 6094
            fused_grad = main_block.create_var(
                name='FusedGrad_{}'.format(grad_segment[0].name),
                dtype=grad_segment[0].dtype,
                persistable=False,
                stop_gradient=False,
            )
6095
            # keep the '.cast_fp16' info in the fuse var name
6096 6097 6098 6099 6100 6101 6102 6103 6104
            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)
            )
6105 6106 6107 6108
            fused_merged_grad = main_block.create_var(
                name=fused_merged_grad_name,
                dtype=merged_grad_segment[0].dtype,
                persistable=True,
6109 6110
                stop_gradient=False,
            )
6111 6112 6113 6114 6115 6116 6117 6118 6119 6120 6121 6122 6123 6124 6125 6126 6127 6128 6129 6130 6131 6132 6133 6134 6135
            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},
6136
                outputs={"Output": grads, "FusedOutput": fused_grad},
6137 6138 6139 6140 6141 6142 6143 6144 6145 6146 6147 6148 6149 6150 6151 6152
                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,
6153 6154 6155 6156 6157 6158 6159
                    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),
6160 6161
                },
            )
6162 6163 6164 6165 6166 6167 6168 6169 6170 6171
            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,
6172
                    "FusedOutput": fused_merged_grad,
6173 6174 6175 6176 6177 6178 6179 6180
                },
                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,
6181 6182 6183
                    self._op_role_key: self._op_role.Optimize.LRSched,
                },
            )
6184 6185 6186 6187 6188 6189 6190 6191 6192
            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
6193
            need_cast = is_fp16_grad is not fp16
6194 6195 6196 6197
            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'
6198 6199 6200 6201 6202 6203 6204 6205 6206 6207 6208 6209 6210 6211 6212 6213 6214
                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,
                    },
                )
6215 6216 6217 6218 6219 6220 6221
                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},
6222 6223
                attrs={self._op_role_key: self._op_role.Backward},
            )
6224 6225 6226 6227 6228 6229 6230 6231 6232 6233
            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'
6234 6235 6236 6237 6238 6239 6240 6241 6242 6243 6244 6245 6246 6247 6248 6249 6250 6251
                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,
                    },
                )
6252 6253 6254 6255 6256 6257
                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

6258
        return fused_merged_gradients, first_opt_op_idx
6259

6260 6261 6262
    def _accumulate_gradients_with_fuse(
        self, main_block, fp16, fused_size, shard=None
    ):
6263 6264 6265 6266 6267 6268 6269 6270 6271 6272 6273 6274 6275 6276 6277 6278 6279 6280 6281
        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

6282 6283 6284
            if self._is_backward_op(op) and (
                self._op_role_var_key in op.attr_names
            ):
6285 6286 6287 6288 6289 6290 6291 6292 6293 6294 6295
                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(
6296 6297
                        (op_role_var[i + 1], op_role_var[i])
                    )
6298 6299 6300 6301 6302 6303 6304 6305 6306 6307 6308 6309 6310

        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:
6311 6312 6313 6314 6315 6316
            (
                fused_merged_gradients,
                first_opt_op_idx,
            ) = self._insert_accumulate_gradients_with_fuse(
                main_block, fp16, fused_size, pairs, first_opt_op_idx
            )
6317 6318 6319 6320
            all_fused_merged_gradients += fused_merged_gradients

        main_block._sync_with_cpp()
        return all_fused_merged_gradients
6321

6322 6323 6324 6325 6326 6327 6328 6329 6330 6331 6332 6333 6334 6335 6336 6337 6338 6339
    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

6340 6341 6342 6343 6344 6345 6346 6347 6348 6349 6350 6351
    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
6352 6353 6354 6355 6356 6357
        return (
            reduce(lambda x, y: x * y, var.shape)
            * dtype_to_size[var.dtype]
            / 1024.0
            / 1024.0
        )
6358

6359 6360
    def _add_sub_blocks(self, main_block, program_list):
        main_program = main_block.program
6361
        for prog in program_list:
6362 6363 6364 6365 6366 6367
            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)
6368 6369
                for sub_op in origin_sub_block.ops:
                    op_desc = sub_op.desc
6370 6371 6372
                    ap_op = new_sub_block.desc.append_op()
                    ap_op.copy_from(op_desc)
                new_sub_block._sync_with_cpp()
6373
                self._create_vars(new_sub_block, origin_sub_block)
6374
                op._set_attr('sub_block', new_sub_block)
6375 6376 6377

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

6383 6384 6385
    def _process_persistable_vars_in_multi_sections(
        self, main_program, startup_prog, program_list
    ):
6386 6387 6388 6389 6390 6391 6392
        """
        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()
6393
        for prog in program_list:
6394 6395
            block = prog.block(0)
            for var_name in block.vars:
6396 6397
                if var_name == "double_buffer_0":
                    continue
6398
                var = block.var(var_name)
6399 6400
                if not var.persistable:
                    continue
6401 6402 6403 6404 6405 6406 6407 6408 6409 6410 6411 6412 6413 6414 6415
                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:
6416 6417 6418 6419 6420 6421
                    if (
                        op.type == "recv_v2"
                        or op.type == "create_py_reader"
                        or op.type == "read"
                        or op.type == "update_loss_scaling"
                    ):
6422
                        continue
6423 6424
                    # We have processed lr related vars
                    if op.attr(self._op_role_key) == int(
6425 6426
                        self._op_role.Optimize.LRSched
                    ):
6427 6428 6429 6430
                        continue
                    if var_name in op.desc.output_arg_names():
                        assert var_name not in write_info, (
                            "two sections write the same var({}): second "
6431 6432
                            "op {}.".format(var_name, op)
                        )
6433 6434 6435 6436 6437
                        write_info[var_name] = prog
                        break

        for var_name in var_info.keys():
            # Case 1: read only variables, no special process
6438 6439
            if not var_name in write_info:
                continue
6440 6441 6442 6443 6444

            # 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)
6445
            write_dev_index = int(write_device.split(':')[1])
6446 6447
            all_progs = var_info[var_name]
            for prog in all_progs:
6448 6449
                if prog == write_prog:
                    continue
6450 6451 6452
                read_block = prog.block(0)
                read_device = self._get_device_info(read_block)
                read_dev_index = int(read_device.split(':')[1])
6453 6454 6455 6456 6457 6458 6459 6460 6461
                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]
6462 6463 6464

                write_block._insert_op(
                    index=0,
6465
                    type='send_v2',
6466 6467 6468
                    inputs={
                        'X': write_block.var(var_name),
                    },
6469
                    attrs={
6470 6471
                        self._op_device_key: write_device,
                        'use_calc_stream': False,
6472 6473
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
6474 6475 6476 6477 6478
                        self._op_role_key: self._op_role.LRSched,
                        'peer': read_dev_index,
                        'ring_id': ring_id,
                    },
                )
6479 6480
                read_block._insert_op(
                    index=0,
6481
                    type='recv_v2',
6482 6483
                    outputs={'Out': [read_block.var(var_name)]},
                    attrs={
6484 6485 6486 6487
                        '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,
6488 6489
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
6490 6491 6492 6493 6494
                        self._op_role_key: self._op_role.LRSched,
                        'peer': write_dev_index,
                        'ring_id': ring_id,
                    },
                )
6495 6496 6497 6498 6499 6500
                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={
6501
                        self._op_device_key: read_device,
6502 6503
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
6504 6505 6506 6507
                        self._op_role_key: self._op_role.LRSched,
                        'ring_id': ring_id,
                    },
                )
6508 6509

    def _is_gradient_clip_op(self, op):
6510 6511 6512
        return op.desc.has_attr("op_namescope") and op.desc.attr(
            "op_namescope"
        ).startswith("/gradient_clip")
6513 6514

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

6519 6520
    def _is_weight_decay_op(self, op):
        # in AdamW namescope is /optimizer_*/weight decay/
6521 6522 6523
        return op.desc.has_attr(
            "op_namescope"
        ) and 'weight decay' in op.desc.attr("op_namescope")
6524

6525 6526 6527 6528 6529
    def _get_input_output_info(self, block):
        '''
        Get info of op input and output.
        '''
        # A map from output var to op which generate it.
6530
        output_var_to_op = defaultdict(list)
6531
        # A map from var to op which takes it as input.
6532
        input_var_to_op = defaultdict(list)
6533

6534
        for index, op in enumerate(block.ops):
6535
            for var_name in op.input_arg_names:
6536
                input_var_to_op[var_name].append([op, index])
6537
            for var_name in op.output_arg_names:
6538 6539 6540 6541 6542 6543 6544 6545
                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
        """
6546 6547
        if self.schedule_mode != '1F1B':
            return
6548 6549 6550

        block = program.block(0)

6551
        recv_type = 'recv_v2' if self.mp_degree == 1 else 'partial_recv'
6552 6553
        backward_recv_index = None
        for index, op in enumerate(block.ops):
6554
            if op.type == recv_type and self._is_backward_op(op):
6555 6556 6557
                backward_recv_index = index
                break

6558
        # last pipeline stage
6559 6560
        if backward_recv_index is None:
            return
6561 6562 6563

        offset = 0
        for index, op in enumerate(list(block.ops)):
6564 6565
            if index >= backward_recv_index:
                break
6566 6567 6568 6569 6570 6571 6572 6573 6574 6575 6576 6577 6578 6579 6580 6581
            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]},
6582 6583
                    attrs={self._op_role_key: self._op_role.Backward},
                )
6584
        block._sync_with_cpp()
6585

6586 6587 6588 6589 6590 6591 6592 6593 6594 6595 6596 6597 6598
    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))
6599 6600 6601 6602
            if (
                op_role == int(self._op_role.Backward)
                and backward_insert_index is None
            ):
6603
                backward_insert_index = i
6604 6605 6606 6607 6608 6609
            if (
                op.type != "partial_recv"
                and op.type != "partial_allgather"
                and op.type != "nop"
                and op.type != "recv_v2"
            ):
6610 6611 6612 6613 6614 6615 6616 6617 6618 6619 6620 6621 6622 6623 6624 6625 6626 6627 6628
                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)
6629 6630 6631 6632 6633 6634 6635
            block._insert_op_without_sync(
                index=insert_index,
                type=op.type,
                inputs=op_inputs,
                outputs=op_outputs,
                attrs=op.all_attrs(),
            )
6636 6637 6638 6639 6640 6641 6642
            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()

6643 6644 6645 6646 6647 6648 6649 6650 6651 6652 6653 6654 6655 6656 6657 6658 6659 6660 6661 6662 6663 6664 6665 6666 6667 6668 6669
    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 "
6670 6671
            "forward and used in backward:\n{}".format(used_in_backward)
        )
6672

6673 6674 6675
    def minimize(
        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
    ):
6676
        main_block = loss.block
6677
        self.origin_main_block = main_block
6678
        main_program = main_block.program
6679 6680
        if startup_program is None:
            startup_program = default_startup_program()
6681

6682 6683
        pipeline_opt = main_program._pipeline_opt
        assert pipeline_opt, 'Please use pipeline with fleet.'
6684 6685 6686 6687 6688 6689 6690
        required_keys = [
            'local_rank',
            'schedule_mode',
            'micro_batch_size',
            'ring_id',
            'global_ring_id',
            'use_sharding',
6691 6692
            'mp_degree',
            'mp_rank',
6693 6694
        ]
        for key in required_keys:
6695 6696 6697
            assert (
                key in pipeline_opt
            ), 'Please use pipeline with fleet to use {}.'.format(key)
6698 6699 6700 6701 6702 6703 6704 6705
        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']
6706
        self.scale_gradient = pipeline_opt.get('scale_gradient', False)
6707 6708
        assert self.mp_degree >= 1
        assert 0 <= self.mp_rank < self.mp_degree
6709 6710

        optimize_ops, params_grads = self._optimizer.minimize(
6711 6712
            loss, startup_program, parameter_list, no_grad_set
        )
6713
        self._param_device_map = self._origin_optimizer._param_device_map
6714

6715 6716 6717 6718
        (
            self.output_var_to_op,
            self.input_var_to_op,
        ) = self._get_input_output_info(main_block)
6719 6720 6721
        # Step1: add default op_device attribute for ops.
        self._add_op_device_attr(main_block)
        device_list = self._check_validation(main_block)
6722 6723 6724 6725 6726 6727 6728 6729 6730 6731 6732

        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

6733 6734 6735
        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 "
6736 6737
            "another in the order of their ids."
        )
6738
        # Step2: add send and recv ops between section boundaries
6739
        self._insert_sendrecv_ops_for_boundaries(main_block)
6740

6741
        # Step3: split program into sections and add pairs of
6742 6743
        # send and recv ops for data var.
        main_program = main_block.program
6744
        program_list = self._split_program(main_program, device_list)
6745
        for p in program_list:
6746
            self._create_vars(p.global_block(), main_block)
6747

L
lilong12 已提交
6748 6749 6750 6751 6752
        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 "
6753 6754
                "stages."
            )
L
lilong12 已提交
6755 6756
        else:
            self.local_rank %= len(device_list)
6757 6758 6759
        # Step3.5: optimize forward send sync_comm to overlap send and recv
        self._optimize_forward_send_sync(program_list[self.local_rank])

6760
        # Step4: Special Case: process persistable vars that exist in
6761
        # multiple sections
6762
        # FIXME
6763 6764
        # self._process_persistable_vars_in_multi_sections(
        #     main_program, startup_program, program_list)
6765

6766
        # Step5: Add sub blocks for section programs
6767 6768
        self._add_sub_blocks(main_block, program_list)

6769
        place_list = []
6770 6771
        for dev in device_list:
            dev_index = int(dev.split(":")[1])
6772 6773 6774 6775
            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))
6776

6777
        # Step6: Split startup program
6778
        new_startup_program = self._split_startup_program(
6779 6780
            startup_program, self.local_rank
        )
6781 6782 6783 6784

        startup_program._pipeline_opt = {
            "startup_program": new_startup_program,
        }
6785
        real_block = program_list[self.local_rank].global_block()
6786 6787
        if not self.scale_gradient:
            self._insert_loss_scale(real_block)
6788
        if not self.use_sharding:
6789
            # Step7: clear gradients before each mini-batch and
6790 6791 6792 6793 6794
            # 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()
6795

6796 6797 6798 6799
        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"))
6800 6801 6802
        # A pass to move the recv op to the beginning of
        # the forward/backward phase
        self._mv_head_recv(program_list[self.local_rank])
6803 6804 6805 6806 6807

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

6808
        main_program._pipeline_opt = {
H
hutuxian 已提交
6809 6810
            "trainer": "PipelineTrainer",
            "device_worker": "Section",
6811
            "pipeline_stage": self.local_rank,
6812
            "num_pipeline_stages": len(device_list),
6813
            "schedule_mode": self.schedule_mode,
6814
            "inner_parallelism": len(device_list),
6815 6816
            "section_program": program_list[self.local_rank],
            "place": place_list[self.local_rank],
6817
            "place_id": place_id,
6818
            "sync_steps": -1,
L
lilong12 已提交
6819
            "num_microbatches": self._num_microbatches,
H
hutuxian 已提交
6820 6821
            "start_cpu_core_id": self._start_cpu_core_id,
        }
6822 6823 6824 6825 6826 6827 6828
        return (
            optimize_ops,
            params_grads,
            program_list,
            self._pipeline_pair,
            self._pp_ring_map,
        )
M
mapingshuo 已提交
6829 6830


M
mapingshuo 已提交
6831 6832
class RecomputeOptimizer(Optimizer):
    """
6833
        :api_attr: Static Graph
S
swtkiwi 已提交
6834

M
mapingshuo 已提交
6835 6836 6837
    Recompute Optimizer Wrapper

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

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

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

M
mapingshuo 已提交
6850 6851 6852 6853 6854 6855 6856 6857 6858 6859 6860 6861 6862 6863 6864 6865 6866 6867 6868 6869 6870 6871 6872 6873 6874 6875 6876 6877 6878 6879 6880 6881 6882 6883 6884 6885 6886 6887 6888 6889 6890 6891 6892 6893 6894
    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 已提交
6895
        if framework._non_static_mode():
Z
zhongpu 已提交
6896
            raise Exception("In dygraph, don't support RecomputeOptimizer.")
M
mapingshuo 已提交
6897 6898
        self._optimizer = optimizer
        self._checkpoints = None
M
mapingshuo 已提交
6899 6900
        self._learning_rate = self._optimizer._learning_rate
        self._learning_rate_map = self._optimizer._learning_rate_map
J
JZ-LIANG 已提交
6901
        self.enable_offload = False
M
mapingshuo 已提交
6902 6903

    def _set_checkpoints(self, checkpoints):
6904 6905
        """
        Args:
6906
            checkpoints (list): List of Variable or string
6907 6908 6909 6910 6911
        """
        assert isinstance(
            checkpoints, list
        ), "_checkpoints should be a list of Variable or a list of String"
        for ckpt in checkpoints:
6912 6913
            assert isinstance(ckpt, str) or isinstance(
                ckpt, Variable
6914
            ), "_checkpoints should be a list of Variable or a list of String"
M
mapingshuo 已提交
6915 6916
        self._checkpoints = checkpoints

6917
    # should enable offload before calling backward
J
JZ-LIANG 已提交
6918 6919 6920
    def _enable_offload(self):
        self.enable_offload = True

6921 6922
    @framework.deprecate_stat_dict
    def load(self, state_dict):
M
mapingshuo 已提交
6923
        """
6924
            :api_attr: Static Graph
S
swtkiwi 已提交
6925

M
mapingshuo 已提交
6926 6927 6928 6929
        load function is not supported by Recompute Optimizer for now.
        :return: None

        Args:
6930
            state_dict: the dict load by load_persistable method
M
mapingshuo 已提交
6931 6932 6933 6934

        Examples:
            .. code-block:: python

6935
                import paddle
M
mapingshuo 已提交
6936
                import paddle.fluid as fluid
6937

6938
                paddle.enable_static()
M
mapingshuo 已提交
6939 6940 6941 6942 6943 6944
                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
6945

M
mapingshuo 已提交
6946 6947 6948 6949
                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")
6950

M
mapingshuo 已提交
6951 6952 6953 6954
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
                sgd._set_checkpoints([fc_1, pred])
                try:
6955 6956
                    state_dict = {}
                    sgd.load(state_dict)
M
mapingshuo 已提交
6957
                except NotImplementedError as e:
6958
                    print(e)
M
mapingshuo 已提交
6959 6960
        """
        raise NotImplementedError(
6961 6962
            "load function is not supported by Recompute Optimizer for now"
        )
M
mapingshuo 已提交
6963 6964 6965 6966 6967 6968 6969 6970 6971 6972 6973 6974 6975 6976 6977 6978 6979 6980 6981 6982 6983 6984 6985 6986 6987 6988 6989 6990 6991 6992 6993 6994

    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)
6995
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
6996 6997 6998 6999
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
7000
                    no_grad_set=None)
M
mapingshuo 已提交
7001 7002 7003 7004 7005 7006 7007 7008 7009 7010

                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 已提交
7011 7012 7013 7014 7015 7016 7017 7018 7019
    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,
7020 7021
            stop_gradient=True,
        )
J
JZ-LIANG 已提交
7022 7023 7024 7025 7026 7027

        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,
7028 7029
            stop_gradient=False,
        )
J
JZ-LIANG 已提交
7030 7031 7032 7033 7034 7035 7036 7037

        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
7038 7039 7040
        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 已提交
7041 7042 7043 7044 7045 7046 7047 7048 7049 7050 7051 7052 7053
        """
        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,
7054 7055 7056 7057 7058 7059 7060 7061 7062 7063 7064 7065 7066
                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 已提交
7067 7068 7069

        return

7070 7071 7072
    def _insert_async_memcpy_op(
        self, insert_idx, src_varname, dst_varname, op_role, dst_place_type
    ):
J
JZ-LIANG 已提交
7073 7074 7075 7076 7077 7078 7079 7080
        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)]
            },
7081 7082
            attrs={"dst_place_type": int(dst_place_type), OP_ROLE_KEY: op_role},
        )
J
JZ-LIANG 已提交
7083 7084

    def _insert_fetch_op(self, idx, varname):
7085 7086 7087 7088 7089
        assert (
            varname in self.checkpoint_name2pinned_name
        ), "Try to fetch {} from Pinned Memory, but it is NOT a checkpoint".format(
            varname
        )
J
JZ-LIANG 已提交
7090 7091 7092

        pinned_varname = self.checkpoint_name2pinned_name[varname]
        fetch_varname = self.checkpoint_name2fetch_name[varname]
7093
        self._insert_async_memcpy_op(idx, pinned_varname, fetch_varname, 1, 1)
J
JZ-LIANG 已提交
7094 7095

    def _insert_offload_op(self, idx, varname):
7096 7097 7098 7099 7100
        assert (
            varname in self.checkpoint_name2pinned_name
        ), "Try to offload {} to Pinned Memory, but it is NOT a checkpoint".format(
            varname
        )
J
JZ-LIANG 已提交
7101
        pinned_varname = self.checkpoint_name2pinned_name[varname]
7102
        self._insert_async_memcpy_op(idx, varname, pinned_varname, 0, 2)
J
JZ-LIANG 已提交
7103 7104

    def _insert_sync_op(self, op_idx, checkpoint_name):
7105
        # single stream offload no need sync
J
JZ-LIANG 已提交
7106 7107 7108
        pass

    def _record_fetch_op(self, idx):
7109 7110 7111
        assert (
            len(self.un_fetch_checkpoint_names) > 0
        ), "Could NOT found checkpoint to fetch"
J
JZ-LIANG 已提交
7112 7113 7114 7115 7116 7117 7118 7119
        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)
7120 7121 7122 7123 7124
        assert (
            checkpoint_name == expected_checkpoint_name
        ), "expected to offload [{}] but got [{}]".format(
            expected_checkpoint_name, checkpoint_name
        )
J
JZ-LIANG 已提交
7125 7126 7127 7128
        logging.debug("Record offload [{}]".format(checkpoint_name))
        self.idx2insertions[idx] = ("offload", checkpoint_name)

    def _record_sync_op(self, idx, checkpoint_name):
7129 7130 7131
        assert (
            checkpoint_name not in self.synced_checkpoints
        ), "Try to sync the checkpoint [{}] twice".format(checkpoint_name)
J
JZ-LIANG 已提交
7132 7133 7134 7135 7136 7137 7138
        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 = {}
7139
        # don't offload the last checkpoints, to favor throughput
J
JZ-LIANG 已提交
7140 7141 7142 7143 7144 7145 7146 7147 7148 7149 7150 7151 7152 7153
        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(
7154 7155
            self.block.ops
        ), "Could NOT found backword op in prog"
J
JZ-LIANG 已提交
7156 7157 7158

        # fetch second to last checkpoint at the beginning of BW
        fetched_checkpoint_varname = self._record_fetch_op(
7159 7160
            self.bw_strart_op_idx
        )
J
JZ-LIANG 已提交
7161 7162
        last_last_fetch_checkpoint = None

7163
        for i, op in enumerate(self.block.ops[self.bw_strart_op_idx :]):
J
JZ-LIANG 已提交
7164 7165 7166 7167 7168 7169 7170 7171 7172
            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
7173 7174 7175
                            second_to_last_fetch_checkpoint = (
                                fetched_checkpoint_varname
                            )
7176
                            # there is NO fetch ahead the first checkpoint
J
JZ-LIANG 已提交
7177
                            if input_var != self.sorted_checkpoint_names[0]:
7178 7179 7180
                                fetched_checkpoint_varname = (
                                    self._record_fetch_op(idx)
                                )
J
JZ-LIANG 已提交
7181

7182
                        # should check the current used checkpoint is ths last fetch one
7183 7184 7185 7186 7187
                        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 已提交
7188 7189 7190
                        # rename
                        self.block.ops[idx]._rename_input(
                            input_var,
7191 7192
                            self.checkpoint_name2fetch_name[input_var],
                        )
J
JZ-LIANG 已提交
7193 7194 7195 7196
                        self.checkpoint_usage_count[input_var] += 1
                    else:
                        raise ValueError(
                            "use checkpoint [{}] before fetch in BW".format(
7197 7198 7199
                                input_var
                            )
                        )
J
JZ-LIANG 已提交
7200

7201 7202 7203 7204 7205
        assert (
            len(self.un_fetch_checkpoint_names) == 0
        ), "{} checkpoints have NOT been Recorded".format(
            self.un_fetch_checkpoint_names
        )
J
JZ-LIANG 已提交
7206 7207 7208 7209 7210 7211 7212 7213 7214 7215

    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)
7216
                    logging.debug(
7217 7218
                        "Insert [{}] fetch op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
7219 7220 7221 7222 7223
                    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()
7224 7225 7226 7227 7228
        assert (
            len(self.idx2insertions) == 0
        ), "{} checkpoints left un-Fecthed".format(
            [ele[1] for ele in self.idx2insertions.values()]
        )
J
JZ-LIANG 已提交
7229 7230 7231 7232

    def _parse_forward(self):

        self.idx2insertions = {}
7233
        # don't offload the last checkpoints, faster, less memory saving
J
JZ-LIANG 已提交
7234 7235 7236 7237 7238 7239 7240
        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,
7241
                'idx': -1,
J
JZ-LIANG 已提交
7242 7243 7244 7245 7246 7247 7248 7249 7250
            }
        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(
7251 7252
            self.block.ops
        ), "Could NOT found Forward op in prog"
J
JZ-LIANG 已提交
7253 7254
        last_offload_checkpoint = None

7255
        for i, op in enumerate(
7256 7257
            self.block.ops[self.fw_strart_op_idx : self.bw_strart_op_idx]
        ):
J
JZ-LIANG 已提交
7258 7259 7260 7261 7262 7263 7264

            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:
7265 7266 7267 7268 7269
                    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 已提交
7270 7271 7272

                    if output_var in self.un_offload_checkpoint_names:
                        # insert sync op if last checkpoint has not been sync
7273
                        if last_offload_checkpoint is not None:
7274 7275 7276 7277 7278 7279 7280 7281 7282
                            if (
                                self.checkpoint_usage_count_and_idx[
                                    last_offload_checkpoint
                                ]['count']
                                == 0
                            ):
                                self._record_sync_op(
                                    idx, last_offload_checkpoint
                                )
J
JZ-LIANG 已提交
7283
                            else:
7284 7285 7286 7287 7288 7289 7290 7291 7292 7293 7294 7295 7296
                                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 已提交
7297 7298 7299 7300 7301
                        # 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(
7302 7303 7304 7305
                            "There should be just ONE op that output checkpoint [{}]".format(
                                output_var
                            )
                        )
J
JZ-LIANG 已提交
7306 7307
                # need to sync the last need to offload checkpoint before the last checkpoint as output op
                if output_var == last_checkpoint:
7308 7309 7310 7311 7312 7313 7314 7315 7316 7317 7318 7319 7320
                    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 已提交
7321
                    # sync if last checkpoint has not been sync
7322 7323 7324 7325 7326 7327
                    if (
                        self.checkpoint_usage_count_and_idx[
                            last_offload_checkpoint
                        ]['idx']
                        == 0
                    ):
J
JZ-LIANG 已提交
7328 7329 7330
                        self._record_sync_op(idx, last_offload_checkpoint)
                    else:
                        last_usage_idx = self.checkpoint_usage_count_and_idx[
7331 7332 7333 7334 7335 7336 7337 7338 7339 7340
                            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
                        )
7341
            # record checkpoint usage
J
JZ-LIANG 已提交
7342 7343
            for input_var in input_vars:
                if input_var in need_offload_checkpoint_names:
7344 7345 7346
                    assert (
                        input_var not in self.synced_checkpoints
                    ), "checkpoint [{}] used after sync".format(input_var)
J
JZ-LIANG 已提交
7347 7348 7349
                    self.checkpoint_usage_count_and_idx[input_var]['count'] += 1
                    self.checkpoint_usage_count_and_idx[input_var]['idx'] = idx

7350 7351 7352 7353 7354
        assert (
            len(self.un_offload_checkpoint_names) == 0
        ), "{} checkpoints have NOT been Recorded".format(
            self.un_fetch_checkpoint_names
        )
J
JZ-LIANG 已提交
7355 7356 7357
        assert len(self.synced_checkpoints) == len(
            need_offload_checkpoint_names
        ), "{} checkpoints have NOT been Recorded".format(
7358 7359
            set(need_offload_checkpoint_names) - set(self.synced_checkpoints)
        )
J
JZ-LIANG 已提交
7360 7361 7362 7363 7364

    def _update_forward(self):
        if len(self.idx2insertions) == 0:
            return
        for op_idx in reversed(
7365 7366
            range(self.fw_strart_op_idx, self.bw_strart_op_idx)
        ):
J
JZ-LIANG 已提交
7367 7368 7369 7370
            if op_idx in self.idx2insertions:
                operation, checkpoint_name = self.idx2insertions[op_idx]
                if operation == "offload":
                    self._insert_offload_op(op_idx, checkpoint_name)
7371
                    logging.debug(
7372 7373
                        "Insert [{}] offload op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
7374 7375 7376
                    del self.idx2insertions[op_idx]
                elif operation == "sync":
                    self._insert_sync_op(op_idx, checkpoint_name)
7377
                    logging.debug(
7378 7379
                        "Insert [{}] offload_sync op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
7380 7381 7382
                    del self.idx2insertions[op_idx]

        self.block._sync_with_cpp()
7383 7384 7385 7386 7387
        assert (
            len(self.idx2insertions) == 0
        ), "{} checkpoints left un-Offloaded".format(
            [ele[1] for ele in self.idx2insertions.values()]
        )
J
JZ-LIANG 已提交
7388 7389 7390 7391 7392 7393 7394 7395

    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
7396
        1. create pinned vars and temp vars
J
JZ-LIANG 已提交
7397 7398 7399 7400 7401 7402
        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
7403
        if startup_program is None:
J
JZ-LIANG 已提交
7404
            startup_program = paddle.static.default_startup_program()
J
JZ-LIANG 已提交
7405 7406

        with program_guard(self._main_program, startup_program):
7407 7408 7409 7410 7411 7412 7413 7414 7415 7416
            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 已提交
7417 7418 7419 7420
            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(
7421 7422
                    checkpoint_varname
                )
J
JZ-LIANG 已提交
7423
                self.checkpoint_name2pinned_name[
7424 7425
                    checkpoint_varname
                ] = pinned_var_name
J
JZ-LIANG 已提交
7426
                self.checkpoint_name2fetch_name[
7427 7428
                    checkpoint_varname
                ] = fetch_var_name
J
JZ-LIANG 已提交
7429 7430 7431 7432 7433 7434 7435 7436 7437 7438 7439 7440 7441
            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

7442 7443 7444 7445 7446 7447 7448 7449
    def backward(
        self,
        loss,
        startup_program=None,
        parameter_list=None,
        no_grad_set=None,
        callbacks=None,
    ):
M
mapingshuo 已提交
7450 7451 7452 7453 7454 7455 7456
        """
        call append_backward with checkpoints.

        Args:
            loss (Variable): loss variable to run optimizations.
            startup_program (Program): startup_program for initializing parameters
                in `parameter_list`.
7457 7458
            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 已提交
7459 7460 7461 7462 7463 7464 7465 7466
            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
7467

M
mapingshuo 已提交
7468 7469 7470 7471 7472 7473
                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
7474 7475


M
mapingshuo 已提交
7476 7477 7478 7479
                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")
7480

M
mapingshuo 已提交
7481 7482
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
7483
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
7484 7485 7486 7487
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
7488
                    no_grad_set=None)
M
mapingshuo 已提交
7489 7490
                print("Finished backward")
        """
7491 7492 7493
        assert (
            self._checkpoints is not None
        ), "You should call _set_checkpoints first"
M
mapingshuo 已提交
7494

J
Jiabin Yang 已提交
7495
        if framework._non_static_mode():
M
mapingshuo 已提交
7496
            raise NotImplementedError(
7497 7498
                "DyGraph current does not support recompute"
            )
M
mapingshuo 已提交
7499 7500 7501 7502

        self._dtype = loss.dtype
        program = loss.block.program
        with program_guard(program, startup_program):
7503 7504 7505 7506 7507 7508 7509
            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 已提交
7510 7511 7512 7513 7514 7515
            # 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,
7516 7517
                    checkpoints=checkpoint_vars,
                )
J
JZ-LIANG 已提交
7518
            else:
7519 7520 7521 7522 7523 7524
                params_grads = append_backward(
                    loss,
                    parameter_list,
                    no_grad_set,
                    checkpoints=checkpoint_vars,
                )
J
JZ-LIANG 已提交
7525 7526 7527 7528 7529

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

M
mapingshuo 已提交
7530 7531 7532 7533 7534 7535 7536 7537 7538 7539 7540 7541 7542
        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
7543

M
mapingshuo 已提交
7544 7545 7546 7547 7548
                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)
7549 7550
                    return sum_cost, fc_1, prediction

M
mapingshuo 已提交
7551 7552 7553 7554
                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")
7555

M
mapingshuo 已提交
7556 7557
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
7558
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
7559 7560 7561 7562
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
7563
                    no_grad_set=None)
7564

M
mapingshuo 已提交
7565 7566
                optimize_ops = sgd.apply_optimize(
                    cost, startup_program=None, params_grads=params_grads)
7567

M
mapingshuo 已提交
7568 7569 7570
                print("Finished apply_optimize")
        """

7571 7572 7573 7574 7575 7576 7577 7578 7579 7580 7581 7582
        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
    ):
7583
        assert isinstance(loss, Variable), "The loss should be an Variable."
7584 7585 7586
        assert (
            self._checkpoints is not None
        ), "You should call _set_checkpoints first"
J
Jiabin Yang 已提交
7587
        if framework._non_static_mode():
M
mapingshuo 已提交
7588
            raise NotImplementedError(
7589 7590 7591 7592 7593 7594 7595 7596
                "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 已提交
7597

7598 7599 7600
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
M
mapingshuo 已提交
7601 7602 7603 7604

        return optimize_ops, params_grads


M
mapingshuo 已提交
7605
class LookaheadOptimizer(object):
7606
    r"""
7607
        :api_attr: Static Graph
S
swtkiwi 已提交
7608

M
mapingshuo 已提交
7609 7610 7611 7612
    This implements the Lookahead optimizer of the
    paper : https://arxiv.org/abs/1907.08610.

    Lookahead keeps two sets of params: the fast_params and
7613 7614
    the slow_params. inner_optimizer update fast_params every
    training step. Lookahead updates the slow_params and fast_params
M
mapingshuo 已提交
7615 7616 7617
    every k training steps as follows:

    .. math::
7618

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

7621
        fast\_param_t &=  slow\_param_t
M
mapingshuo 已提交
7622 7623

    Args:
7624
        inner_optimizer (Optimizer): The optimizer that update fast params step by step.
M
mapingshuo 已提交
7625 7626 7627 7628 7629 7630 7631 7632 7633
        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
7634
            import numpy.random as random
M
mapingshuo 已提交
7635

7636
            paddle.enable_static()
7637

7638 7639 7640 7641
            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)
7642
            loss = paddle.mean(x=loss)
7643 7644 7645 7646 7647 7648 7649 7650 7651
            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 已提交
7652

7653 7654 7655
            def train_reader(limit=5):
                for i in range(limit):
                    yield random.random([2]).astype('float32'), random.random([1]).astype('int64')
7656

7657 7658
            feeder = fluid.DataFeeder(feed_list=[x, label], place=place)
            reader = paddle.batch(paddle.reader.shuffle(train_reader, buf_size=50000),batch_size=1)
7659

7660 7661 7662
            for batch_data in reader():
                exe.run(fluid.default_main_program(),
                feed=feeder.feed(batch_data))
M
mapingshuo 已提交
7663 7664 7665 7666 7667

    """

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

J
Jiabin Yang 已提交
7668
        if framework._non_static_mode():
Z
zhongpu 已提交
7669
            raise Exception("In dygraph, don't support LookaheadOptimizer.")
7670
        assert inner_optimizer is not None, "inner optimizer can not be None"
M
mapingshuo 已提交
7671 7672 7673
        assert (
            0.0 <= alpha <= 1.0
        ), "alpha should be larger or equal to 0.0, and less or equal than 1.0"
7674
        assert isinstance(k, int) and k > 0, "k should be a positive integer"
M
mapingshuo 已提交
7675 7676 7677 7678 7679 7680 7681 7682 7683 7684

        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(
7685 7686
            loss, startup_program=startup_program
        )
M
mapingshuo 已提交
7687 7688 7689 7690 7691 7692 7693 7694 7695 7696 7697

        # 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)
7698 7699 7700 7701 7702 7703 7704
            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 已提交
7705 7706 7707 7708 7709 7710
            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)
7711 7712 7713 7714 7715 7716 7717
            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 已提交
7718

7719 7720 7721
            startup_block.append_op(
                type="assign", inputs={"X": fast_var}, outputs={"Out": slow_var}
            )
M
mapingshuo 已提交
7722

7723 7724
        with framework.program_guard(main_block.program, startup_program):
            # Add Var k to main prog and startup prog
7725 7726 7727 7728 7729 7730 7731
            k = layers.create_global_var(
                name="lookahead_k",
                shape=[1],
                value=int(self.k),
                dtype='int32',
                persistable=True,
            )
M
mapingshuo 已提交
7732

7733
            # Add Var alpha to main prog and startup prog
7734 7735 7736 7737 7738 7739 7740
            alpha = layers.create_global_var(
                name="lookahead_alpha",
                shape=[1],
                value=float(self.alpha),
                dtype='float32',
                persistable=True,
            )
M
mapingshuo 已提交
7741

7742
            # Add Var step
7743 7744 7745 7746 7747 7748 7749
            step = layers.create_global_var(
                name="lookahead_step",
                shape=[1],
                value=int(0),
                dtype='int32',
                persistable=True,
            )
7750 7751 7752
            layers.increment(x=step, value=1.0, in_place=True)

            # lookahead
7753 7754 7755
            zero_var = layers.fill_constant(
                shape=[1], dtype='float32', value=0.0
            )
7756

7757 7758 7759
            one_var = layers.fill_constant(
                shape=[1], dtype='float32', value=1.0
            )
7760 7761 7762

            mod = layers.elementwise_mod(step, k)
            with layers.control_flow.Switch() as switch:
7763 7764 7765 7766 7767
                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)
7768 7769 7770 7771 7772 7773 7774
                with switch.case(mod == zero_var):
                    for param_name in params:
                        fast_var = main_block.var(param_name)
                        slow_var = param_to_slow[param_name]
                        tmp_var = layers.elementwise_add(
                            layers.elementwise_mul(fast_var, alpha),
                            layers.elementwise_mul(
7775 7776 7777
                                slow_var, layers.elementwise_sub(one_var, alpha)
                            ),
                        )
7778 7779 7780 7781
                        layers.assign(input=tmp_var, output=slow_var)
                        layers.assign(input=tmp_var, output=fast_var)
                with switch.default():
                    pass
M
mapingshuo 已提交
7782
        return mini_out
7783 7784 7785 7786 7787 7788 7789 7790 7791 7792 7793 7794 7795 7796 7797 7798 7799 7800 7801 7802 7803 7804 7805 7806 7807 7808 7809 7810 7811 7812 7813 7814 7815 7816 7817 7818 7819 7820 7821 7822 7823 7824 7825 7826 7827 7828 7829 7830 7831 7832 7833 7834 7835 7836 7837 7838 7839


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

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

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

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

    Examples:
        .. code-block:: python

        import paddle.fluid as fluid
        import numpy as np

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

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

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

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

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

7840 7841
    GRAD_MERGE_COND_NAME = "grad_merge_cond_name"

7842
    def __init__(self, inner_optimizer, k_steps=1, avg=True):
J
Jiabin Yang 已提交
7843
        if framework._non_static_mode():
7844 7845 7846
            raise Exception(
                "In dygraph, we don't support GradientMergeOptimizer."
                "You can do Gradient merge by yourself with k-times forward + backward, "
7847 7848
                "and one-time optimizer.minimize()"
            )
7849

7850 7851 7852 7853
        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"
7854 7855 7856 7857 7858

        self.inner_optimizer = inner_optimizer
        self.k_steps = k_steps
        self.type = "gradient_merge"
        self.avg = avg
7859
        self._optimize_ops = None
7860

7861 7862 7863 7864 7865 7866
    def _set_k_steps(self, k_steps):
        self.k_steps = k_steps

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

7867 7868 7869 7870 7871 7872 7873 7874
    def backward(
        self,
        loss,
        startup_program=None,
        parameter_list=None,
        no_grad_set=None,
        callbacks=None,
    ):
7875 7876 7877 7878 7879 7880 7881 7882 7883
        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(
7884 7885
            loss, startup_program=startup_program
        )
7886 7887 7888 7889 7890 7891 7892 7893 7894 7895 7896
        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
7897 7898 7899
        if op_maker.kOpRoleVarAttrName() in op.attr_names and int(
            op.all_attrs()[op_maker.kOpRoleAttrName()]
        ) == int(backward):
7900 7901 7902 7903 7904 7905
            return True
        return False

    def _remove_op_role_var(self, param, grad):
        op_maker = core.op_proto_and_checker_maker
        op = grad.op
7906 7907 7908 7909 7910
        assert self._is_the_backward_op(
            op
        ), 'grad.op={} is not the backward op which produces the grad={}'.format(
            op, grad.name
        )
7911 7912 7913

        block = grad.block
        var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
7914 7915 7916 7917 7918 7919 7920 7921 7922 7923
        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
        )
7924 7925 7926 7927 7928 7929 7930 7931 7932 7933 7934 7935 7936 7937 7938 7939 7940 7941 7942 7943 7944 7945 7946 7947 7948 7949

        # 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
7950 7951 7952 7953 7954 7955 7956 7957 7958 7959 7960 7961 7962 7963 7964 7965 7966
        k_step_var = layers.create_global_var(
            name="gradient_merge_k",
            shape=[1],
            value=int(self.k_steps),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )

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

        # Add step var & cond var
7969 7970 7971 7972 7973 7974 7975 7976
        step_var = layers.create_global_var(
            name="gradient_merge_step",
            shape=[1],
            value=int(0),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )
7977

7978 7979 7980
        cond_var = main_block.create_var(
            name="gradient_merge_cond", shape=[1], dtype='bool'
        )
7981 7982 7983 7984

        with device_guard("cpu"):
            # step_var = (step_var + 1) % k_step
            layers.increment(x=step_var, value=1.0, in_place=True)
7985 7986 7987 7988 7989 7990
            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},
            )
7991 7992

            # cond_var = (step_var == 0)
7993 7994 7995 7996 7997
            main_block.append_op(
                type='equal',
                inputs={'X': step_var, 'Y': zero_var},
                outputs={'Out': cond_var},
            )
7998 7999 8000 8001 8002 8003 8004 8005 8006 8007

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

8009
        # TODO(mapingshuo) support sparse embedding
8010 8011
        # step1: remove grad.op's op_role_var
        for param, grad in params_grads:
8012
            assert (
8013
                param.type != core.VarDesc.VarType.SELECTED_ROWS
8014 8015
            ), "SELECTED_ROWS is not supported in GradientMergeOptimizer for now"

8016
            self._remove_op_role_var(param, grad)
8017

8018
        param_to_grad = {k.name: v for (k, v) in params_grads}
8019 8020 8021
        param_names = param_to_grad.keys()
        param_to_gradient_merge = {}

8022 8023 8024 8025 8026
        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
8027
            param_var = main_block.var(param_name)
8028 8029 8030 8031 8032 8033 8034
            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,
            )
8035
            param_to_gradient_merge[param_name] = gradient_merge_var
8036

8037 8038 8039 8040
            startup_gradient_merge_var = startup_block.create_var(
                name=param_name + "@GRAD@GradientMerge",
                shape=param_var.shape,
                dtype=param_var.dtype,
8041 8042 8043 8044 8045 8046 8047 8048 8049 8050 8051
                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),
                },
            )
8052

8053 8054 8055
            # grad_merge += grad
            new_grad_op = main_block.append_op(
                type="elementwise_add",
8056
                inputs={'X': grad, 'Y': gradient_merge_var},
8057
                outputs={'Out': gradient_merge_var},
8058 8059 8060 8061 8062
                attrs={'axis': -1, 'use_mkldnn': False},
            )
            self._add_gm_op_role_var(
                new_grad_op, param, gradient_merge_var, cond
            )
8063 8064 8065 8066 8067 8068 8069 8070
            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)
8071
            op_maker = core.op_proto_and_checker_maker
8072 8073 8074 8075

            if self.avg:
                for param, new_grad in new_params_grads:
                    # grad /= k_steps
8076 8077 8078 8079 8080 8081 8082 8083 8084 8085 8086 8087 8088
                    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
                    )
8089

8090 8091 8092 8093 8094 8095
            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
8096

8097
            self._optimize_ops = self.inner_optimizer.apply_gradients(
8098 8099
                new_params_grads
            )
8100

8101 8102
            # clear gradient_merge_vars
            for param, new_grad in new_params_grads:
8103 8104 8105 8106 8107 8108 8109 8110 8111
                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
                )
8112 8113 8114 8115 8116 8117

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

        return self._optimize_ops

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

8123 8124 8125 8126 8127 8128
        params_grads = self.backward(
            loss,
            startup_program=startup_program,
            parameter_list=parameter_list,
            no_grad_set=no_grad_set,
        )
8129

8130 8131 8132
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
8133 8134

        return optimize_ops, params_grads