optimizer.py 299.8 KB
Newer Older
1
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
14

15
import numpy as np
16
import os
17
import logging
18
from collections import defaultdict
19

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

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

69
__all__ = [
70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96
    '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',
97
]
Q
Qiao Longfei 已提交
98 99


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

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

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

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

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

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

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

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

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

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

                import paddle.fluid as fluid
216
                import paddle
217 218

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

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

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

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

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

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

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

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

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

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

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

                paddle.disable_static()

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

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

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

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

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

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

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

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

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

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

339 340 341 342 343
            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
            )
344

345 346 347 348 349
            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
            )
350 351 352

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

                    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
501
                    lr_var = paddle.static.create_global_var(
502 503 504 505 506 507 508 509 510 511 512 513 514
                        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."
515 516
                % (type(value))
            )
517 518 519 520 521 522 523 524
        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:
525 526
                if in_dygraph_mode():
                    place = _current_expected_place()
527 528 529 530 531 532 533
                    _C_ops.full_(
                        current_lr,
                        list(current_lr.shape),
                        float(value),
                        current_lr.dtype,
                        place,
                    )
534
                else:
535 536 537 538 539 540 541 542 543 544 545 546 547
                    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,
                    )
548
        else:
549 550 551
            assert (
                len(value.shape) == 1 and value.shape[0] == 1
            ), "optimizer's learning rate must be 1-D Tensor with shape[1]"
552 553
            self._learning_rate_map[framework.default_main_program()] = value

554 555 556
    @framework.dygraph_only
    def current_step_lr(self):
        """
557
        :api_attr: imperative
558

559 560 561 562 563 564 565 566 567 568 569
        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
570
                import paddle
571 572 573

                # example1: LearningRateDecay is not used, return value is all the same
                with fluid.dygraph.guard():
574
                    emb = paddle.nn.Embedding(10, 10)
575 576 577 578 579 580 581
                    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")
582
                    linear = paddle.nn.Linear(10, 10)
583 584
                    inp = fluid.dygraph.to_variable(inp)
                    out = linear(inp)
585
                    loss = paddle.mean(out)
586

587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603
                    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()
604
        if isinstance(current_lr, framework.Variable):
605 606 607 608
            return self._global_learning_rate().numpy()[0]

        if isinstance(self._learning_rate, float):
            return self._learning_rate
609 610 611
        elif isinstance(self._learning_rate, _LearningRateEpochDecay):
            step_lr = self._learning_rate()
            return step_lr.numpy()[0]
612 613 614 615 616 617 618
        else:
            step_lr = self._learning_rate.step()
            if isinstance(step_lr, (float, int)):
                return step_lr
            else:
                return step_lr.numpy()[0]

Y
yuyang18 已提交
619
    def _global_learning_rate(self, program=None):
Q
Qiao Longfei 已提交
620 621 622 623
        """
        get global decayed learning rate
        :return:
        """
624 625
        if program is None:
            program = framework.default_main_program()
Q
qiaolongfei 已提交
626
        return self._learning_rate_map.get(program, None)
Q
Qiao Longfei 已提交
627

Q
Qiao Longfei 已提交
628
    def _append_optimize_op(self, block, param_and_grad):
629
        """append optimize operator to block and return all the added optimize_op"""
Q
Qiao Longfei 已提交
630 631
        raise NotImplementedError()

632 633 634 635
    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 已提交
636 637
        if type(param_lr) == Variable:
            return param_lr
Q
qiaolongfei 已提交
638
        else:
W
Wu Yi 已提交
639
            if param_lr == 1.0:
Y
yuyang18 已提交
640
                return self._global_learning_rate()
W
Wu Yi 已提交
641
            else:
X
Xin Pan 已提交
642
                with default_main_program()._lr_schedule_guard(
643 644
                    is_with_opt=True
                ), framework.name_scope('scale_with_param_lr'):
645
                    return self._global_learning_rate() * param_lr
646 647 648 649 650 651 652

    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 已提交
653
        """
654 655
        pass

656
    def _finish_update(self, block, parameters_and_grads):
657 658 659 660 661 662 663 664
        """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 已提交
665
            None
666 667 668
        """
        pass

669 670 671 672 673 674 675 676 677 678
    def _add_accumulator(
        self,
        name,
        param,
        dtype=None,
        fill_value=0.0,
        shape=None,
        type=None,
        device=None,
    ):
679 680 681 682 683 684 685 686 687
        """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 已提交
688 689
        if self._name is not None:
            name = self._name + "_" + name
690 691 692 693
        if (
            name in self._accumulators
            and param.name in self._accumulators[name]
        ):
姜永久 已提交
694
            if in_dygraph_mode():
X
polish  
Xin Pan 已提交
695
                return self._accumulators[name][param.name]
696 697
            raise Exception(
                "Accumulator {} already exists for parameter {}".format(
698 699 700
                    name, param.name
                )
            )
701
        if shape is None:
702
            shape = param.shape
Q
Qiao Longfei 已提交
703
        assert isinstance(self.helper, LayerHelper)
704 705 706 707 708

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

Q
Qiao Longfei 已提交
709
        var = self.helper.create_global_variable(
710
            name=var_name,
Q
Qiao Longfei 已提交
711
            persistable=True,
F
fengjiayi 已提交
712
            dtype=dtype or param.dtype,
713
            type=core.VarDesc.VarType.LOD_TENSOR
姜永久 已提交
714
            if in_dygraph_mode()
715
            else (param.type if type is None else type),
H
hong 已提交
716
            shape=shape,
717 718
            belong_to_optimizer=True,
        )
719 720 721 722
        if device is None:
            device = self._get_device_for_param(param.name)
        with device_guard(device):
            self.helper.set_variable_initializer(
723 724
                var, initializer=Constant(value=float(fill_value))
            )
H
hong 已提交
725

姜永久 已提交
726
        if in_dygraph_mode():
H
hong 已提交
727
            if len(self._accumulators_holder) > 0:
728 729 730 731 732
                assert (
                    var_name in self._accumulators_holder
                ), "Optimizer set error, {} should in state dict".format(
                    var_name
                )
H
hong 已提交
733 734
                var.set_value(self._accumulators_holder[var_name])

Q
Qiao Longfei 已提交
735
        self._accumulators[name][param.name] = var
736
        return var
737

738 739 740 741 742 743 744 745 746
    def _add_global_accumulator(
        self,
        name,
        dtype=None,
        fill_value=0.0,
        shape=None,
        type=None,
        device=None,
    ):
747 748 749 750 751 752 753 754 755 756 757 758 759
        """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
760
        if name in self._global_accumulators:
姜永久 已提交
761
            if in_dygraph_mode():
762 763
                return self._global_accumulators[name]
            raise Exception("Global accumulator {} already exists".format(name))
764
        if shape is None:
765 766 767 768 769 770 771 772 773 774 775 776 777
            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,
778 779
            belong_to_optimizer=True,
        )
780 781 782 783
        if device is None:
            device = 'cpu'
        with device_guard(device):
            self.helper.set_variable_initializer(
784 785
                var, initializer=Constant(value=float(fill_value))
            )
786

姜永久 已提交
787
        if in_dygraph_mode():
788
            if len(self._accumulators_holder) > 0:
789 790 791 792 793
                assert (
                    var_name in self._accumulators_holder
                ), "Optimizer set error, {} should in state dict".format(
                    var_name
                )
794 795 796 797 798
                var.set_value(self._accumulators_holder[var_name])

        self._global_accumulators[name] = var
        return var

799 800 801 802 803 804 805 806
    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:
807
            accumulator variable
808
        """
W
whs 已提交
809 810
        if self._name is not None:
            name = self._name + "_" + name
811 812 813 814
        if (
            name not in self._accumulators
            or param.name not in self._accumulators[name]
        ):
815 816
            raise Exception(
                "Accumulator {} does not exist for parameter {}".format(
817 818 819
                    name, param.name
                )
            )
820 821
        return self._accumulators[name][param.name]

822 823 824 825 826 827 828 829 830 831 832
    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
833
        if name not in self._global_accumulators:
834 835 836
            raise Exception("Global accumulator {} does not exist".format(name))
        return self._global_accumulators[name]

837 838 839 840 841
    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
842 843
                device_attr_name = (
                    core.op_proto_and_checker_maker.kOpDeviceAttrName()
844 845 846 847 848
                )
                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(
849 850
                            device_attr_name
                        )
851
                        break
852 853 854 855 856 857 858

    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

859
    def _create_optimization_pass(self, parameters_and_grads):
Q
Qiao Longfei 已提交
860 861 862
        """Add optimization operators to update gradients to variables.

        Args:
Q
qiaolongfei 已提交
863
          parameters_and_grads(list(tuple(Variable, Variable))):
864
            a list of (variable, gradient) pair to update.
Q
Qiao Longfei 已提交
865 866

        Returns:
867
          return_op_list: a list of operators that will complete one step of
868 869 870
            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 已提交
871
        """
872 873 874 875 876
        # 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
877
        # for parameters and extend _finish_update method to add custom ops.
878

879
        # Allways called under program_guard use global block as loss block
880 881 882
        # But if current block is in control flow, append optimize op in the
        # grad block of current block

883
        global_block = framework.default_main_program().global_block()
884 885 886
        target_block = global_block
        current_block = framework.default_main_program().current_block()
        if current_block.idx != global_block.idx:
887 888 889
            assert (
                current_block.backward_block_idx != -1
            ), "current block is not global_block, but it doesn't have backward block."
890
            target_block = framework.default_main_program().blocks[
891 892
                current_block.backward_block_idx
            ]
893 894

        start = len(target_block.ops)
895

896
        self._update_param_device_map(parameters_and_grads, target_block)
C
chengduo 已提交
897
        self._create_accumulators(
898 899
            target_block, [p[0] for p in parameters_and_grads if p[0].trainable]
        )
900 901
        self._create_global_learning_rate()

姜永久 已提交
902
        if in_dygraph_mode():
903 904 905
            for param_and_grad in parameters_and_grads:
                if param_and_grad[1] is None:
                    continue
906 907
                if param_and_grad[0].trainable is True:
                    self._append_optimize_op(target_block, param_and_grad)
908 909 910 911 912
        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(
913 914
                    param_and_grad
                ), name_scope("optimizer"):
915
                    if param_and_grad[0].trainable is True:
916
                        device = self._get_device_for_param(
917 918
                            param_and_grad[0].name
                        )
919 920
                        with device_guard(device):
                            optimize_op = self._append_optimize_op(
921 922
                                target_block, param_and_grad
                            )
923 924 925

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

928 929
        end = len(target_block.ops)
        return target_block._slice_ops(start, end)
930 931

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

976 977 978 979 980 981 982 983
    def backward(
        self,
        loss,
        startup_program=None,
        parameter_list=None,
        no_grad_set=None,
        callbacks=None,
    ):
984
        """
985
        The first part of ``minimize``, do auto-diff to append backward operations for
986 987 988
        the current program.

        Args:
989 990 991 992
            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 已提交
993
            parameter_list (Iterable, optional): Iterable of ``Variable`` or ``Variable.name`` to update
994 995
                to minimize ``loss``. The default value is None, at this time all parameters
                will be updated.
996
            no_grad_set (set, optional): Set of ``Variable``  or ``Variable.name`` that don't need
997 998 999
                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 已提交
1000

1001
        Return:
1002 1003
            list: list of (param, grad) variable pairs, param is ``Parameter``,
                grad is the gradient value corresponding to the parameter.
M
minqiyang 已提交
1004

1005
        Examples:
1006
            See examples in ``apply_gradients``.
1007
        """
1008
        act_no_grad_set = None
姜永久 已提交
1009
        if in_dygraph_mode():
1010
            pass
L
Leo Chen 已提交
1011 1012
        else:
            act_no_grad_set = self._get_no_grad_set(loss, no_grad_set)
G
gongweibao 已提交
1013

L
Leo Chen 已提交
1014 1015 1016 1017
        # Infer dtype by loss if None
        if self._dtype is None:
            self._dtype = loss.dtype

姜永久 已提交
1018
        if in_dygraph_mode():
1019 1020 1021
            parameter_list = (
                parameter_list if parameter_list else self._parameter_list
            )
1022

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

1052
    def _create_regularization_of_grad(self, param, grad, regularization=None):
1053
        """Create and add backward regularization Operators
1054

1055 1056 1057
        Function helper of append_regularization_ops.
        """
        # If no gradient or no regularization is specified,  then we don't need to do anything
1058
        if grad is None or (
1059 1060 1061 1062 1063 1064
            (
                not hasattr(param, 'regularizer')
                or (hasattr(param, 'regularizer') and param.regularizer is None)
            )
            and regularization is None
        ):
1065 1066 1067 1068 1069 1070 1071 1072 1073 1074
            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

姜永久 已提交
1075
        if in_dygraph_mode():
1076
            return _legacy_C_ops.sum([grad, regularization_term])
1077

1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088
        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,
1089 1090
                type=core.VarDesc.VarType.LOD_TENSOR,
            )
1091 1092 1093

        inputs = {"X": [grad, regularization_term]}
        outputs = {"Out": [new_grad]}
1094
        grad.block.append_op(type='sum', inputs=inputs, outputs=outputs)
1095 1096 1097

        return new_grad

1098 1099 1100
    def append_regularization_ops(
        self, parameters_and_grads, regularization=None
    ):
1101
        r"""Create and add backward regularization Operators
1102

1103 1104 1105 1106
        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.
1107

1108 1109 1110 1111 1112
        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.
1113

1114 1115 1116
        Returns:
            list[(Variable, Variable)]: list of (parameters, gradients) \
            pair with the regularized gradient
1117

1118 1119 1120 1121
        Raises:
            Exception: Unknown regularization type
        """
        params_and_grads = []
姜永久 已提交
1122
        if in_dygraph_mode():
1123
            for param, grad in parameters_and_grads:
1124
                new_grad = self._create_regularization_of_grad(
1125 1126
                    param, grad, regularization
                )
1127 1128 1129 1130 1131
                params_and_grads.append((param, new_grad))
        else:
            repeate_regularizer = False
            with framework.name_scope('regularization'):
                for param, grad in parameters_and_grads:
1132 1133 1134 1135 1136
                    if (
                        not repeate_regularizer
                        and getattr(param, 'regularizer', None) is not None
                        and regularization is not None
                    ):
1137 1138 1139 1140
                        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!"
1141 1142
                            % regularization.__str__()
                        )
1143 1144
                    with param.block.program._optimized_guard([param, grad]):
                        new_grad = self._create_regularization_of_grad(
1145 1146
                            param, grad, regularization
                        )
1147 1148 1149
                        params_and_grads.append((param, new_grad))
        return params_and_grads

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

        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)],
1191 1192
            belong_to_optimizer=True,
        )
1193 1194

        with program_guard(default_main_program()):
1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223
            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,
                },
            )
1224

1225
        # NOTE(zhiqiu): the initializer should be set after coalesce_tensor op,
1226
        # so the shape of flatten_param and flatten_grad will be inferred.
1227 1228 1229 1230 1231 1232
        self.helper.set_variable_initializer(
            flatten_param, initializer=Constant(0.0)
        )
        self.helper.set_variable_initializer(
            flatten_grad, initializer=Constant(0.0)
        )
1233 1234 1235

        return [(flatten_param, flatten_grad)]

1236 1237 1238 1239 1240 1241 1242
    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 已提交
1243

1244 1245
        Returns:
            list: A list of operators appended to the current program.
M
minqiyang 已提交
1246

1247 1248 1249
        Examples:
            .. code-block:: python

1250
                import paddle.fluid as fluid
1251 1252 1253 1254 1255 1256 1257 1258 1259
                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)

1260 1261
        # NOTE(zhiqiu): currently, only support ClipGradByGlobalNorm and without regularization.
        if self._flatten_param_grads and self.regularization is None:
1262
            if self._grad_clip is None or isinstance(
1263 1264
                self._grad_clip, ClipGradByGlobalNorm
            ):
1265 1266
                params_grads = self.flatten_param_grads(params_grads)

1267
        # 'optimizer(grad_clip)' or 'set_gradient_clip'
1268 1269 1270 1271
        if self._grad_clip is not None:
            params_grads = self._grad_clip(params_grads)
        else:
            params_grads = append_gradient_clip_ops(params_grads)
1272 1273

        # Add regularization if any
1274 1275 1276
        params_grads = self.append_regularization_ops(
            params_grads, self.regularization
        )
1277 1278 1279 1280

        optimize_ops = self._create_optimization_pass(params_grads)
        return optimize_ops

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

        return no_grad_set

1321 1322 1323 1324
    @framework.dygraph_only
    def clear_gradients(self):
        """
        Clear the gradients of all optimized parameters for model.
1325 1326

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

1328 1329
        Returns:
            None
1330

1331 1332 1333 1334
        Examples:
            .. code-block:: python

                import paddle.fluid as fluid
1335
                import paddle
1336 1337 1338 1339 1340
                import numpy as np

                with fluid.dygraph.guard():
                    value = np.arange(26).reshape(2, 13).astype("float32")
                    a = fluid.dygraph.to_variable(value)
1341
                    linear = paddle.nn.Linear(13, 5)
1342
                    # This can be any optimizer supported by dygraph.
1343
                    adam = fluid.optimizer.Adam(learning_rate = 0.01,
1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354
                                                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()

1355
    @imperative_base.no_grad
1356 1357 1358
    def minimize(
        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
    ):
1359
        """
1360
        Add operations to minimize ``loss`` by updating ``parameter_list``.
M
minqiyang 已提交
1361

1362
        Args:
1363 1364 1365 1366
            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 已提交
1367
            parameter_list (Iterable, optional): Iterable of ``Variable`` or ``Variable.name`` to update
1368 1369
                to minimize ``loss``. The default value is None, at this time all parameters
                will be updated.
1370
            no_grad_set (set, optional): Set of ``Variable``  or ``Variable.name`` that don't need
1371
                to be updated. The default value is None.
Q
Qiao Longfei 已提交
1372

1373
        Returns:
1374 1375 1376
            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.
1377 1378
            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
1379
            ``fetch_list`` before run, see details in ``Executor``.
1380 1381 1382

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

1386 1387 1388
        parameter_list = (
            parameter_list if parameter_list else self._parameter_list
        )
1389

1390 1391 1392 1393 1394 1395
        params_grads = self.backward(
            loss,
            startup_program=startup_program,
            parameter_list=parameter_list,
            no_grad_set=no_grad_set,
        )
1396

1397 1398 1399
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
M
minqiyang 已提交
1400

Q
Qiao Longfei 已提交
1401
        return optimize_ops, params_grads
Q
Qiao Longfei 已提交
1402 1403 1404


class SGDOptimizer(Optimizer):
1405
    r"""
Q
qiaolongfei 已提交
1406 1407 1408 1409 1410 1411
    Optimizer of the stochastic gradient descent algorithm.

    .. math::

        param\_out = param - learning\_rate * grad

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

    Examples:
        .. code-block:: python

1433 1434 1435 1436
            import paddle
            import paddle.fluid as fluid
            import numpy as np

1437
            paddle.enable_static()
1438 1439 1440 1441 1442 1443
            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)
1444
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
1445
                avg_cost = paddle.mean(cost)
1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458

                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 已提交
1459 1460
    """

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

1530
    @no_grad
1531
    def _append_optimize_op(self, block, param_and_grad):
1532

1533 1534 1535 1536 1537 1538 1539 1540 1541
        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
        )
1542

1543
        lr = self._create_param_lr(param_and_grad)
Z
zyfncg 已提交
1544
        if in_dygraph_mode():
1545 1546 1547 1548 1549 1550 1551
            _C_ops.sgd_(
                param_and_grad[0],
                lr,
                param_and_grad[1],
                master_weight,
                find_master,
            )
Z
zyfncg 已提交
1552
            return None
姜永久 已提交
1553 1554 1555 1556 1557 1558 1559 1560
        else:
            assert isinstance(block, framework.Block)
            # create the optimize op
            inputs = {
                "Param": param_and_grad[0],
                "Grad": param_and_grad[1],
                "LearningRate": lr,
            }
1561

姜永久 已提交
1562
            outputs = {"ParamOut": param_and_grad[0]}
1563

姜永久 已提交
1564
            attrs = {"multi_precision": find_master}
1565

姜永久 已提交
1566 1567 1568
            if find_master:
                inputs["MasterParam"] = master_weight
                outputs["MasterParamOut"] = master_weight
1569

姜永久 已提交
1570 1571 1572 1573 1574 1575 1576
            sgd_op = block.append_op(
                type=self.type,
                inputs=inputs,
                outputs=outputs,
                attrs=attrs,
                stop_gradient=True,
            )
Q
Qiao Longfei 已提交
1577

姜永久 已提交
1578
            return sgd_op
1579 1580 1581


class MomentumOptimizer(Optimizer):
1582
    r"""
Q
qiaolongfei 已提交
1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595

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

1596
        &\quad   param = param - (gradient + mu * velocity) * learning\_rate
Q
qiaolongfei 已提交
1597 1598 1599

        & else:

Q
qiaolongfei 已提交
1600
        &\quad   param = param - learning\_rate * velocity
Q
qiaolongfei 已提交
1601

1602 1603 1604 1605
    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 已提交
1606
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1607
            This parameter is required in dygraph mode. \
1608
            The default value is None in static graph mode, at this time all parameters will be updated.
1609
        use_nesterov (bool, optional): Enables Nesterov momentum, default is false.
1610 1611 1612 1613 1614
        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.
1615 1616 1617
        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` ,
1618
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
1619 1620
        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 已提交
1621 1622 1623 1624

    Examples:
        .. code-block:: python

1625 1626 1627 1628
            import paddle
            import paddle.fluid as fluid
            import numpy as np

1629
            paddle.enable_static()
1630 1631 1632 1633 1634 1635
            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)
1636
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
1637
                avg_cost = paddle.mean(cost)
1638 1639 1640 1641 1642 1643 1644 1645 1646 1647 1648 1649 1650

                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)

1651 1652 1653
    """
    _velocity_acc_str = "velocity"

1654 1655 1656 1657 1658 1659 1660 1661 1662 1663
    def __init__(
        self,
        learning_rate,
        momentum,
        parameter_list=None,
        use_nesterov=False,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
1664 1665
        assert learning_rate is not None
        assert momentum is not None
1666
        super().__init__(
1667 1668 1669 1670 1671 1672
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
1673 1674
        self.type = "momentum"
        self._momentum = momentum
1675
        self._use_nesterov = bool(use_nesterov)
1676 1677 1678 1679 1680

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

        for p in parameters:
Q
Qiao Longfei 已提交
1681
            self._add_accumulator(self._velocity_acc_str, p)
1682 1683 1684 1685

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

1686 1687 1688
        velocity_acc = self._get_accumulator(
            self._velocity_acc_str, param_and_grad[0]
        )
1689
        lr = self._create_param_lr(param_and_grad)
1690
        master_weight = None
姜永久 已提交
1691
        if in_dygraph_mode():
1692
            _, _, _ = _legacy_C_ops.momentum(
1693 1694 1695 1696 1697 1698 1699 1700 1701 1702 1703 1704 1705
                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,
            )
1706
            return None
姜永久 已提交
1707 1708 1709 1710 1711 1712 1713 1714
        else:
            attrs = {"mu": self._momentum, "use_nesterov": self._use_nesterov}
            inputs = {
                "Param": [param_and_grad[0]],
                "Grad": [param_and_grad[1]],
                "Velocity": [velocity_acc],
                "LearningRate": [lr],
            }
1715

姜永久 已提交
1716 1717 1718 1719 1720 1721 1722 1723 1724 1725 1726 1727
            outputs = {
                "ParamOut": [param_and_grad[0]],
                "VelocityOut": [velocity_acc],
            }
            # create the momentum optimize op
            momentum_op = block.append_op(
                type=self.type,
                inputs=inputs,
                outputs=outputs,
                attrs=attrs,
                stop_gradient=True,
            )
1728

姜永久 已提交
1729
            return momentum_op
1730 1731


1732
class LarsMomentumOptimizer(Optimizer):
1733
    r"""
1734 1735 1736 1737 1738 1739 1740 1741 1742
    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||}

1743
        & velocity = mu * velocity + local\_learning\_rate * (gradient + lars\_weight\_decay * param + epsilon)
1744 1745 1746

        & param = param - velocity

1747 1748 1749 1750 1751 1752
    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 已提交
1753
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
1754
            This parameter is required in dygraph mode. \
1755
            The default value is None in static graph mode, at this time all parameters will be updated.
1756 1757 1758 1759 1760
        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.
1761 1762 1763
        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` ,
1764
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
1765 1766
        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.
1767 1768
        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.
1769 1770 1771
        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`.
1772

1773 1774 1775
    Examples:
        .. code-block:: python

1776
            import paddle
1777 1778 1779
            import paddle.fluid as fluid
            import numpy as np

1780
            paddle.enable_static()
1781 1782 1783 1784
            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)
1785
            out = paddle.sum(out)
1786 1787 1788 1789 1790 1791 1792 1793
            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])
1794 1795 1796
    """
    _velocity_acc_str = "velocity"

1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811
    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,
    ):
1812 1813
        assert learning_rate is not None
        assert momentum is not None
1814
        super().__init__(
1815 1816 1817 1818 1819 1820
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
1821 1822 1823 1824
        self.type = "lars_momentum"
        self._momentum = momentum
        self._lars_coeff = float(lars_coeff)
        self._lars_weight_decay = float(lars_weight_decay)
1825 1826 1827 1828 1829
        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
1830 1831 1832 1833 1834
        self._multi_precision = multi_precision
        self._rescale_grad = float(rescale_grad)
        self._master_weights = {}

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

1840 1841
            var_name = param.name + '_fp32_master'
            var_name = unique_name.generate(var_name)
1842
            var = paddle.static.create_global_var(
1843 1844 1845 1846 1847 1848
                name=var_name,
                shape=param.shape,
                value=0,
                dtype='float32',
                persistable=True,
            )
1849
            block = self.helper.startup_program.global_block()
1850 1851 1852 1853 1854 1855 1856 1857 1858
            block.append_op(
                type="cast",
                inputs={"X": [param]},
                outputs={"Out": [var]},
                attrs={
                    "in_dtype": param.dtype,
                    "out_dtype": core.VarDesc.VarType.FP32,
                },
            )
1859
            self._master_weights[param.name] = var
1860 1861 1862 1863 1864 1865 1866 1867 1868 1869 1870 1871
        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
1872 1873 1874 1875 1876 1877
        find_master = (
            self._multi_precision and param.dtype == core.VarDesc.VarType.FP16
        )
        target_param = (
            self._master_weights[param.name] if find_master else param
        )
1878
        target_name = target_param.name
1879 1880 1881 1882
        if (
            name not in self._accumulators
            or target_name not in self._accumulators[name]
        ):
1883 1884
            raise Exception(
                "Accumulator {} does not exist for parameter {}".format(
1885 1886 1887
                    name, target_name
                )
            )
1888
        return self._accumulators[name][target_name]
1889 1890 1891 1892 1893

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

        for p in parameters:
1894 1895 1896 1897
            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
1898 1899 1900 1901
            if (
                p.dtype == core.VarDesc.VarType.FP16
                and not self._multi_precision
            ):
1902 1903 1904 1905
                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."
                )
1906 1907 1908 1909
            self._add_accumulator(self._velocity_acc_str, p)

    def _append_optimize_op(self, block, param_and_grad):
        assert isinstance(block, framework.Block)
1910 1911 1912 1913 1914 1915 1916 1917
        _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

1918 1919 1920
        velocity_acc = self._get_accumulator(
            self._velocity_acc_str, param_and_grad[0]
        )
1921 1922
        lr = self._create_param_lr(param_and_grad)

1923 1924 1925 1926 1927 1928 1929 1930 1931
        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
        )
1932 1933 1934

        attrs = {
            "mu": self._momentum,
1935
            "lars_coeff": self._lars_coeff,
L
limingshu 已提交
1936
            "lars_weight_decay": [_lars_weight_decay],
1937
            "multi_precision": find_master,
L
limingshu 已提交
1938
            "epsilon": self._epsilon,
1939
            "rescale_grad": self._rescale_grad,
1940 1941 1942 1943 1944 1945
        }

        inputs = {
            "Param": param_and_grad[0],
            "Grad": param_and_grad[1],
            "Velocity": velocity_acc,
1946
            "LearningRate": lr,
1947 1948 1949 1950 1951 1952 1953 1954
        }

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

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

姜永久 已提交
1955
        if in_dygraph_mode():
1956
            tmp, tmp2 = _legacy_C_ops.lars_momentum(
1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967 1968 1969 1970 1971 1972 1973 1974 1975
                [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,
            )
1976 1977
        else:
            # create the momentum optimize op
1978 1979 1980 1981 1982 1983 1984
            momentum_op = block.append_op(
                type=self.type,
                inputs=inputs,
                outputs=outputs,
                attrs=attrs,
                stop_gradient=True,
            )
1985

1986
            return momentum_op
1987 1988


1989
class AdagradOptimizer(Optimizer):
1990
    r"""
1991 1992
    The Adaptive Gradient optimizer (Adagrad for short) can adaptively assign
    different learning rates to individual parameters.
Q
qiaolongfei 已提交
1993

1994
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
1995 1996 1997 1998 1999 2000 2001

    .. math::

        moment\_out &= moment + grad * grad

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

2002 2003 2004 2005 2006 2007
    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 已提交
2008 2009 2010
    for numerical stability to avoid the division by zero error.

    Args:
2011 2012 2013 2014
        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 已提交
2015
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2016
            This parameter is required in dygraph mode. \
2017
            The default value is None in static graph mode, at this time all parameters will be updated.
2018 2019 2020 2021 2022
        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.
2023 2024 2025
        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` ,
2026
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2027 2028 2029 2030 2031
        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 已提交
2032 2033 2034 2035

    Examples:
        .. code-block:: python

2036
            import paddle
2037
            import numpy as np
2038
            import paddle.fluid as fluid
2039

2040
            paddle.enable_static()
2041
            np_inp = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
2042
            inp = fluid.data(name="inp", shape=[2, 2])
2043
            out = fluid.layers.fc(inp, size=3)
2044
            out = paddle.sum(out)
2045
            optimizer = fluid.optimizer.AdagradOptimizer(learning_rate=0.2)
2046 2047 2048 2049 2050 2051 2052
            optimizer.minimize(out)

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

2056 2057 2058 2059 2060 2061 2062 2063 2064 2065
    def __init__(
        self,
        learning_rate,
        epsilon=1.0e-6,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
        initial_accumulator_value=0.0,
    ):
2066 2067
        assert learning_rate is not None
        assert epsilon is not None
2068
        super().__init__(
2069 2070 2071 2072 2073 2074
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
2075 2076
        self.type = "adagrad"
        self._epsilon = epsilon
2077
        self.initial_accumulator_value = initial_accumulator_value
2078 2079 2080 2081 2082

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

        for p in parameters:
2083 2084 2085 2086 2087
            self._add_accumulator(
                self._moment_acc_str,
                p,
                fill_value=self.initial_accumulator_value,
            )
2088 2089 2090 2091

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

2092 2093 2094
        moment_acc = self._get_accumulator(
            self._moment_acc_str, param_and_grad[0]
        )
C
caozhou 已提交
2095
        if in_dygraph_mode():
2096 2097 2098 2099 2100 2101 2102
            _C_ops.adagrad_(
                param_and_grad[0],
                param_and_grad[1],
                moment_acc,
                self._create_param_lr(param_and_grad),
                self._epsilon,
            )
C
caozhou 已提交
2103
            return None
2104 2105 2106 2107 2108 2109 2110 2111
        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,
2112
                    "LearningRate": self._create_param_lr(param_and_grad),
2113 2114 2115
                },
                outputs={
                    "ParamOut": param_and_grad[0],
2116
                    "MomentOut": moment_acc,
2117 2118
                },
                attrs={"epsilon": self._epsilon},
2119 2120
                stop_gradient=True,
            )
2121

2122
            return adagrad_op
2123 2124 2125


class AdamOptimizer(Optimizer):
2126
    r"""
T
tianshuo78520a 已提交
2127
    The Adam optimizer uses an optimization described at the end
2128 2129 2130
    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.
2131

2132
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2133 2134 2135 2136 2137 2138 2139 2140 2141 2142 2143 2144 2145 2146

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

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

Q
qiaolongfei 已提交
2149
    Args:
2150 2151
        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.
2152 2153
        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.
2154
            The default value is 0.9.
2155 2156
        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.
2157
            The default value is 0.999.
2158 2159
        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.
2160
            The default value is 1e-08.
H
hong 已提交
2161
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2162
            This parameter is required in dygraph mode. \
2163
            The default value is None in static graph mode, at this time all parameters will be updated.
2164 2165 2166 2167 2168
        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.
2169 2170 2171
        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` ,
2172
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2173 2174 2175 2176 2177 2178 2179 2180 2181 2182
        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.
2183
        use_global_beta_pow (bool, optional): Whether to use global beta_pow. If true, Adam will use global beta_pow
2184
            for whole model instead of creating beta_pow for each parameter. Default is false.
2185 2186
        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
2187
            use same align_size as allocator.
Q
qiaolongfei 已提交
2188 2189 2190 2191

    Examples:
        .. code-block:: python

2192 2193 2194
            import paddle
            import paddle.fluid as fluid

2195
            paddle.enable_static()
2196 2197 2198
            place = fluid.CPUPlace()
            main = fluid.Program()
            with fluid.program_guard(main):
2199 2200
                x = fluid.data(name='x', shape=[None, 13], dtype='float32')
                y = fluid.data(name='y', shape=[None, 1], dtype='float32')
2201
                y_predict = fluid.layers.fc(input=x, size=1, act=None)
2202
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
2203
                avg_cost = paddle.mean(cost)
2204 2205 2206 2207 2208 2209 2210 2211 2212 2213 2214 2215

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

2217 2218 2219 2220 2221 2222 2223
        .. 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

2224
            paddle.enable_static()
2225 2226 2227 2228 2229 2230
            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)
2231
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
2232
                avg_cost = paddle.mean(cost)
2233 2234

                # define beta decay variable
2235
                def get_decayed_betas(beta1_init, beta2_init, decay_steps, decay_rate, epsilon_init):
2236 2237
                    global_step = lr_scheduler._decay_step_counter()

2238
                    beta1 = paddle.static.create_global_var(
2239 2240 2241 2242 2243 2244
                        shape=[1],
                        value=float(beta1_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="beta1")
2245
                    beta2 = paddle.static.create_global_var(
2246 2247 2248 2249 2250 2251
                        shape=[1],
                        value=float(beta2_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="beta2")
2252
                    epsilon = paddle.static.create_global_var(
2253 2254 2255 2256 2257 2258
                        shape=[1],
                        value=float(epsilon_init),
                        dtype='float32',
                        # set persistable for save checkpoints and resume
                        persistable=True,
                        name="epsilon")
2259 2260 2261 2262 2263 2264 2265

                    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)

2266
                    return beta1, beta2, epsilon
2267

2268
                beta1, beta2, epsilon = get_decayed_betas(0.9, 0.99, 1e5, 0.9, 1e-8)
2269 2270
                adam_optimizer = fluid.optimizer.AdamOptimizer(
                                                    learning_rate=0.01,
2271
                                                    beta1=beta1,
2272 2273
                                                    beta2=beta2,
                                                    epsilon=epsilon)
2274 2275 2276 2277 2278 2279 2280 2281 2282 2283
                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)
2284 2285 2286
    """
    _moment1_acc_str = "moment1"
    _moment2_acc_str = "moment2"
Q
qiaolongfei 已提交
2287 2288
    _beta1_pow_acc_str = "beta1_pow_acc"
    _beta2_pow_acc_str = "beta2_pow_acc"
2289

2290 2291 2292 2293 2294 2295 2296 2297 2298 2299 2300 2301 2302 2303 2304
    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,
    ):
2305 2306 2307 2308
        assert learning_rate is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
2309
        super().__init__(
2310 2311 2312 2313 2314 2315 2316 2317
            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,
        )
2318 2319 2320 2321
        self.type = "adam"
        self._beta1 = beta1
        self._beta2 = beta2
        self._epsilon = epsilon
Q
Qiao Longfei 已提交
2322
        self._lazy_mode = lazy_mode
2323
        self._use_global_beta_pow = use_global_beta_pow
2324 2325 2326 2327 2328 2329

    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 已提交
2330 2331
            self._add_accumulator(self._moment1_acc_str, p)
            self._add_accumulator(self._moment2_acc_str, p)
2332 2333 2334 2335
            if not self._use_global_beta_pow:
                self._add_accumulator(
                    name=self._beta1_pow_acc_str,
                    param=p,
2336 2337 2338
                    fill_value=0.9
                    if isinstance(self._beta1, Variable)
                    else self._beta1,
2339
                    shape=[1],
2340 2341 2342
                    type=core.VarDesc.VarType.LOD_TENSOR,
                    device='cpu',
                )
2343 2344 2345
                self._add_accumulator(
                    name=self._beta2_pow_acc_str,
                    param=p,
2346 2347 2348
                    fill_value=0.999
                    if isinstance(self._beta2, Variable)
                    else self._beta2,
2349
                    shape=[1],
2350 2351 2352
                    type=core.VarDesc.VarType.LOD_TENSOR,
                    device='cpu',
                )
2353 2354
        if self._use_global_beta_pow:
            self._add_global_accumulator(
Q
qiaolongfei 已提交
2355
                name=self._beta1_pow_acc_str,
2356 2357 2358
                fill_value=0.9
                if isinstance(self._beta1, Variable)
                else self._beta1,
2359
                shape=[1],
2360 2361 2362
                type=core.VarDesc.VarType.LOD_TENSOR,
                device='cpu',
            )
2363
            self._add_global_accumulator(
Q
qiaolongfei 已提交
2364
                name=self._beta2_pow_acc_str,
2365 2366 2367
                fill_value=0.999
                if isinstance(self._beta2, Variable)
                else self._beta2,
2368
                shape=[1],
2369 2370 2371
                type=core.VarDesc.VarType.LOD_TENSOR,
                device='cpu',
            )
2372 2373 2374 2375

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

2376 2377 2378 2379 2380 2381
        moment1 = self._get_accumulator(
            self._moment1_acc_str, param_and_grad[0]
        )
        moment2 = self._get_accumulator(
            self._moment2_acc_str, param_and_grad[0]
        )
2382 2383
        if self._use_global_beta_pow:
            beta1_pow_acc = self._get_global_accumulator(
2384 2385
                self._beta1_pow_acc_str
            )
2386
            beta2_pow_acc = self._get_global_accumulator(
2387 2388
                self._beta2_pow_acc_str
            )
2389
        else:
2390 2391 2392 2393 2394 2395
            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]
            )
2396
        lr = self._create_param_lr(param_and_grad)
2397
        # create the adam optimize op
2398

姜永久 已提交
2399
        if in_dygraph_mode():
2400 2401 2402 2403 2404 2405 2406 2407 2408 2409
            _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)
            )
2410
            master_weight = None
2411
            _, _, _, _, _, _ = _legacy_C_ops.adam(
2412 2413 2414 2415 2416 2417 2418 2419 2420 2421 2422 2423 2424 2425 2426 2427 2428 2429 2430 2431 2432 2433 2434 2435 2436 2437 2438
                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,
            )
2439 2440 2441

            return None

2442
        inputs = {
2443 2444
            "Param": [param_and_grad[0]],
            "Grad": [param_and_grad[1]],
2445
            "LearningRate": [lr],
2446 2447 2448
            "Moment1": [moment1],
            "Moment2": [moment2],
            "Beta1Pow": [beta1_pow_acc],
2449
            "Beta2Pow": [beta2_pow_acc],
2450
        }
2451 2452 2453 2454 2455 2456 2457

        # 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

2458
        outputs = {
2459 2460 2461 2462 2463
            "ParamOut": [param_and_grad[0]],
            "Moment1Out": [moment1],
            "Moment2Out": [moment2],
            "Beta1PowOut": [beta1_pow_acc],
            "Beta2PowOut": [beta2_pow_acc],
2464 2465 2466
        }
        attrs = {
            "lazy_mode": self._lazy_mode,
2467
            "min_row_size_to_use_multithread": 1000,
2468
            'use_global_beta_pow': self._use_global_beta_pow,
2469 2470 2471 2472 2473 2474 2475 2476 2477 2478
        }

        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
2479 2480 2481 2482
        if isinstance(self._epsilon, Variable):
            inputs['EpsilonTensor'] = self._epsilon
        else:
            attrs['epsilon'] = self._epsilon
2483

2484 2485 2486 2487 2488 2489 2490
        adam_op = block.append_op(
            type=self.type,
            inputs=inputs,
            outputs=outputs,
            attrs=attrs,
            stop_gradient=True,
        )
2491 2492 2493

        return adam_op

2494
    def _finish_update(self, block, parameters_and_grads):
2495
        r"""Update beta1_pow and beta2_pow accumulator"""
2496 2497 2498
        assert isinstance(block, framework.Block)
        if self._use_global_beta_pow:
            beta1_pow_acc = self._get_global_accumulator(
2499 2500
                self._beta1_pow_acc_str
            )
2501
            beta2_pow_acc = self._get_global_accumulator(
2502 2503
                self._beta2_pow_acc_str
            )
2504 2505 2506

            with block.program._optimized_guard([]):
                inputs = {"X": beta1_pow_acc}
2507
                outputs = {"Out": beta1_pow_acc}
2508 2509
                attrs = {}
                if isinstance(self._beta1, Variable):
2510 2511
                    inputs["Y"] = self._beta1
                    # use elementwise_mul for better performance
2512 2513 2514 2515 2516 2517 2518
                    block.append_op(
                        type="elementwise_mul",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
2519 2520
                else:
                    attrs['scale'] = self._beta1
2521 2522 2523 2524 2525 2526 2527
                    block.append_op(
                        type="scale",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
2528 2529

                inputs = {"X": beta2_pow_acc}
2530
                outputs = {"Out": beta2_pow_acc}
2531 2532
                attrs = {}
                if isinstance(self._beta2, Variable):
2533 2534
                    inputs["Y"] = self._beta2
                    # use elementwise_mul for better performance
2535 2536 2537 2538 2539 2540 2541
                    block.append_op(
                        type="elementwise_mul",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
2542 2543
                else:
                    attrs['scale'] = self._beta2
2544 2545 2546 2547 2548 2549 2550
                    block.append_op(
                        type="scale",
                        inputs=inputs,
                        outputs=outputs,
                        attrs=attrs,
                        stop_gradient=True,
                    )
2551

2552 2553

class AdamaxOptimizer(Optimizer):
2554
    r"""
2555
    The Adamax optimizer is implemented based on the Adamax Optimization
2556 2557 2558
    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 已提交
2559

2560
    The parameter ``param_out`` update rule with gradient ``grad``:
Q
qiaolongfei 已提交
2561 2562 2563 2564 2565 2566 2567 2568 2569 2570 2571 2572 2573

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

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

2576 2577 2578 2579 2580 2581 2582 2583 2584 2585 2586 2587
    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 已提交
2588
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2589
            This parameter is required in dygraph mode. \
2590
            The default value is None in static graph mode, at this time all parameters will be updated.
2591 2592 2593 2594 2595
        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.
2596 2597 2598
        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` ,
2599
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2600 2601 2602 2603 2604 2605
        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 已提交
2606

2607 2608 2609 2610 2611
    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          import numpy
2
201716010711 已提交
2612 2613
          import paddle
          paddle.enable_static()
2614 2615 2616 2617 2618 2619 2620 2621

          # 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):
2622
              data = fluid.data(name='X', shape=[None, 1], dtype='float32')
2623
              hidden = fluid.layers.fc(input=data, size=10)
2
201716010711 已提交
2624
              loss = paddle.mean(hidden)
2625
              adam = fluid.optimizer.AdamaxOptimizer(learning_rate=0.2)
2626 2627 2628 2629 2630 2631 2632 2633 2634
              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])
2635 2636 2637
    """
    _moment_acc_str = "moment"
    _inf_norm_acc_str = "inf_norm"
Q
qiaolongfei 已提交
2638
    _beta1_pow_acc_str = "beta1_pow_acc"
2639

2640 2641 2642 2643 2644 2645 2646 2647 2648 2649 2650
    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,
    ):
2651 2652 2653 2654
        assert learning_rate is not None
        assert beta1 is not None
        assert beta2 is not None
        assert epsilon is not None
2655
        super().__init__(
2656 2657 2658 2659 2660 2661
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
2662 2663 2664 2665 2666 2667 2668 2669
        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 已提交
2670 2671
            self._add_accumulator(self._moment_acc_str, p)
            self._add_accumulator(self._inf_norm_acc_str, p)
2672 2673 2674 2675 2676 2677
            self._add_accumulator(
                name=self._beta1_pow_acc_str,
                param=p,
                fill_value=self._beta1,
                shape=[1],
            )
2678 2679 2680 2681 2682

    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])
2683 2684 2685 2686 2687 2688
        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]
        )
2689

姜永久 已提交
2690
        if in_dygraph_mode():
2691 2692 2693 2694 2695 2696 2697 2698 2699 2700 2701
            _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,
            )
2702 2703 2704 2705 2706 2707 2708 2709 2710 2711
        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,
2712
                    "Beta1Pow": beta1_pow_acc,
2713 2714 2715 2716
                },
                outputs={
                    "ParamOut": param_and_grad[0],
                    "MomentOut": moment,
2717
                    "InfNormOut": inf_norm,
2718 2719 2720 2721
                },
                attrs={
                    "beta1": self._beta1,
                    "beta2": self._beta2,
2722
                    "epsilon": self._epsilon,
2723
                },
2724 2725
                stop_gradient=True,
            )
2726

2727
            return adamax_op
2728

2729
    def _finish_update(self, block, parameters_and_grads):
2730
        """Update Beta1 Power accumulator"""
2731
        assert isinstance(block, framework.Block)
2732
        for param, grad in parameters_and_grads:
C
chengduo 已提交
2733
            if grad is None or param.trainable is False:
2734
                continue
2735 2736 2737 2738 2739 2740
            with param.block.program._optimized_guard(
                [param, grad]
            ), name_scope('adamx'):
                beta1_pow_acc = self._get_accumulator(
                    self._beta1_pow_acc_str, param
                )
姜永久 已提交
2741 2742
                if in_dygraph_mode():
                    tmp = _C_ops.scale(beta1_pow_acc, self._beta1, 0.0, True)
2743 2744
                    beta1_pow_acc.copy_(tmp, False)
                else:
2745 2746 2747 2748 2749 2750 2751
                    block.append_op(
                        type="scale",
                        inputs={"X": beta1_pow_acc},
                        outputs={"Out": beta1_pow_acc},
                        attrs={"scale": self._beta1},
                        stop_gradient=True,
                    )
2752 2753


2754
class DpsgdOptimizer(Optimizer):
2755
    r"""
2756 2757 2758 2759 2760 2761 2762 2763
    We implement the Dpsgd optimizer according to CCS16 paper -
    Deep Learning with Differential Privacy.

    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          import numpy
2
201716010711 已提交
2764 2765
          import paddle
          paddle.enable_static()
2766 2767 2768 2769 2770 2771 2772 2773 2774 2775

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

          train_program = fluid.Program()
          startup_program = fluid.Program()
          with fluid.program_guard(train_program, startup_program):
              data = fluid.layers.data(name='X', shape=[1], dtype='float32')
              hidden = fluid.layers.fc(input=data, size=10)
2
201716010711 已提交
2776
              loss = paddle.mean(hidden)
2777 2778 2779 2780 2781 2782 2783 2784 2785 2786 2787 2788 2789 2790 2791 2792 2793
              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 已提交
2794
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2795
            This parameter is required in dygraph mode. \
2796
            The default value is None in static graph mode, at this time all parameters will be updated.
2797 2798 2799 2800
    Notes:
       Currently, DpsgdOptimizer doesn't support sparse parameter optimization.
    """

2801 2802 2803 2804 2805 2806 2807 2808
    def __init__(
        self,
        learning_rate=0.001,
        clip=0.9,
        batch_size=0.999,
        sigma=1e-8,
        parameter_list=None,
    ):
2809 2810 2811 2812
        assert learning_rate is not None
        assert clip is not None
        assert batch_size is not None
        assert sigma is not None
2813
        super().__init__(
2814 2815
            learning_rate=learning_rate, parameter_list=parameter_list
        )
2816 2817 2818 2819
        self.type = "dpsgd"
        self._clip = clip
        self._batch_size = batch_size
        self._sigma = sigma
Z
zhongpu 已提交
2820 2821 2822 2823 2824 2825 2826
        '''
        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
2827 2828 2829 2830 2831

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

        # create the dpsgd optimize op
2832
        if self._seed is None:
Z
zhongpu 已提交
2833 2834
            self._seed = 0

姜永久 已提交
2835
        if in_dygraph_mode():
2836 2837 2838 2839 2840 2841 2842 2843 2844 2845 2846 2847 2848 2849
            _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,
            )
2850
        else:
2851 2852 2853 2854 2855 2856 2857 2858 2859 2860 2861 2862 2863 2864 2865 2866
            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,
            )
2867

2868
            return dpsgd_op
2869 2870


2871
class DecayedAdagradOptimizer(Optimizer):
2872
    r"""
2873 2874 2875
    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.
2876

2877
    The parameter ``param_out`` update rule with gradient ``grad``:
2878 2879 2880 2881 2882 2883 2884

    .. math::

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

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

2885 2886 2887 2888
    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
2889 2890 2891
    stability to avoid the division by zero error.

    Args:
2892 2893 2894 2895 2896
        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 已提交
2897
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
2898
            This parameter is required in dygraph mode. \
2899
            The default value is None in static graph mode, at this time all parameters will be updated.
2900 2901 2902 2903 2904
        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.
2905 2906 2907
        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` ,
2908
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
2909 2910 2911 2912 2913 2914
        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.**
2915 2916 2917 2918

    Examples:
        .. code-block:: python

2919 2920
            import paddle.fluid as fluid

2921 2922 2923 2924
            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)
2925
            optimizer.minimize(cost)
2926 2927 2928
    """
    _moment_acc_str = "moment"

2929 2930 2931 2932 2933 2934 2935 2936 2937 2938
    def __init__(
        self,
        learning_rate,
        decay=0.95,
        epsilon=1.0e-6,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
2939 2940 2941 2942
        assert learning_rate is not None
        assert decay is not None
        assert epsilon is not None

2943
        super().__init__(
2944 2945 2946 2947 2948 2949
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
2950 2951 2952 2953 2954 2955 2956 2957 2958 2959 2960 2961 2962
        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)

2963 2964 2965
        moment_acc = self._get_accumulator(
            self._moment_acc_str, param_and_grad[0]
        )
2966

姜永久 已提交
2967
        if in_dygraph_mode():
2968 2969 2970 2971 2972 2973 2974 2975 2976 2977 2978 2979
            _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,
            )
2980 2981 2982 2983 2984 2985 2986 2987
        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,
2988
                    "LearningRate": self._create_param_lr(param_and_grad),
2989 2990 2991
                },
                outputs={
                    "ParamOut": param_and_grad[0],
2992
                    "MomentOut": moment_acc,
2993
                },
2994 2995 2996
                attrs={"epsilon": self._epsilon, "decay": self._decay},
                stop_gradient=True,
            )
2997

2998
            return decayed_adagrad_op
2999 3000


3001
class AdadeltaOptimizer(Optimizer):
3002
    r"""
Z
Zeng Jinle 已提交
3003
    **Notes: This API does not support sparse parameter optimization.**
Q
qiaolongfei 已提交
3004

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

    The update is done as follows:
3009

Z
Zeng Jinle 已提交
3010 3011
    .. math::

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

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

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

    Args:
Z
Zeng Jinle 已提交
3019 3020 3021
        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 已提交
3022
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3023
            This parameter is required in dygraph mode. \
3024
            The default value is None in static graph mode, at this time all parameters will be updated.
3025 3026 3027 3028 3029
        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.
3030 3031 3032
        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` ,
3033
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3034 3035 3036
        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` .
3037 3038 3039 3040

    Examples:
        .. code-block:: python

3041
            import paddle.fluid as fluid
Z
Zeng Jinle 已提交
3042

3043
            image = fluid.data(name='image', shape=[None, 28], dtype='float32')
Z
Zeng Jinle 已提交
3044 3045
            fc = fluid.layers.fc(image, size=10)
            cost = fluid.layers.reduce_mean(fc)
3046 3047
            optimizer = fluid.optimizer.Adadelta(
                learning_rate=0.0003, epsilon=1.0e-6, rho=0.95)
C
chengduo 已提交
3048

Z
Zeng Jinle 已提交
3049 3050 3051 3052
            # 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)
3053
    """
3054

3055 3056 3057
    _avg_squared_grad_acc_str = "_avg_squared_grad"
    _avg_squared_update_acc_str = "_avg_squared_update"

3058 3059 3060 3061 3062 3063 3064 3065 3066 3067
    def __init__(
        self,
        learning_rate,
        epsilon=1.0e-6,
        rho=0.95,
        parameter_list=None,
        regularization=None,
        grad_clip=None,
        name=None,
    ):
3068 3069 3070 3071 3072 3073
        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.")
3074
        super().__init__(
3075 3076 3077 3078 3079 3080
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
3081 3082 3083 3084 3085
        self.type = "adadelta"
        self._epsilon = epsilon
        self._rho = rho

    def _create_accumulators(self, block, parameters):
3086 3087
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")
3088 3089 3090 3091 3092 3093

        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):
3094 3095
        if not isinstance(block, framework.Block):
            raise TypeError("block is not instance of framework.Block.")
3096 3097

        avg_squared_grad_acc = self._get_accumulator(
3098 3099
            self._avg_squared_grad_acc_str, param_and_grad[0]
        )
3100
        avg_squared_update_acc = self._get_accumulator(
3101 3102
            self._avg_squared_update_acc_str, param_and_grad[0]
        )
3103

姜永久 已提交
3104
        if in_dygraph_mode():
3105 3106 3107 3108 3109 3110 3111 3112
            _C_ops.adadelta_(
                param_and_grad[0],
                param_and_grad[1],
                avg_squared_grad_acc,
                avg_squared_update_acc,
                self._rho,
                self._epsilon,
            )
3113 3114
        else:
            # Create the adadelta optimizer op
3115 3116 3117 3118 3119 3120 3121 3122 3123 3124 3125 3126 3127 3128 3129 3130
            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,
            )
3131

3132
            return adadelta_op
3133 3134


Q
qingqing01 已提交
3135
class RMSPropOptimizer(Optimizer):
3136
    r"""
Q
qingqing01 已提交
3137 3138 3139 3140 3141 3142 3143 3144
    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 已提交
3145
        r(w, t) & = \\rho r(w, t-1) + (1 - \\rho)(\\nabla Q_{i}(w))^2
Q
qingqing01 已提交
3146 3147 3148 3149

        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 已提交
3150
    each weight. Then dividing the gradient by :math:`sqrt{v(w,t)}`.
Q
qingqing01 已提交
3151 3152 3153 3154 3155 3156

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

    ..  math::

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

3159 3160 3161 3162 3163 3164 3165 3166 3167 3168 3169 3170 3171 3172
        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 已提交
3173 3174 3175 3176
            \\epsilon}} \\nabla Q_{i}(w)

        w & = w - v(w, t)

Q
qiaolongfei 已提交
3177
    where, :math:`\\rho` is a hyperparameter and typical values are 0.9, 0.95
Q
qingqing01 已提交
3178 3179 3180 3181 3182
    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.


3183 3184 3185
    Parameters:
        learning_rate(float): Global learning rate.
        rho(float): rho is :math: `\\rho` in equation, default is 0.95.
Q
qingqing01 已提交
3186
        epsilon(float): :math: `\\epsilon` in equation is smoothing term to
3187
            avoid division by zero, default is 1e-6.
Q
qiaolongfei 已提交
3188
        momentum(float): :math:`\\beta` in equation is the momentum term,
3189
            default is 0.0.
3190 3191 3192 3193
        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 已提交
3194
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3195
            This parameter is required in dygraph mode. \
3196
            The default value is None in static graph mode, at this time all parameters will be updated.
3197 3198 3199 3200 3201
        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.
3202 3203 3204
        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` ,
3205
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3206 3207
        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 已提交
3208 3209 3210 3211 3212 3213 3214

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

    Examples:
          .. code-block:: python

3215 3216 3217 3218
            import paddle
            import paddle.fluid as fluid
            import numpy as np

3219
            paddle.enable_static()
3220 3221 3222 3223 3224 3225
            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)
3226
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
3227
                avg_cost = paddle.mean(cost)
3228 3229 3230 3231 3232 3233 3234 3235 3236 3237 3238 3239 3240

                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 已提交
3241 3242 3243 3244
    """

    _momentum_acc_str = "momentum"
    _mean_square_acc_str = "mean_square"
3245
    _mean_grad_acc_str = "mean_grad"
Q
qingqing01 已提交
3246

3247 3248 3249 3250 3251 3252 3253 3254 3255 3256 3257 3258
    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,
    ):
3259
        super().__init__(
3260 3261 3262 3263 3264 3265
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
Q
qingqing01 已提交
3266 3267 3268 3269 3270 3271 3272 3273 3274 3275 3276 3277 3278
        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
3279
        self._centered = centered
Q
qingqing01 已提交
3280 3281 3282 3283 3284 3285 3286 3287

    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)
3288
            self._add_accumulator(self._mean_grad_acc_str, p)
Q
qingqing01 已提交
3289 3290 3291 3292 3293

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

3294 3295 3296 3297 3298 3299 3300 3301 3302
        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 已提交
3303
        if in_dygraph_mode():
3304 3305 3306 3307 3308 3309 3310 3311 3312 3313 3314 3315
            _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 已提交
3316
            return None
3317 3318 3319 3320 3321 3322 3323 3324 3325 3326 3327 3328 3329 3330 3331
        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,
3332
                    "MeanGradOut": mean_grad_acc,
3333 3334 3335 3336 3337
                },
                attrs={
                    "epsilon": self._epsilon,
                    "decay": self._rho,
                    "momentum": self._momentum,
3338
                    "centered": self._centered,
3339
                },
3340 3341
                stop_gradient=True,
            )
Q
qingqing01 已提交
3342

3343
            return rmsprop_op
Q
qingqing01 已提交
3344 3345


Q
qiaolongfei 已提交
3346
class FtrlOptimizer(Optimizer):
3347
    r"""
Q
qiaolongfei 已提交
3348 3349 3350 3351 3352 3353 3354 3355 3356 3357 3358 3359 3360 3361 3362 3363 3364 3365 3366 3367 3368 3369 3370 3371 3372 3373 3374 3375 3376 3377 3378 3379 3380 3381 3382 3383 3384 3385
    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

3386 3387 3388 3389 3390
    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 已提交
3391
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3392
            This parameter is required in dygraph mode. \
3393
            The default value is None in static graph mode, at this time all parameters will be updated.
3394 3395 3396 3397 3398
        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.
3399 3400 3401
        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` ,
3402
            :ref:`api_fluid_clip_GradientClipByValue` ). Default None, meaning there is no gradient clipping.
3403 3404
        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 已提交
3405 3406 3407 3408 3409 3410 3411

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

    Examples:
          .. code-block:: python

3412 3413 3414 3415
            import paddle
            import paddle.fluid as fluid
            import numpy as np

3416 3417
            paddle.enable_static()

3418 3419 3420 3421 3422 3423
            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)
3424
                cost = paddle.nn.functional.square_error_cost(input=y_predict, label=y)
2
201716010711 已提交
3425
                avg_cost = paddle.mean(cost)
3426 3427 3428 3429 3430 3431 3432 3433 3434 3435 3436 3437

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

3439
    NOTE:
C
chengduo 已提交
3440
       Currently, FtrlOptimizer doesn't support sparse parameter optimization.
Q
qiaolongfei 已提交
3441 3442 3443 3444 3445
    """

    _squared_acc_str = "squared"
    _linear_acc_str = "linear"

3446 3447 3448 3449 3450 3451 3452 3453 3454 3455 3456
    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,
    ):
3457
        super().__init__(
3458 3459 3460 3461 3462 3463
            learning_rate=learning_rate,
            parameter_list=parameter_list,
            regularization=regularization,
            grad_clip=grad_clip,
            name=name,
        )
Q
qiaolongfei 已提交
3464 3465 3466 3467 3468 3469 3470 3471 3472 3473 3474 3475 3476 3477 3478 3479 3480 3481 3482 3483
        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.")

3484 3485 3486 3487 3488 3489
        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]
        )
姜永久 已提交
3490
        if in_dygraph_mode():
3491 3492 3493 3494 3495 3496 3497 3498 3499 3500 3501 3502 3503 3504 3505 3506
            _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,
            )
3507 3508

        else:
3509 3510 3511 3512 3513 3514 3515 3516 3517 3518 3519 3520 3521 3522 3523 3524 3525 3526 3527 3528 3529
            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 已提交
3530

3531
            return ftrl_op
Q
qiaolongfei 已提交
3532 3533


Y
Yibing Liu 已提交
3534
class LambOptimizer(AdamOptimizer):
3535
    r"""
Y
Yibing Liu 已提交
3536 3537
    LAMB (Layer-wise Adaptive Moments optimizer for Batching training) Optimizer.

3538 3539 3540
    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 已提交
3541
    Deep Learning: Training BERT in 76 minutes <https://arxiv.org/abs/1904.00962>`_ .
Y
Yibing Liu 已提交
3542 3543 3544 3545 3546

    The updating of parameters follows:

    ..  math::

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

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

3551 3552 3553 3554
        m_t &= \\frac{m_t}{\\beta_1^t}

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

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

Y
Yibing Liu 已提交
3557
        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 已提交
3558 3559


3560
    where :math:`m` is the 1st moment, and :math:`v` the 2nd moment, :math:`\\eta` the
Y
Yibing Liu 已提交
3561 3562 3563
    learning rate, :math:`\\lambda` the LAMB weight decay rate.

    Args:
Y
Yibing Liu 已提交
3564 3565 3566 3567 3568 3569 3570 3571
        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 已提交
3572
        parameter_list (Iterable, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
3573
            This parameter is required in dygraph mode. \
3574
            The default value is None in static graph mode, at this time all parameters will be updated.
3575 3576 3577 3578 3579
        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.
3580 3581
        grad_clip (GradientClipBase, optional): Gradient cliping strategy, it's an instance of
            some derived class of ``GradientClipBase`` . There are three cliping strategies
3582 3583 3584
            ( :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.
3585 3586
        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 已提交
3587
            Default None.
3588
        name(str|None): For detailed information, please refer to
Y
Yibing Liu 已提交
3589
            :ref:`api_guide_Name` . Usually name is no need to set and None by default.
Y
Yibing Liu 已提交
3590 3591 3592

    Examples:
        .. code-block:: python
3593

2
201716010711 已提交
3594
            import paddle
3595
            import paddle.fluid as fluid
2
201716010711 已提交
3596
            paddle.enable_static()
Y
Yibing Liu 已提交
3597

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

Y
Yibing Liu 已提交
3602 3603 3604 3605 3606
            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 已提交
3607 3608 3609 3610 3611 3612 3613
            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"

3614 3615 3616 3617 3618 3619 3620 3621 3622 3623 3624 3625 3626
    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 已提交
3627 3628 3629 3630 3631
        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
3632
        super().__init__(
3633 3634 3635 3636 3637 3638 3639 3640 3641
            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 已提交
3642 3643
        self.type = "lamb"
        self._weight_decay = lamb_weight_decay
Y
Yibing Liu 已提交
3644
        self._exclude_from_weight_decay_fn = exclude_from_weight_decay_fn
Y
Yibing Liu 已提交
3645 3646 3647

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

3650 3651 3652 3653 3654 3655 3656 3657 3658 3659 3660 3661 3662 3663 3664 3665 3666
        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 已提交
3667 3668 3669
            weight_decay = 0.0
        else:
            weight_decay = self._weight_decay
3670
        lr = self._create_param_lr(param_and_grad)
3671
        master_weight = None
姜永久 已提交
3672
        if in_dygraph_mode():
3673 3674 3675 3676 3677 3678 3679 3680 3681 3682 3683 3684 3685 3686 3687 3688 3689 3690 3691 3692 3693 3694 3695 3696
            _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,
            )
3697
            return None
Y
Yibing Liu 已提交
3698

Y
Yibing Liu 已提交
3699
        # create the lamb optimize op
3700 3701 3702 3703 3704 3705 3706 3707 3708 3709 3710 3711 3712 3713 3714 3715 3716 3717 3718 3719 3720 3721 3722 3723 3724 3725
        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 已提交
3726 3727 3728 3729

        return lamb_op


3730 3731 3732 3733 3734 3735 3736 3737 3738 3739 3740 3741 3742
# 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
3743
Dpsgd = DpsgdOptimizer
3744
DecayedAdagrad = DecayedAdagradOptimizer
3745
Adadelta = AdadeltaOptimizer
Q
qingqing01 已提交
3746
RMSProp = RMSPropOptimizer
Q
qiaolongfei 已提交
3747
Ftrl = FtrlOptimizer
3748
LarsMomentum = LarsMomentumOptimizer
Y
Yibing Liu 已提交
3749
Lamb = LambOptimizer
3750 3751 3752


class ModelAverage(Optimizer):
3753
    r"""
3754
	:api_attr: Static Graph
S
swtkiwi 已提交
3755

3756 3757 3758 3759 3760 3761 3762 3763 3764 3765 3766 3767 3768 3769 3770 3771 3772 3773
    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:

    ::
3774

3775 3776 3777 3778 3779 3780 3781 3782 3783
        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.
3784 3785

    Args:
3786 3787 3788
        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.
3789 3790 3791 3792 3793
        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.
3794 3795 3796
        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.
3797

3798
    Examples:
Q
qiaolongfei 已提交
3799 3800 3801

      .. code-block:: python

2
201716010711 已提交
3802
        import paddle
3803 3804
        import paddle.fluid as fluid
        import numpy
2
201716010711 已提交
3805
        paddle.enable_static()
3806 3807 3808 3809

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

3811 3812 3813 3814
        train_program = fluid.Program()
        startup_program = fluid.Program()
        with fluid.program_guard(train_program, startup_program):
            # build net
3815
            data = fluid.data(name='X', shape=[None, 1], dtype='float32')
3816
            hidden = fluid.layers.fc(input=data, size=10)
2
201716010711 已提交
3817
            loss = paddle.mean(hidden)
3818 3819 3820 3821 3822 3823
            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,
3824
                                                         max_average_window=12500)
3825 3826

            exe.run(startup_program)
3827 3828 3829 3830 3831
            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])
3832 3833

            # apply ModelAverage
3834
            with model_average.apply(exe):
3835 3836 3837 3838
                x = numpy.random.random(size=(10, 1)).astype('float32')
                exe.run(program=train_program,
                        feed={'X': x},
                        fetch_list=[loss.name])
3839 3840
    """

3841 3842 3843 3844 3845 3846 3847 3848
    def __init__(
        self,
        average_window_rate,
        min_average_window=10000,
        max_average_window=10000,
        regularization=None,
        name=None,
    ):
姜永久 已提交
3849
        if in_dygraph_mode():
Z
zhongpu 已提交
3850
            raise Exception("In dygraph, don't support ModelAverage.")
3851
        super().__init__(0.0, regularization=regularization, name=name)
3852 3853 3854
        self.average_window = average_window_rate
        self.min_average_window = min_average_window
        self.max_average_window = max_average_window
3855

3856
        self.params_grads = []
3857 3858 3859
        for param in (
            framework.default_main_program().global_block().all_parameters()
        ):
3860
            if param.do_model_average != False:
3861
                grad = param.block.create_var(
3862 3863 3864
                    name=unique_name.generate_with_ignorable_key(
                        ".".join([param.name, 'tmp'])
                    ),
3865 3866
                    dtype=param.dtype,
                    persistable=False,
3867 3868
                    stop_gradient=True,
                )
3869
                self.params_grads.append((param, grad))
3870

3871
        for param, grad in self.params_grads:
3872 3873
            if grad is None:
                continue
X
Xin Pan 已提交
3874
            with param.block.program._optimized_guard(
3875 3876
                [param, grad]
            ), name_scope('move_average'):
3877
                self._append_average_accumulate_op(param)
3878

3879 3880 3881 3882
        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:
3883
                self._add_average_apply_op(block, param_grad)
3884 3885 3886 3887 3888

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

3891
    def _add_average_apply_op(self, block, param_grad):
L
Luo Tao 已提交
3892 3893 3894 3895 3896 3897
        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(
3898 3899
            self._get_accumulator('num_accumulates', param)
        )
L
Luo Tao 已提交
3900
        old_num_accumulates = block._clone_variable(
3901 3902
            self._get_accumulator('old_num_accumulates', param)
        )
L
Luo Tao 已提交
3903
        num_updates = block._clone_variable(
3904 3905
            self._get_accumulator('num_updates', param)
        )
3906 3907 3908
        # backup param value to grad
        layers.assign(input=param, output=grad)
        # param = (sum_1 + sum_2 + sum_3) / (num_accumulates + old_num_accumulates)
3909 3910
        tmp = paddle.add_n([num_accumulates, old_num_accumulates])
        sum = paddle.add_n([sum_1, sum_2, sum_3])
D
dzhwinter 已提交
3911
        tmp = layers.cast(
3912
            x=tmp, dtype='float32' if self._dtype is None else self._dtype
3913
        )
D
dzhwinter 已提交
3914
        sum = layers.cast(
3915
            x=sum, dtype='float32' if self._dtype is None else self._dtype
3916
        )
3917
        paddle.assign(paddle.divide(sum, tmp), output=param)
3918 3919

    def _add_average_restore_op(self, block, param_grad):
L
Luo Tao 已提交
3920 3921
        param = block._clone_variable(param_grad[0])
        grad = block._clone_variable(param_grad[1])
3922 3923 3924 3925 3926 3927 3928
        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)
3929 3930 3931 3932 3933 3934 3935 3936 3937 3938 3939 3940 3941 3942 3943 3944 3945 3946 3947 3948 3949 3950 3951 3952 3953 3954 3955 3956 3957 3958 3959 3960 3961 3962 3963 3964
        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,
        )
3965

S
rename  
sneaxiy 已提交
3966
    @signature_safe_contextmanager
3967
    def apply(self, executor, need_restore=True):
3968 3969
        """
        Apply the average of the cumulative ``Parameter`` to the parameters of the current model.
3970 3971

        Args:
3972 3973 3974 3975 3976 3977 3978 3979 3980 3981 3982
            executor(fluid.Executor): The current network executor.
            need_restore(bool): Restore flag variable, if set to True, the network will restore
                the parameters of the network to the default value, if set to False,
                it will not be restored. The default value is True.

        Examples:

          .. code-block:: python

            import paddle.fluid as fluid
            import numpy
2
201716010711 已提交
3983 3984
            import paddle
            paddle.enable_static()
3985 3986 3987 3988 3989 3990 3991 3992 3993 3994 3995

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

            train_program = fluid.Program()
            startup_program = fluid.Program()
            with fluid.program_guard(train_program, startup_program):
                # build net
                data = fluid.data(name='X', shape=[None, 1], dtype='float32')
                hidden = fluid.layers.fc(input=data, size=10)
2
201716010711 已提交
3996
                loss = paddle.mean(hidden)
3997 3998 3999 4000 4001 4002 4003 4004 4005 4006 4007 4008 4009 4010 4011 4012 4013 4014 4015 4016 4017
                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])
4018
        """
4019 4020 4021 4022 4023 4024
        executor.run(self.apply_program)
        try:
            yield
        finally:
            if need_restore:
                self.restore(executor)
4025 4026

    def restore(self, executor):
4027 4028
        """
        Restore ``Parameter`` values of current model.
4029

4030
        Args:
4031 4032 4033 4034 4035 4036 4037 4038
            executor(fluid.Executor): The current network executor.

        Examples:

          .. code-block:: python

            import paddle.fluid as fluid
            import numpy
2
201716010711 已提交
4039 4040
            import paddle
            paddle.enable_static()
4041 4042 4043 4044 4045 4046 4047 4048 4049 4050 4051

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

            train_program = fluid.Program()
            startup_program = fluid.Program()
            with fluid.program_guard(train_program, startup_program):
                # build net
                data = fluid.data(name='X', shape=[None, 1], dtype='float32')
                hidden = fluid.layers.fc(input=data, size=10)
2
201716010711 已提交
4052
                loss = paddle.mean(hidden)
4053 4054 4055 4056 4057 4058 4059 4060 4061 4062 4063 4064 4065 4066 4067 4068 4069 4070 4071 4072 4073 4074 4075 4076
                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)
4077
        """
4078
        executor.run(self.restore_program)
4079 4080


4081
class ExponentialMovingAverage:
4082
    r"""
S
swtkiwi 已提交
4083

4084 4085 4086 4087 4088 4089
    Compute the moving average of parameters with exponential decay.
    Given a parameter :math:`\\theta`, its exponential moving average (EMA)
    will be

    ..  math::

4090
        \text{EMA}_0 & = 0
4091

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

4094 4095 4096
    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 已提交
4097
    the **restore()** method is used to restore the parameters.
4098

4099 4100
    **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
4101
    :math:`(1 - \text{decay}^t)` , i.e., the actual EMAs applied to parameters
4102
    when calling **apply()** method would be
4103 4104

    ..  math::
4105

4106
        \widehat{\text{EMA}}_t = \frac{\text{EMA}_t}{1 - \text{decay}^t}
4107

4108 4109
    **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
4110
    relative smaller decay rate in the very beginning. The argument **thres_steps**
4111
    allows users to pass a Variable to schedule the decay rate, in this case,
4112
    the actual decay rate becomes
4113

4114
    ..  math::
4115

4116
        \min(\text{decay}, \frac{1 + \text{thres_steps}}{10 + \text{thres_steps}})
4117 4118

    Usually **thres_steps** can be the global training steps.
4119 4120 4121


    Args:
4122 4123 4124
        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.
4125 4126 4127 4128


    Examples:

4129 4130 4131 4132 4133 4134 4135 4136 4137 4138 4139 4140 4141 4142 4143 4144 4145 4146 4147 4148 4149 4150 4151 4152 4153 4154 4155 4156
        .. 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(),
4157
                    feed={'x': data},
4158 4159 4160 4161 4162 4163
                    fetch_list=[cost.name])

                # usage 1
                with ema.apply(exe):
                    data = numpy.random.random(size=(10, 5)).astype('float32')
                    exe.run(program=test_program,
4164
                        feed={'x': data},
4165 4166 4167 4168 4169 4170
                        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,
4171
                        feed={'x': data},
4172 4173 4174
                        fetch_list=[hidden.name])
                ema.restore(exe)

4175 4176
    """

4177
    def __init__(self, decay=0.999, thres_steps=None, name=None):
姜永久 已提交
4178
        if in_dygraph_mode():
Z
zhongpu 已提交
4179
            raise Exception(
4180 4181
                "In dygraph, don't support ExponentialMovingAverage."
            )
4182
        self._decay = decay
4183
        self._thres_steps = thres_steps
4184
        self._name = name if name is not None else ''
4185 4186
        self._decay_var = self._get_ema_decay()

4187
        self._step_counter_name = "@EMA_STEP_COUNTER@"
Y
Yibing Liu 已提交
4188
        self._params_tmps = []
4189
        for param in default_main_program().global_block().all_parameters():
4190
            if param.do_model_average != False:
4191 4192 4193 4194 4195 4196 4197 4198
                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 已提交
4199
                self._params_tmps.append((param, tmp))
4200

Y
Yibing Liu 已提交
4201 4202
        self._ema_vars = {}
        for param, tmp in self._params_tmps:
4203 4204 4205
            with param.block.program._optimized_guard([param, tmp]), name_scope(
                'moving_average'
            ):
Y
Yibing Liu 已提交
4206
                self._ema_vars[param.name] = self._create_ema_vars(param)
4207 4208 4209 4210

        self.apply_program = Program()
        block = self.apply_program.global_block()
        with program_guard(main_program=self.apply_program):
4211
            decay_pow, global_step = self._get_decay_pow(block)
Y
Yibing Liu 已提交
4212
            for param, tmp in self._params_tmps:
4213 4214
                param = block._clone_variable(param)
                tmp = block._clone_variable(tmp)
Y
Yibing Liu 已提交
4215
                ema = block._clone_variable(self._ema_vars[param.name])
4216
                layers.assign(input=param, output=tmp)
4217
                # bias correction
4218 4219
                with layers.control_flow.Switch() as switch:
                    with switch.case(global_step > 0):
4220 4221 4222
                        layers.assign(
                            output=param, input=ema / (1.0 - decay_pow)
                        )
4223 4224
                    with switch.default():
                        layers.assign(output=param, input=ema)
4225 4226 4227 4228

        self.restore_program = Program()
        block = self.restore_program.global_block()
        with program_guard(main_program=self.restore_program):
Y
Yibing Liu 已提交
4229
            for param, tmp in self._params_tmps:
4230 4231 4232 4233
                tmp = block._clone_variable(tmp)
                param = block._clone_variable(param)
                layers.assign(input=tmp, output=param)

4234 4235
    def _get_ema_decay(self):
        with default_main_program()._lr_schedule_guard():
4236
            decay_var = paddle.static.create_global_var(
4237 4238 4239 4240
                shape=[1],
                value=self._decay,
                dtype='float32',
                persistable=True,
4241 4242
                name="scheduled_ema_decay_rate",
            )
4243 4244 4245 4246 4247 4248 4249 4250

            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(
4251 4252
                            np.array([self._decay], dtype=np.float32), decay_var
                        )
4253 4254 4255
        return decay_var

    def _get_decay_pow(self, block):
4256
        global_step = paddle.static.create_global_var(
4257 4258 4259 4260 4261 4262
            name=self._step_counter_name,
            shape=[1],
            value=0,
            dtype='int64',
            persistable=True,
        )
4263
        global_step = layers.cast(global_step, "float32")
4264
        decay_var = block._clone_variable(self._decay_var)
4265
        decay_pow_acc = paddle.pow(decay_var, global_step)
4266
        return decay_pow_acc, global_step
4267

Y
Yibing Liu 已提交
4268
    def _create_ema_vars(self, param):
4269
        param_ema = paddle.static.create_global_var(
4270 4271 4272 4273
            name=unique_name.generate(self._name + param.name + '_ema'),
            shape=param.shape,
            value=0.0,
            dtype=param.dtype,
4274 4275
            persistable=True,
        )
4276 4277 4278

        return param_ema

Y
Yibing Liu 已提交
4279
    def update(self):
4280 4281
        """
        Update Exponential Moving Average. Should only call this method in
Y
Yibing Liu 已提交
4282 4283
        train program.
        """
4284
        global_step = layers.autoincreased_step_counter(
4285 4286
            counter_name=self._step_counter_name
        )
4287
        param_master_emas = []
Y
Yibing Liu 已提交
4288
        for param, tmp in self._params_tmps:
4289 4290 4291
            with param.block.program._optimized_guard([param, tmp]), name_scope(
                'moving_average'
            ):
Y
Yibing Liu 已提交
4292
                param_ema = self._ema_vars[param.name]
4293
                if param.name + '.master' in self._ema_vars:
4294 4295 4296 4297
                    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 * (
4298 4299
                        1 - self._decay_var
                    )
4300 4301 4302 4303 4304 4305 4306 4307 4308 4309
                    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,
4310 4311 4312
                    "out_dtype": param_ema.dtype,
                },
            )
Y
Yibing Liu 已提交
4313

4314 4315 4316 4317
    @signature_safe_contextmanager
    def apply(self, executor, need_restore=True):
        """
        Apply moving average to parameters for evaluation.
4318

4319 4320
        Args:
            executor (Executor): The Executor to execute applying.
4321
            need_restore (bool, optional): Whether to restore parameters after
Y
Yibing Liu 已提交
4322
                applying. Default True.
4323 4324 4325 4326 4327 4328 4329 4330 4331 4332
        """
        executor.run(self.apply_program)
        try:
            yield
        finally:
            if need_restore:
                self.restore(executor)

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

4334 4335 4336 4337
        Args:
            executor (Executor): The Executor to execute restoring.
        """
        executor.run(self.restore_program)
H
hutuxian 已提交
4338 4339


4340
class PipelineOptimizer:
4341
    """
4342
        :api_attr: Static Graph
S
swtkiwi 已提交
4343

4344 4345 4346 4347
    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 已提交
4348

4349
    Args:
4350 4351 4352
        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].
4353

4354 4355
    Examples:
        .. code-block:: python
H
hutuxian 已提交
4356

4357
            import paddle.fluid as fluid
H
hutuxian 已提交
4358 4359
            import paddle.fluid.layers as layers

4360 4361 4362 4363 4364 4365 4366 4367 4368 4369 4370 4371 4372 4373 4374 4375
            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 已提交
4376
            optimizer = fluid.optimizer.SGD(learning_rate=0.5)
4377
            optimizer = fluid.optimizer.PipelineOptimizer(optimizer)
H
hutuxian 已提交
4378
            optimizer.minimize(loss)
4379 4380 4381 4382 4383 4384 4385 4386 4387

            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 已提交
4388 4389
            exe = fluid.Executor(place)
            exe.run(fluid.default_startup_program())
4390 4391
            batch_size = 1
            data_loader.start()
H
hutuxian 已提交
4392
            exe.train_from_dataset(
4393
                    fluid.default_main_program())
4394
            data_loader.reset()
4395 4396
    """

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

        # 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

4425 4426 4427
        assert (
            num_microbatches >= 1
        ), "num_microbatches must be a positive value."
4428
        self._num_microbatches = num_microbatches
4429 4430 4431
        assert (
            start_cpu_core_id >= 0
        ), "start_cpu_core_id must be a non-negative integer."
H
hutuxian 已提交
4432
        self._start_cpu_core_id = start_cpu_core_id
4433 4434 4435 4436 4437 4438
        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()
4439
        self._param_device_map = None
4440 4441
        self._pipeline_pair = []
        self._pp_ring_map = dict()
4442 4443
        self.output_var_to_op = None
        self.input_var_to_op = None
4444 4445 4446 4447 4448 4449 4450 4451 4452 4453 4454 4455 4456 4457 4458

    # 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")
4459 4460 4461 4462 4463 4464 4465 4466 4467 4468 4469 4470 4471 4472
            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,
                },
            )
4473 4474 4475 4476
            offset += 1
        block._insert_op(
            op_idx + 1 + offset,
            type='c_allreduce_max'
4477 4478
            if op.type == "reduce_any"
            else 'c_allreduce_sum',
4479 4480 4481
            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={
4482
                'ring_id': self.global_ring_id,
4483
                self._op_role_key: self._op_role.Optimize,
4484 4485 4486
                'use_calc_stream': True,
            },
        )
4487 4488
        offset += 1
        if op.type == "reduce_any":
4489 4490 4491 4492 4493 4494 4495 4496 4497 4498 4499
            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,
                },
            )
4500
            offset += 1
4501
        return offset
H
hutuxian 已提交
4502

4503
    def _create_vars(self, block, ori_block):
4504
        # Create vars for block, copied from ori_block
H
hutuxian 已提交
4505
        used_var_set = set()
4506 4507 4508 4509 4510 4511 4512 4513 4514
        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]
4515
            # For op process vars on all devices, remove its input
4516 4517 4518 4519 4520 4521 4522 4523 4524 4525 4526 4527 4528 4529 4530
            # 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)
4531 4532 4533 4534 4535 4536 4537 4538 4539 4540
            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
4541 4542 4543 4544 4545 4546 4547 4548
            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 已提交
4549
            for var in vars:
4550 4551
                # a var whose name contains "blocking_queue"
                # only exists in startup program
4552
                if var in used_var_set or "_blocking_queue" in var:
H
hutuxian 已提交
4553 4554
                    continue
                used_var_set.add(var)
4555 4556
                if block._find_var_recursive(str(var)):
                    continue
4557
                source_var = ori_block._var_recursive(str(var))
4558
                if source_var.type == core.VarDesc.VarType.READER:
4559
                    dest_var = block.create_var(
4560 4561
                        name=var,
                        type=core.VarDesc.VarType.READER,
4562 4563
                        persistable=source_var.persistable,
                    )
4564 4565 4566 4567 4568 4569 4570 4571 4572 4573 4574
                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,
4575 4576
                        error_clip=source_var.error_clip,
                    )
4577
                else:
4578
                    dest_var = block._clone_variable(source_var, False)
4579
                self._clone_var_attr(dest_var, source_var)
4580 4581 4582
            # When use with sharding, allreduce_sum and allreduce_max
            # used for global gradient clip and amp will be added by sharding.
            op_idx += 1
4583 4584
            if self.use_sharding or not should_insert:
                continue
4585 4586 4587 4588
            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 已提交
4589

4590
    def _is_loss_grad_op(self, op):
4591 4592
        assert self._op_role_key in op.attr_names
        op_role = int(op.attr(self._op_role_key))
4593
        return op_role & int(self._op_role.Backward) and op_role & int(
4594 4595
            self._op_role.Loss
        )
4596

4597
    def _is_forward_op(self, op):
4598 4599 4600
        return self._op_role_key in op.attr_names and (
            int(op.attr(self._op_role_key)) == int(self._op_role.Forward)
        )
4601

4602
    def _is_backward_op(self, op):
4603
        return self._op_role_key in op.attr_names and (
4604 4605
            int(op.attr(self._op_role_key)) & int(self._op_role.Backward)
        )
4606 4607 4608 4609

    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)
4610 4611

    def _is_optimize_op(self, op):
4612
        return self._op_role_key in op.attr_names and (
4613 4614
            int(op.attr(self._op_role_key)) & int(self._op_role.Optimize)
        )
4615 4616

    def _is_update_op(self, op):
4617 4618 4619 4620 4621
        return (
            'Param' in op.input_names
            and 'Grad' in op.input_names
            and ("LearningRate" in op.input_names)
        )
4622

4623
    def _split_program(self, main_program, devices):
H
hutuxian 已提交
4624
        """
4625
        Split a program into sections according to devices that ops run on.
4626
        The op whose op_device attr is "gpu:all" is copied to all sections.
4627 4628 4629

        Args:
            main_program (Program): the main program
4630
            devices: all used devices
H
hutuxian 已提交
4631
        """
4632
        # Map from device to its corresponding section program info
4633
        device_program_map = defaultdict(Program)
4634

4635
        block = main_program.block(0)
4636 4637
        for op in block.ops:
            device = op.attr(self._op_device_key)
4638
            # Copy ops whose op_device set to "gpu:all" to all sections.
4639
            if device == f"{self._device}:all":
4640
                for device in devices:
4641 4642
                    program = device_program_map[device]
                    op_desc = op.desc
4643
                    ap_op = program.global_block().desc.append_op()
4644
                    ap_op.copy_from(op_desc)
4645
                    ap_op._set_attr(self._op_device_key, "")
4646 4647 4648
            else:
                program = device_program_map[device]
                op_desc = op.desc
4649
                ap_op = program.global_block().desc.append_op()
4650
                ap_op.copy_from(op_desc)
4651
                ap_op._set_attr(self._op_device_key, "")
4652

4653
        program_list = []
4654
        for key in devices:
4655
            program = device_program_map[key]
4656 4657
            program._sync_with_cpp()
            program_list.append(program)
H
hutuxian 已提交
4658

4659
        return program_list
H
hutuxian 已提交
4660

4661 4662 4663 4664 4665 4666 4667
    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.
        """
4668 4669
        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 '
4670
            'or beta2_pow_acc.'
4671 4672
        )
        param_name = var_name[0 : var_name.index('_beta')]
4673 4674 4675
        device = self._param_device_map[param_name]
        return device

4676 4677
    def _split_startup_program(self, startup_program, device_id):
        block = startup_program.global_block()
4678 4679 4680
        new_startup_program = Program()
        for op in block.ops:
            device = op.attr(self._op_device_key)
4681 4682
            if device == "cpu":
                assert op.type == "fill_constant", (
4683
                    "For ops in startup program with the op_device attribute "
4684 4685
                    "of cpu, they must be of type fill_constant."
                )
4686 4687 4688
                output_var = op.output_arg_names[0]
                device = self._get_op_device_for_startup_program(output_var)

4689
            if device:
4690
                device_index = int(device.split(':')[1])
4691
            else:
4692 4693
                # LR related ops
                device = None
4694 4695
            if device and device_index != device_id:
                continue
4696
            op_desc = op.desc
4697
            ap_op = new_startup_program.global_block().desc.append_op()
4698 4699 4700
            ap_op.copy_from(op_desc)
            ap_op._set_attr(self._op_device_key, "")
        new_startup_program._sync_with_cpp()
4701
        self._create_vars(new_startup_program.global_block(), block)
4702 4703
        return new_startup_program

4704
    def _find_post_op(self, index, var_name):
H
hutuxian 已提交
4705
        """
4706
        Find the post op that has variable named var_name as input.
H
hutuxian 已提交
4707
        """
4708 4709 4710 4711 4712 4713
        # 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', '')

4714
        post_ops = self.input_var_to_op[var_name]
4715
        if post_ops is None:
4716
            return None
4717 4718 4719 4720 4721 4722
        result_op = None
        for post_op, post_idx in reversed(post_ops):
            if post_idx > index:
                result_op = post_op
                break
        return result_op
4723

4724
    def _find_prev_op(self, index, var_name):
H
hutuxian 已提交
4725
        """
4726 4727
        Find the previous op of op with index that outputs
        variable named var_name.
H
hutuxian 已提交
4728
        """
4729
        prev_ops = self.output_var_to_op[var_name]
4730
        if prev_ops is None:
4731
            return None
4732 4733 4734 4735
        result_op = None
        for prev_op, prev_idx in reversed(prev_ops):
            if prev_idx < index:
                result_op = prev_op
4736
                break
4737
        return result_op
4738 4739

    def _rename_arg(self, op, old_name, new_name):
4740 4741
        op._rename_input(old_name, new_name)
        op._rename_output(old_name, new_name)
4742

4743
    def _create_var(self, block, ref_var, name, dtype=None):
4744 4745 4746 4747 4748 4749 4750 4751
        """
        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,
4752
            dtype=ref_var.dtype if dtype is None else dtype,
4753 4754
            type=ref_var.type,
            lod_level=ref_var.lod_level,
4755 4756
            persistable=ref_var.persistable,
            is_data=ref_var.is_data,
4757 4758
            need_check_feed=ref_var.desc.need_check_feed(),
        )
4759
        self._clone_var_attr(new_var, ref_var)
4760 4761
        return new_var

4762 4763 4764 4765 4766
    def _clone_var_attr(self, dest, src):
        dest.stop_gradient = src.stop_gradient
        if hasattr(src, 'is_distributed'):
            dest.is_distributed = src.is_distributed

4767 4768 4769 4770 4771 4772
    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 已提交
4773

4774 4775 4776 4777 4778 4779
    def _append_grad_suffix(self, name):
        """
        Append grad suffix to the given variable name
        """
        return name + core.grad_var_suffix()

4780
    def _get_op_device_attr(self, op):
H
hutuxian 已提交
4781
        """
4782
        Get the op_device attribute of a op.
H
hutuxian 已提交
4783
        """
4784 4785 4786 4787 4788
        device = (
            op.attr(self._op_device_key)
            if op.has_attr(self._op_device_key)
            else None
        )
4789
        if device:
4790 4791
            assert device[0:3] == 'gpu' or device[0:3] == 'npu', (
                "Now, only gpu and npu devices are "
4792
                "supported in pipeline parallemism."
4793
            )
4794 4795 4796 4797 4798 4799 4800 4801 4802 4803 4804 4805 4806
        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
4807
            op._set_attr(self._op_device_key, f"{self._device}:all")
4808 4809 4810 4811
        # 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():
4812 4813 4814
                assert (
                    '@RENAME@' in name
                ), "The op must be sum used to accumulate renamed vars."
4815 4816 4817 4818
            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(
4819 4820 4821 4822
                'op_device'
            ), "{} has no op_device attr for var {}".format(
                post_op.type, out_name
            )
4823 4824 4825
            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)
4826 4827 4828
        elif (op.type == "cast" or op.type == "scale") and self._is_backward_op(
            op
        ):
4829
            prev_op = self._find_prev_op(idx, op.desc.input("X")[0])
4830 4831
            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):
4832
            # for checkpoint offloading
4833 4834 4835
            assert (
                len(op.input_arg_names) == 1 and len(op.output_arg_names) == 1
            )
4836 4837 4838
            input_name = op.input_arg_names[0]
            output_name = op.output_arg_names[0]
            if '@Fetch' in output_name:
4839
                post_op = self._find_post_op(idx, output_name)
4840 4841 4842
                op._set_attr(
                    self._op_device_key, post_op.attr(self._op_device_key)
                )
4843
            else:
4844
                prev_op = self._find_prev_op(idx, op.desc.input("X")[0])
4845 4846 4847
                op._set_attr(
                    self._op_device_key, prev_op.attr(self._op_device_key)
                )
4848 4849 4850
        elif self._is_loss_op(op):
            # For loss * loss_scaling op added by AMP
            offset = 1
4851 4852 4853
            while not block.ops[idx + offset].has_attr(
                self._op_device_key
            ) or not block.ops[idx + offset].attr(self._op_device_key):
4854 4855 4856 4857 4858 4859 4860 4861 4862
                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
4863
            param_name = self._strip_grad_suffix(grad_name[0])
4864 4865 4866 4867 4868
            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.
4869 4870
            assert self._op_role_var_key in op.attr_names, (
                "gradient_clip "
4871
                "and regularization ops must have op_role_var attribute."
4872
            )
4873
            op_role_var = op.attr(self._op_role_var_key)
4874 4875
            assert len(op_role_var) == 2, (
                "op_role_var for gradient_clip "
4876
                "regularization ops must have two elements."
4877
            )
4878 4879
            param_name = op_role_var[0]
            device = self._param_device_map[param_name]
4880
            # For sum op added by global gradient clip, it must be
4881
            # put on all devices
4882 4883 4884 4885 4886 4887 4888
            if (
                op.type == 'sum'
                or op.type == 'sqrt'
                or op.type == 'fill_constant'
                or op.type == 'elementwise_max'
                or op.type == 'elementwise_div'
            ):
4889
                device = f"{self._device}:all"
4890
            op._set_attr(self._op_device_key, device)
R
Roc 已提交
4891
        elif op.type == "alloc_float_status" or op.type == "clear_float_status":
4892
            op._set_attr(self._op_device_key, f"{self._device}:all")
4893 4894 4895 4896 4897 4898 4899 4900 4901 4902
            # 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
4903 4904
        else:
            other_known_ops = [
4905 4906 4907 4908 4909 4910
                'update_loss_scaling',
                'reduce_any',
                'concat',
                'sum',
                'check_finite_and_unscale',
                'memcpy',
4911
            ]
4912 4913 4914
            assert op.type in other_known_ops, (
                "For other ops without "
                "op_device set, they must be one of {}, but it "
4915
                "is {}".format(other_known_ops, op.type)
4916
            )
4917
            assert self._is_optimize_op(op)
4918
            op._set_attr(self._op_device_key, f"{self._device}:all")
4919 4920

    def _add_op_device_attr(self, block):
4921
        """
4922
        Add op_device attrribute for ops in block that have
4923
        not that attribute set.
4924
        """
4925
        for idx, op in enumerate(list(block.ops)):
4926 4927 4928 4929 4930
            if (
                op.type == "create_py_reader"
                or op.type == "read"
                or op.type == "create_double_buffer_reader"
            ):
4931
                # Copy read related ops to all section to make them exit
4932 4933 4934 4935
                # 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.
4936
                op._set_attr(self._op_device_key, f"{self._device}:all")
4937 4938
                continue
            # op_device attribute has been set
4939 4940
            if self._get_op_device_attr(op):
                continue
4941
            self._add_op_device_attr_for_op(op, idx, block)
H
hutuxian 已提交
4942

4943 4944
    def _check_validation(self, block):
        """
4945
        Check whether ops in a block have both the op_device and the
4946 4947
        op_role attributes set.
        Then, return all devices in order.
4948
        """
4949 4950 4951 4952 4953 4954 4955 4956 4957 4958
        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),
        ]
4959
        for op in block.ops:
4960
            if not op._has_kernel(op.type):
4961 4962 4963 4964 4965 4966
                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."
                )
4967
            assert op.has_attr(
4968 4969
                self._op_role_key
            ), "op ({}) has no {} attribute.".format(op.type, self._op_role_key)
4970
            op_role = op.attr(self._op_role_key)
4971 4972 4973 4974 4975
            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
            )
4976

4977
            assert op.has_attr(
4978 4979 4980 4981
                self._op_device_key
            ), "op ({}) has no {} attribute.".format(
                op.type, self._op_device_key
            )
4982 4983

            device = op.attr(self._op_device_key)
4984 4985 4986 4987 4988 4989 4990
            assert (
                device
            ), "op_device attribute for op " "{} has not been set.".format(
                op.type
            )
            if device == f"{self._device}:all":
                continue
4991

4992
            dev_type = device.split(':')[0]
B
Baibaifan 已提交
4993 4994
            assert dev_type == "gpu" or dev_type == 'npu', (
                "Now only gpu and npu devices are supported "
4995 4996
                "for pipeline parallelism."
            )
4997 4998

            if device not in device_list:
4999
                device_list.append(device)
5000

5001
        return device_list
5002

5003
    def _insert_sendrecv_ops_for_boundaries(self, block):
5004
        """
5005
        Insert a pair of send and recv ops for every two
5006 5007
        consecutive ops on different devices.
        """
5008
        # A map from var to device where op takes it as input,
5009
        # avoiding multiple send and recv ops.
5010
        input_var_to_device = dict()
5011 5012 5013 5014 5015 5016 5017 5018
        # 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,
5019
            'first_optimize_index': first_optimize_index,
5020
        }
5021

5022
        for index, op in enumerate(list(block.ops)):
5023
            cur_device = op.attr(self._op_device_key)
5024 5025
            if cur_device == f"{self._device}:all":
                continue
5026 5027
            for var_name in op.input_arg_names:
                var = block.var(var_name)
5028
                # skip data var
5029 5030
                if var.is_data:
                    continue
5031
                prev_device = None
5032 5033 5034

                prev_op = self._find_prev_op(index, var_name)
                if prev_op is None:
5035 5036
                    if var_name not in self._param_device_map:
                        continue
5037
                    prev_device = self._param_device_map[var_name]
5038

5039
                if not prev_device:
5040 5041 5042
                    prev_device = (
                        prev_op.attr(self._op_device_key) if prev_op else None
                    )
5043

5044 5045
                if prev_device is None or prev_device == f"{self._device}:all":
                    continue
5046

5047 5048
                if prev_device == cur_device:
                    continue
5049

5050 5051 5052 5053 5054 5055 5056
                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] + ':'

5057 5058 5059 5060
                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)
5061 5062
                    assert is_forward or is_backward, (
                        'send/recv in pipeline should only be inserted in forward or backward,'
5063
                        'please check the op_role of op={}'.format(op)
5064
                    )
5065 5066

                    if is_forward:
5067 5068
                        assert prev_id < cur_id, (
                            "In forward, send/recv can only be passed forward, but now "
5069
                            "prev_stage={} great than cur_stage={}, please check op_device of op={}".format(
5070 5071 5072
                                prev_id, cur_id, op
                            )
                        )
5073
                    elif is_backward:
5074 5075
                        assert prev_id > cur_id, (
                            "In backward, send/recv can only be passed backward, but now "
5076
                            "prev_stage={} less than cur_stage={}, please check op_device of op={}".format(
5077 5078 5079
                                prev_id, cur_id, op
                            )
                        )
5080

5081 5082 5083 5084 5085 5086 5087 5088 5089 5090
                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(
5091 5092
                            (cur_dev, prev_dev)
                        )
5093 5094 5095 5096 5097
                        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(
5098 5099
                            (cur_dev, prev_dev)
                        )
5100 5101 5102 5103 5104 5105
                        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)
5106
                    var = block.vars[var_name]
5107 5108 5109
                    pair = (prev_id, cur_id)
                    # 1000 is just a magic number
                    pair_key = prev_id * 1000 + cur_id
5110 5111 5112 5113 5114 5115 5116
                    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]
5117

5118
                    if self.schedule_mode == 'F-then-B':  # F-then-B
F
fangshuixun007 已提交
5119
                        block._insert_op_without_sync(
5120
                            index=index + extra_index_info['index'],
5121 5122 5123
                            type='send_v2',
                            inputs={'X': var},
                            attrs={
5124
                                self._op_device_key: prev_dev,
5125 5126 5127
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 1,
5128 5129 5130
                                'ring_id': ring_id,
                            },
                        )
5131
                        extra_index_info['index'] += 1
5132
                        var_shape = list(var.shape)
5133 5134 5135 5136 5137
                        var_shape[0] = (
                            self.micro_batch_size
                            if var_shape[0] < 0
                            else var_shape[0]
                        )
F
fangshuixun007 已提交
5138
                        block._insert_op_without_sync(
5139
                            index=index + extra_index_info['index'],
5140 5141 5142
                            type='recv_v2',
                            outputs={'Out': [var]},
                            attrs={
5143
                                'out_shape': var_shape,
5144
                                'dtype': var.dtype,
5145
                                self._op_device_key: cur_dev,
5146 5147 5148
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 0,
5149 5150 5151
                                'ring_id': ring_id,
                            },
                        )
5152
                        extra_index_info['index'] += 1
5153
                    elif self.schedule_mode == '1F1B':  # 1F1B
5154
                        var_shape = list(var.shape)
5155 5156 5157 5158 5159
                        var_shape[0] = (
                            self.micro_batch_size
                            if var_shape[0] < 0
                            else var_shape[0]
                        )
5160

5161
                        numel = np.prod(var_shape)
5162 5163 5164
                        use_mp = (self.mp_degree > 1) and (
                            numel % self.mp_degree == 0
                        )
5165 5166 5167 5168 5169 5170 5171 5172 5173 5174 5175 5176 5177 5178 5179 5180 5181 5182 5183 5184 5185 5186

                        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,
5187 5188
                                },
                            )
5189 5190 5191
                            extra_index_info['index'] += 1
                            return

5192 5193
                        _check_stage(cur_id, prev_id)

5194 5195 5196 5197 5198 5199 5200 5201 5202 5203
                        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,
                            },
                        )
5204
                        extra_index_info['index'] += 1
5205 5206
                        prefix_name = var.name.split('@')[0]
                        prefix_var = block.var(prefix_name)
5207 5208 5209
                        is_param = (
                            True if isinstance(prefix_var, Parameter) else False
                        )
F
fangshuixun007 已提交
5210
                        block._insert_op_without_sync(
5211
                            index=index + extra_index_info['index'],
5212
                            type='send_v2'
5213 5214
                            if not use_mp or is_param
                            else 'partial_send',
5215 5216
                            inputs={'X': var},
                            attrs={
5217
                                self._op_device_key: prev_dev,
5218 5219 5220 5221
                                self._op_role_key: op_role,
                                'use_calc_stream': False,
                                'ring_id': ring_id,
                                'peer': 1,
5222 5223 5224
                                # if send_v2, num&id attr is not in op_attrs, will not insert
                                'num': self.mp_degree,
                                'id': self.mp_rank,
5225 5226
                            },
                        )
5227
                        extra_index_info['index'] += 1
5228 5229 5230
                        insert_index = None
                        if int(op_role) == int(self._op_role.Backward):
                            insert_index = extra_index_info[
5231 5232
                                'first_optimize_index'
                            ]
5233 5234 5235 5236
                            new_op_role = self._op_role.Optimize
                        else:
                            insert_index = index
                            new_op_role = self._op_role.Backward
5237
                        sync_comm_op = block._insert_op_without_sync(
5238
                            index=insert_index + extra_index_info['index'],
5239 5240 5241 5242
                            type='c_sync_comm_stream',
                            inputs={'X': [var]},
                            outputs={'Out': [var]},
                            attrs={
5243
                                self._op_device_key: prev_dev,
5244
                                self._op_role_key: new_op_role,
5245
                                'ring_id': ring_id,
5246 5247
                            },
                        )
5248
                        if int(op_role) == int(self._op_role.Forward):
5249
                            sync_comm_op._set_attr('pipeline_flag', '')
5250
                            extra_index_info['index'] += 1
F
fangshuixun007 已提交
5251
                        block._insert_op_without_sync(
5252
                            index=index + extra_index_info['index'],
5253
                            type='recv_v2'
5254 5255
                            if not use_mp or is_param
                            else 'partial_recv',
5256 5257 5258 5259
                            outputs={'Out': [var]},
                            attrs={
                                'out_shape': var_shape,
                                'dtype': var.dtype,
5260
                                self._op_device_key: cur_dev,
5261 5262 5263
                                self._op_role_key: op_role,
                                'use_calc_stream': True,
                                'peer': 0,
5264 5265 5266 5267
                                '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,
5268 5269
                            },
                        )
5270
                        extra_index_info['index'] += 1
5271
                        if use_mp and not is_param:
5272 5273 5274 5275 5276 5277 5278 5279 5280 5281 5282 5283 5284
                            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,
5285 5286
                                },
                            )
5287
                            extra_index_info['index'] += 1
5288 5289 5290
                    else:
                        raise ValueError(
                            "Now only 'F-then-B' and '1F1B' are supported."
5291 5292
                            "The given value is {}.".format(self.schedule_mode)
                        )
5293

5294 5295 5296 5297
                _insert_send_recv(
                    int(cur_device.split(':')[1]),
                    int(prev_device.split(':')[1]),
                )
5298 5299
        block._sync_with_cpp()

5300
    def _insert_loss_scale(self, block):
5301
        """
5302
        Scale the loss corresponding to number of micro-batches.
5303
        """
5304 5305
        if self._num_microbatches == 1:
            return
5306
        for index, op in reversed(tuple(enumerate(list(block.ops)))):
5307
            if self._is_loss_grad_op(op):
5308 5309
                assert op.type == 'fill_constant', (
                    "loss_grad_op must be fill_constant op, "
5310
                    "but this op is {}".format(op.type)
5311
                )
5312 5313 5314 5315
                assert op.has_attr('value')
                loss_scale = float(op.attr('value'))
                loss_scale = loss_scale / self._num_microbatches
                op._set_attr('value', loss_scale)
5316 5317
                break

5318 5319
    def _rename_gradient_var_name(self, block):
        for index, op in enumerate(block.ops):
5320 5321
            if not self._is_optimize_op(op):
                continue
5322 5323 5324
            input_names = op.input_arg_names
            output_names = op.output_arg_names
            in_out_names = input_names + output_names
5325 5326
            if op.type == 'cast' or op.type == "c_sync_comm_stream":
                continue
5327 5328 5329
            # 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:
5330 5331
                if not core.grad_var_suffix() in name:
                    continue
5332 5333 5334 5335
                param_name = name.strip(core.grad_var_suffix())
                new_grad_name = name + "@MERGED"
                self._rename_arg(op, name, new_grad_name)

5336 5337 5338
    def _accumulate_gradients(
        self, block, pp_allreduce_in_optimize=False, strategy=None, shard=None
    ):
5339 5340 5341 5342
        """
        Create a new merged gradient for each parameter and accumulate the
        corresponding gradient to it.
        """
5343 5344
        fp16_allreduce = strategy.fp16_allreduce if strategy else False
        if strategy and strategy.fuse_grad_merge:
5345
            fused_gradient_names = self._accumulate_gradients_with_fuse(
5346 5347
                block, fp16_allreduce, strategy.fuse_grad_size_in_MB, shard
            )
5348 5349
            return fused_gradient_names

5350 5351 5352
        merged_gradient_names = []
        first_opt_op_idx = None

5353 5354 5355
        merged_suffix = '@MERGED@FP16' if fp16_allreduce else '@MERGED'
        dtype = paddle.float16 if fp16_allreduce else None

5356 5357 5358 5359 5360 5361 5362 5363
        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)
5364
                    continue
5365

5366
            if self._is_backward_op(op) and first_opt_op_idx is None:
5367
                first_opt_op_idx = index + 1
5368 5369
                # maybe have no optimize
                # if first_opt_op_idx == len(block.ops): return
5370

5371 5372 5373
            if self._is_backward_op(op) and (
                self._op_role_var_key in op.attr_names
            ):
5374
                op_role_var = op.attr(self._op_role_var_key)
5375 5376
                if len(op_role_var) == 0:
                    continue
5377 5378
                assert len(op_role_var) % 2 == 0
                for i in range(0, len(op_role_var), 2):
5379 5380
                    offset = 0
                    param_name = op_role_var[i]
5381 5382 5383 5384
                    if not block.has_var(param_name):
                        continue
                    if '@BroadCast' in param_name:
                        continue
5385

5386
                    param_grad_name = param_name + core.grad_var_suffix()
5387
                    merged_param_grad_name = param_grad_name + merged_suffix
5388
                    if not block.has_var(merged_param_grad_name):
5389 5390 5391 5392 5393 5394
                        self._create_var(
                            block,
                            block.vars[param_name],
                            merged_param_grad_name,
                            dtype,
                        )
5395
                    assert block.has_var(merged_param_grad_name)
5396

5397 5398 5399
                    param_grad_var = block.var(param_grad_name)
                    merged_param_grad_var = block.var(merged_param_grad_name)
                    merged_param_grad_var.persistable = True
5400
                    block._insert_op(
5401 5402 5403 5404
                        index=first_opt_op_idx + offset,
                        type='fill_constant',
                        inputs={},
                        outputs={'Out': [merged_param_grad_var]},
5405
                        attrs={
5406 5407 5408
                            'shape': merged_param_grad_var.shape,
                            'dtype': merged_param_grad_var.dtype,
                            'value': float(0),
5409
                            # a trick to run this op once per mini-batch
5410 5411 5412
                            self._op_role_key: self._op_role.Optimize.LRSched,
                        },
                    )
5413
                    offset += 1
5414 5415
                    grad_name = op_role_var[i + 1]
                    grad_var = block.vars[grad_name]
5416 5417

                    is_fp16_grad = 'cast_fp16' in grad_name
5418
                    need_cast = is_fp16_grad is not fp16_allreduce
5419 5420 5421 5422 5423 5424

                    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
5425
                        cast_grad_var_name = param_grad_name + '@TMP'
5426
                        cast_grad_var = self._create_var(
5427 5428
                            block, param_grad_var, cast_grad_var_name, dtype
                        )
5429
                        cast_grad_var.persistable = False
5430 5431 5432 5433 5434 5435 5436 5437 5438 5439 5440
                        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,
                            },
                        )
5441
                        offset += 1
5442 5443 5444 5445 5446 5447 5448
                        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},
5449 5450
                        attrs={
                            self._op_role_key: self._op_role.Backward,
5451 5452
                        },
                    )
5453 5454 5455
                    offset += 1
                    merged_gradient_names.append(merged_param_grad_name)

5456 5457
        if not fp16_allreduce:
            return merged_gradient_names
5458 5459 5460 5461 5462 5463 5464 5465 5466 5467 5468 5469 5470 5471 5472 5473 5474 5475 5476 5477 5478 5479 5480

        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

5481 5482 5483 5484 5485 5486 5487 5488 5489 5490 5491
            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,
                },
            )
5492

5493
        return merged_gradient_names
5494

5495 5496 5497
    def _insert_accumulate_gradients_with_fuse(
        self, main_block, fp16, fused_size, grad_param_pairs, first_opt_op_idx
    ):
5498
        grad_param_pairs = self._sort_grad_param_by_dtype(
5499 5500
            main_block, grad_param_pairs
        )
5501

5502 5503 5504
        grad_param_segments = []
        merged_suffix = '@MERGED@FP16' if fp16 else '@MERGED'
        dtype = paddle.float16 if fp16 else paddle.float32
5505
        cur_size = 0.0
5506 5507 5508 5509 5510 5511 5512 5513 5514 5515
        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,
5516 5517
                stop_gradient=False,
            )
5518
            real_param = main_block.var(param)
5519 5520
            if hasattr(real_param, 'is_distributed'):
                merged_grad_var.is_distributed = real_param.is_distributed
5521 5522 5523 5524
            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
5525 5526 5527 5528 5529
            if (
                len(grad_param_segments) == 0
                or cur_size + tmp_size > fused_size
                or real_grad.dtype != last_dtype
            ):
5530
                grad_param_segments.append(
5531 5532
                    ([real_grad], [real_param], [merged_grad_var])
                )
5533
                last_dtype = real_grad.dtype
5534
                cur_size = 0.0
5535 5536 5537 5538 5539 5540 5541 5542 5543 5544 5545 5546
            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]
5547 5548 5549 5550 5551 5552
            fused_grad = main_block.create_var(
                name='FusedGrad_{}'.format(grad_segment[0].name),
                dtype=grad_segment[0].dtype,
                persistable=False,
                stop_gradient=False,
            )
5553
            # keep the '.cast_fp16' info in the fuse var name
5554 5555 5556 5557 5558 5559 5560 5561 5562
            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)
            )
5563 5564 5565 5566
            fused_merged_grad = main_block.create_var(
                name=fused_merged_grad_name,
                dtype=merged_grad_segment[0].dtype,
                persistable=True,
5567 5568
                stop_gradient=False,
            )
5569 5570 5571 5572 5573 5574 5575 5576 5577 5578 5579 5580 5581 5582 5583 5584 5585 5586 5587 5588 5589 5590 5591 5592 5593
            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},
5594
                outputs={"Output": grads, "FusedOutput": fused_grad},
5595 5596 5597 5598 5599 5600 5601 5602 5603 5604 5605 5606 5607 5608 5609 5610
                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,
5611 5612 5613 5614 5615 5616 5617
                    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),
5618 5619
                },
            )
5620 5621 5622 5623 5624 5625 5626 5627 5628 5629
            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,
5630
                    "FusedOutput": fused_merged_grad,
5631 5632 5633 5634 5635 5636 5637 5638
                },
                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,
5639 5640 5641
                    self._op_role_key: self._op_role.Optimize.LRSched,
                },
            )
5642 5643 5644 5645 5646 5647 5648 5649 5650
            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
5651
            need_cast = is_fp16_grad is not fp16
5652 5653 5654 5655
            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'
5656 5657 5658 5659 5660 5661 5662 5663 5664 5665 5666 5667 5668 5669 5670 5671 5672
                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,
                    },
                )
5673 5674 5675 5676 5677 5678 5679
                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},
5680 5681
                attrs={self._op_role_key: self._op_role.Backward},
            )
5682 5683 5684 5685 5686 5687 5688 5689 5690 5691
            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'
5692 5693 5694 5695 5696 5697 5698 5699 5700 5701 5702 5703 5704 5705 5706 5707 5708 5709
                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,
                    },
                )
5710 5711 5712 5713 5714 5715
                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

5716
        return fused_merged_gradients, first_opt_op_idx
5717

5718 5719 5720
    def _accumulate_gradients_with_fuse(
        self, main_block, fp16, fused_size, shard=None
    ):
5721 5722 5723 5724 5725 5726 5727 5728 5729 5730 5731 5732 5733 5734 5735 5736 5737 5738 5739
        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

5740 5741 5742
            if self._is_backward_op(op) and (
                self._op_role_var_key in op.attr_names
            ):
5743 5744 5745 5746 5747 5748 5749 5750 5751 5752 5753
                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(
5754 5755
                        (op_role_var[i + 1], op_role_var[i])
                    )
5756 5757 5758 5759 5760 5761 5762 5763 5764 5765 5766 5767 5768

        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:
5769 5770 5771 5772 5773 5774
            (
                fused_merged_gradients,
                first_opt_op_idx,
            ) = self._insert_accumulate_gradients_with_fuse(
                main_block, fp16, fused_size, pairs, first_opt_op_idx
            )
5775 5776 5777 5778
            all_fused_merged_gradients += fused_merged_gradients

        main_block._sync_with_cpp()
        return all_fused_merged_gradients
5779

5780 5781 5782 5783 5784 5785 5786 5787 5788 5789 5790 5791 5792 5793 5794 5795 5796 5797
    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

5798 5799 5800 5801 5802 5803 5804 5805 5806 5807 5808 5809
    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
5810 5811 5812 5813 5814 5815
        return (
            reduce(lambda x, y: x * y, var.shape)
            * dtype_to_size[var.dtype]
            / 1024.0
            / 1024.0
        )
5816

5817 5818
    def _add_sub_blocks(self, main_block, program_list):
        main_program = main_block.program
5819
        for prog in program_list:
5820 5821 5822 5823 5824 5825
            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)
5826 5827
                for sub_op in origin_sub_block.ops:
                    op_desc = sub_op.desc
5828 5829 5830
                    ap_op = new_sub_block.desc.append_op()
                    ap_op.copy_from(op_desc)
                new_sub_block._sync_with_cpp()
5831
                self._create_vars(new_sub_block, origin_sub_block)
5832
                op._set_attr('sub_block', new_sub_block)
5833 5834 5835

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

5841 5842 5843
    def _process_persistable_vars_in_multi_sections(
        self, main_program, startup_prog, program_list
    ):
5844 5845 5846 5847 5848 5849 5850
        """
        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()
5851
        for prog in program_list:
5852 5853
            block = prog.block(0)
            for var_name in block.vars:
5854 5855
                if var_name == "double_buffer_0":
                    continue
5856
                var = block.var(var_name)
5857 5858
                if not var.persistable:
                    continue
5859 5860 5861 5862 5863 5864 5865 5866 5867 5868 5869 5870 5871 5872 5873
                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:
5874 5875 5876 5877 5878 5879
                    if (
                        op.type == "recv_v2"
                        or op.type == "create_py_reader"
                        or op.type == "read"
                        or op.type == "update_loss_scaling"
                    ):
5880
                        continue
5881 5882
                    # We have processed lr related vars
                    if op.attr(self._op_role_key) == int(
5883 5884
                        self._op_role.Optimize.LRSched
                    ):
5885 5886 5887 5888
                        continue
                    if var_name in op.desc.output_arg_names():
                        assert var_name not in write_info, (
                            "two sections write the same var({}): second "
5889 5890
                            "op {}.".format(var_name, op)
                        )
5891 5892 5893 5894 5895
                        write_info[var_name] = prog
                        break

        for var_name in var_info.keys():
            # Case 1: read only variables, no special process
5896 5897
            if not var_name in write_info:
                continue
5898 5899 5900 5901 5902

            # 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)
5903
            write_dev_index = int(write_device.split(':')[1])
5904 5905
            all_progs = var_info[var_name]
            for prog in all_progs:
5906 5907
                if prog == write_prog:
                    continue
5908 5909 5910
                read_block = prog.block(0)
                read_device = self._get_device_info(read_block)
                read_dev_index = int(read_device.split(':')[1])
5911 5912 5913 5914 5915 5916 5917 5918 5919
                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]
5920 5921 5922

                write_block._insert_op(
                    index=0,
5923
                    type='send_v2',
5924 5925 5926
                    inputs={
                        'X': write_block.var(var_name),
                    },
5927
                    attrs={
5928 5929
                        self._op_device_key: write_device,
                        'use_calc_stream': False,
5930 5931
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
5932 5933 5934 5935 5936
                        self._op_role_key: self._op_role.LRSched,
                        'peer': read_dev_index,
                        'ring_id': ring_id,
                    },
                )
5937 5938
                read_block._insert_op(
                    index=0,
5939
                    type='recv_v2',
5940 5941
                    outputs={'Out': [read_block.var(var_name)]},
                    attrs={
5942 5943 5944 5945
                        '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,
5946 5947
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
5948 5949 5950 5951 5952
                        self._op_role_key: self._op_role.LRSched,
                        'peer': write_dev_index,
                        'ring_id': ring_id,
                    },
                )
5953 5954 5955 5956 5957 5958
                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={
5959
                        self._op_device_key: read_device,
5960 5961
                        # A trick to make the role LRSched to avoid copy every
                        # microbatch
5962 5963 5964 5965
                        self._op_role_key: self._op_role.LRSched,
                        'ring_id': ring_id,
                    },
                )
5966 5967

    def _is_gradient_clip_op(self, op):
5968 5969 5970
        return op.desc.has_attr("op_namescope") and op.desc.attr(
            "op_namescope"
        ).startswith("/gradient_clip")
5971 5972

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

5977 5978
    def _is_weight_decay_op(self, op):
        # in AdamW namescope is /optimizer_*/weight decay/
5979 5980 5981
        return op.desc.has_attr(
            "op_namescope"
        ) and 'weight decay' in op.desc.attr("op_namescope")
5982

5983 5984 5985 5986 5987
    def _get_input_output_info(self, block):
        '''
        Get info of op input and output.
        '''
        # A map from output var to op which generate it.
5988
        output_var_to_op = defaultdict(list)
5989
        # A map from var to op which takes it as input.
5990
        input_var_to_op = defaultdict(list)
5991

5992
        for index, op in enumerate(block.ops):
5993
            for var_name in op.input_arg_names:
5994
                input_var_to_op[var_name].append([op, index])
5995
            for var_name in op.output_arg_names:
5996 5997 5998 5999 6000 6001 6002 6003
                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
        """
6004 6005
        if self.schedule_mode != '1F1B':
            return
6006 6007 6008

        block = program.block(0)

6009
        recv_type = 'recv_v2' if self.mp_degree == 1 else 'partial_recv'
6010 6011
        backward_recv_index = None
        for index, op in enumerate(block.ops):
6012
            if op.type == recv_type and self._is_backward_op(op):
6013 6014 6015
                backward_recv_index = index
                break

6016
        # last pipeline stage
6017 6018
        if backward_recv_index is None:
            return
6019 6020 6021

        offset = 0
        for index, op in enumerate(list(block.ops)):
6022 6023
            if index >= backward_recv_index:
                break
6024 6025 6026 6027 6028 6029 6030 6031 6032 6033 6034 6035 6036 6037 6038 6039
            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]},
6040 6041
                    attrs={self._op_role_key: self._op_role.Backward},
                )
6042
        block._sync_with_cpp()
6043

6044 6045 6046 6047 6048 6049 6050 6051 6052 6053 6054 6055 6056
    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))
6057 6058 6059 6060
            if (
                op_role == int(self._op_role.Backward)
                and backward_insert_index is None
            ):
6061
                backward_insert_index = i
6062 6063 6064 6065 6066 6067
            if (
                op.type != "partial_recv"
                and op.type != "partial_allgather"
                and op.type != "nop"
                and op.type != "recv_v2"
            ):
6068 6069 6070 6071 6072 6073 6074 6075 6076 6077 6078 6079 6080 6081 6082 6083 6084 6085 6086
                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)
6087 6088 6089 6090 6091 6092 6093
            block._insert_op_without_sync(
                index=insert_index,
                type=op.type,
                inputs=op_inputs,
                outputs=op_outputs,
                attrs=op.all_attrs(),
            )
6094 6095 6096 6097 6098 6099 6100
            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()

6101 6102 6103 6104 6105 6106 6107 6108 6109 6110 6111 6112 6113 6114 6115 6116 6117 6118 6119 6120 6121 6122 6123 6124 6125 6126 6127
    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 "
6128 6129
            "forward and used in backward:\n{}".format(used_in_backward)
        )
6130

6131 6132 6133
    def minimize(
        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
    ):
6134
        main_block = loss.block
6135
        self.origin_main_block = main_block
6136
        main_program = main_block.program
6137 6138
        if startup_program is None:
            startup_program = default_startup_program()
6139

6140 6141
        pipeline_opt = main_program._pipeline_opt
        assert pipeline_opt, 'Please use pipeline with fleet.'
6142 6143 6144 6145 6146 6147 6148
        required_keys = [
            'local_rank',
            'schedule_mode',
            'micro_batch_size',
            'ring_id',
            'global_ring_id',
            'use_sharding',
6149 6150
            'mp_degree',
            'mp_rank',
6151 6152
        ]
        for key in required_keys:
6153 6154 6155
            assert (
                key in pipeline_opt
            ), 'Please use pipeline with fleet to use {}.'.format(key)
6156 6157 6158 6159 6160 6161 6162 6163
        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']
6164
        self.scale_gradient = pipeline_opt.get('scale_gradient', False)
6165 6166
        assert self.mp_degree >= 1
        assert 0 <= self.mp_rank < self.mp_degree
6167 6168

        optimize_ops, params_grads = self._optimizer.minimize(
6169 6170
            loss, startup_program, parameter_list, no_grad_set
        )
6171
        self._param_device_map = self._origin_optimizer._param_device_map
6172

6173 6174 6175 6176
        (
            self.output_var_to_op,
            self.input_var_to_op,
        ) = self._get_input_output_info(main_block)
6177 6178 6179
        # Step1: add default op_device attribute for ops.
        self._add_op_device_attr(main_block)
        device_list = self._check_validation(main_block)
6180 6181 6182 6183 6184 6185 6186 6187 6188 6189 6190

        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

6191 6192 6193
        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 "
6194 6195
            "another in the order of their ids."
        )
6196
        # Step2: add send and recv ops between section boundaries
6197
        self._insert_sendrecv_ops_for_boundaries(main_block)
6198

6199
        # Step3: split program into sections and add pairs of
6200 6201
        # send and recv ops for data var.
        main_program = main_block.program
6202
        program_list = self._split_program(main_program, device_list)
6203
        for p in program_list:
6204
            self._create_vars(p.global_block(), main_block)
6205

L
lilong12 已提交
6206 6207 6208 6209 6210
        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 "
6211 6212
                "stages."
            )
L
lilong12 已提交
6213 6214
        else:
            self.local_rank %= len(device_list)
6215 6216 6217
        # Step3.5: optimize forward send sync_comm to overlap send and recv
        self._optimize_forward_send_sync(program_list[self.local_rank])

6218
        # Step4: Special Case: process persistable vars that exist in
6219
        # multiple sections
6220
        # FIXME
6221 6222
        # self._process_persistable_vars_in_multi_sections(
        #     main_program, startup_program, program_list)
6223

6224
        # Step5: Add sub blocks for section programs
6225 6226
        self._add_sub_blocks(main_block, program_list)

6227
        place_list = []
6228 6229
        for dev in device_list:
            dev_index = int(dev.split(":")[1])
6230 6231 6232 6233
            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))
6234

6235
        # Step6: Split startup program
6236
        new_startup_program = self._split_startup_program(
6237 6238
            startup_program, self.local_rank
        )
6239 6240 6241 6242

        startup_program._pipeline_opt = {
            "startup_program": new_startup_program,
        }
6243
        real_block = program_list[self.local_rank].global_block()
6244 6245
        if not self.scale_gradient:
            self._insert_loss_scale(real_block)
6246
        if not self.use_sharding:
6247
            # Step7: clear gradients before each mini-batch and
6248 6249 6250 6251 6252
            # 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()
6253

6254 6255 6256 6257
        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"))
6258 6259 6260
        # A pass to move the recv op to the beginning of
        # the forward/backward phase
        self._mv_head_recv(program_list[self.local_rank])
6261 6262 6263 6264 6265

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

6266
        main_program._pipeline_opt = {
H
hutuxian 已提交
6267 6268
            "trainer": "PipelineTrainer",
            "device_worker": "Section",
6269
            "pipeline_stage": self.local_rank,
6270
            "num_pipeline_stages": len(device_list),
6271
            "schedule_mode": self.schedule_mode,
6272
            "inner_parallelism": len(device_list),
6273 6274
            "section_program": program_list[self.local_rank],
            "place": place_list[self.local_rank],
6275
            "place_id": place_id,
6276
            "sync_steps": -1,
L
lilong12 已提交
6277
            "num_microbatches": self._num_microbatches,
H
hutuxian 已提交
6278 6279
            "start_cpu_core_id": self._start_cpu_core_id,
        }
6280 6281 6282 6283 6284 6285 6286
        return (
            optimize_ops,
            params_grads,
            program_list,
            self._pipeline_pair,
            self._pp_ring_map,
        )
M
mapingshuo 已提交
6287 6288


M
mapingshuo 已提交
6289 6290
class RecomputeOptimizer(Optimizer):
    """
6291
        :api_attr: Static Graph
S
swtkiwi 已提交
6292

M
mapingshuo 已提交
6293 6294 6295
    Recompute Optimizer Wrapper

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

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

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

M
mapingshuo 已提交
6308 6309 6310 6311 6312 6313 6314 6315 6316
    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

6317
            import paddle
M
mapingshuo 已提交
6318 6319
            import paddle.fluid as fluid
            import numpy as np
6320 6321 6322

            paddle.enable_static()

M
mapingshuo 已提交
6323 6324 6325 6326 6327 6328 6329
            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')
6330 6331 6332 6333 6334
                cost = paddle.nn.functional.cross_entropy(
                    input=prediction, label=input_y,
                    reduction='none', use_softmax=False
                )
                sum_cost = paddle.mean(cost)
M
mapingshuo 已提交
6335 6336 6337 6338 6339 6340 6341 6342 6343 6344 6345 6346 6347 6348 6349 6350 6351 6352 6353 6354 6355 6356 6357 6358 6359
                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):
姜永久 已提交
6360
        if in_dygraph_mode():
Z
zhongpu 已提交
6361
            raise Exception("In dygraph, don't support RecomputeOptimizer.")
M
mapingshuo 已提交
6362 6363
        self._optimizer = optimizer
        self._checkpoints = None
M
mapingshuo 已提交
6364 6365
        self._learning_rate = self._optimizer._learning_rate
        self._learning_rate_map = self._optimizer._learning_rate_map
J
JZ-LIANG 已提交
6366
        self.enable_offload = False
M
mapingshuo 已提交
6367 6368

    def _set_checkpoints(self, checkpoints):
6369 6370
        """
        Args:
6371
            checkpoints (list): List of Variable or string
6372 6373 6374 6375 6376
        """
        assert isinstance(
            checkpoints, list
        ), "_checkpoints should be a list of Variable or a list of String"
        for ckpt in checkpoints:
6377 6378
            assert isinstance(ckpt, str) or isinstance(
                ckpt, Variable
6379
            ), "_checkpoints should be a list of Variable or a list of String"
M
mapingshuo 已提交
6380 6381
        self._checkpoints = checkpoints

6382
    # should enable offload before calling backward
J
JZ-LIANG 已提交
6383 6384 6385
    def _enable_offload(self):
        self.enable_offload = True

6386 6387
    @framework.deprecate_stat_dict
    def load(self, state_dict):
M
mapingshuo 已提交
6388
        """
6389
            :api_attr: Static Graph
S
swtkiwi 已提交
6390

M
mapingshuo 已提交
6391 6392 6393 6394
        load function is not supported by Recompute Optimizer for now.
        :return: None

        Args:
6395
            state_dict: the dict load by load_persistable method
M
mapingshuo 已提交
6396 6397 6398 6399

        Examples:
            .. code-block:: python

6400
                import paddle
M
mapingshuo 已提交
6401
                import paddle.fluid as fluid
6402

6403
                paddle.enable_static()
M
mapingshuo 已提交
6404 6405 6406
                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')
6407 6408 6409 6410 6411
                    cost = paddle.nn.functional.cross_entropy(
                        input=prediction, label=input_y,
                        reduction='none', use_softmax=False
                    )
                    sum_cost = paddle.mean(cost)
M
mapingshuo 已提交
6412
                    return sum_cost, fc_1, prediction
6413

M
mapingshuo 已提交
6414 6415 6416 6417
                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")
6418

M
mapingshuo 已提交
6419 6420 6421 6422
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
                sgd._set_checkpoints([fc_1, pred])
                try:
6423 6424
                    state_dict = {}
                    sgd.load(state_dict)
M
mapingshuo 已提交
6425
                except NotImplementedError as e:
6426
                    print(e)
M
mapingshuo 已提交
6427 6428
        """
        raise NotImplementedError(
6429 6430
            "load function is not supported by Recompute Optimizer for now"
        )
M
mapingshuo 已提交
6431 6432 6433 6434 6435 6436 6437 6438 6439 6440 6441 6442 6443 6444

    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

6445
                import paddle
M
mapingshuo 已提交
6446 6447 6448
                import paddle.fluid as fluid
                import paddle.fluid.framework as framework

6449 6450
                paddle.enable_static()

M
mapingshuo 已提交
6451 6452 6453
                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')
6454 6455 6456 6457 6458
                    cost = paddle.nn.functional.cross_entropy(
                        input=prediction, label=input_y,
                        reduction='none', use_softmax=False
                    )
                    sum_cost = paddle.mean(cost)
M
mapingshuo 已提交
6459 6460 6461 6462 6463 6464 6465 6466 6467 6468
                    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)
6469
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
6470 6471 6472 6473
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
6474
                    no_grad_set=None)
M
mapingshuo 已提交
6475 6476 6477 6478 6479 6480 6481 6482 6483 6484

                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 已提交
6485 6486 6487 6488 6489 6490 6491 6492 6493
    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,
6494 6495
            stop_gradient=True,
        )
J
JZ-LIANG 已提交
6496 6497 6498 6499 6500 6501

        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,
6502 6503
            stop_gradient=False,
        )
J
JZ-LIANG 已提交
6504 6505 6506 6507 6508 6509 6510 6511

        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
6512 6513 6514
        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 已提交
6515 6516 6517 6518 6519 6520 6521 6522 6523 6524 6525 6526 6527
        """
        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,
6528 6529 6530 6531 6532 6533 6534 6535 6536 6537 6538 6539 6540
                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 已提交
6541 6542 6543

        return

6544 6545 6546
    def _insert_async_memcpy_op(
        self, insert_idx, src_varname, dst_varname, op_role, dst_place_type
    ):
J
JZ-LIANG 已提交
6547 6548 6549 6550 6551 6552 6553 6554
        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)]
            },
6555 6556
            attrs={"dst_place_type": int(dst_place_type), OP_ROLE_KEY: op_role},
        )
J
JZ-LIANG 已提交
6557 6558

    def _insert_fetch_op(self, idx, varname):
6559 6560 6561 6562 6563
        assert (
            varname in self.checkpoint_name2pinned_name
        ), "Try to fetch {} from Pinned Memory, but it is NOT a checkpoint".format(
            varname
        )
J
JZ-LIANG 已提交
6564 6565 6566

        pinned_varname = self.checkpoint_name2pinned_name[varname]
        fetch_varname = self.checkpoint_name2fetch_name[varname]
6567
        self._insert_async_memcpy_op(idx, pinned_varname, fetch_varname, 1, 1)
J
JZ-LIANG 已提交
6568 6569

    def _insert_offload_op(self, idx, varname):
6570 6571 6572 6573 6574
        assert (
            varname in self.checkpoint_name2pinned_name
        ), "Try to offload {} to Pinned Memory, but it is NOT a checkpoint".format(
            varname
        )
J
JZ-LIANG 已提交
6575
        pinned_varname = self.checkpoint_name2pinned_name[varname]
6576
        self._insert_async_memcpy_op(idx, varname, pinned_varname, 0, 2)
J
JZ-LIANG 已提交
6577 6578

    def _insert_sync_op(self, op_idx, checkpoint_name):
6579
        # single stream offload no need sync
J
JZ-LIANG 已提交
6580 6581 6582
        pass

    def _record_fetch_op(self, idx):
6583 6584 6585
        assert (
            len(self.un_fetch_checkpoint_names) > 0
        ), "Could NOT found checkpoint to fetch"
J
JZ-LIANG 已提交
6586 6587 6588 6589 6590 6591 6592 6593
        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)
6594 6595 6596 6597 6598
        assert (
            checkpoint_name == expected_checkpoint_name
        ), "expected to offload [{}] but got [{}]".format(
            expected_checkpoint_name, checkpoint_name
        )
J
JZ-LIANG 已提交
6599 6600 6601 6602
        logging.debug("Record offload [{}]".format(checkpoint_name))
        self.idx2insertions[idx] = ("offload", checkpoint_name)

    def _record_sync_op(self, idx, checkpoint_name):
6603 6604 6605
        assert (
            checkpoint_name not in self.synced_checkpoints
        ), "Try to sync the checkpoint [{}] twice".format(checkpoint_name)
J
JZ-LIANG 已提交
6606 6607 6608 6609 6610 6611 6612
        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 = {}
6613
        # don't offload the last checkpoints, to favor throughput
J
JZ-LIANG 已提交
6614 6615 6616 6617 6618 6619 6620 6621 6622 6623 6624 6625 6626 6627
        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(
6628 6629
            self.block.ops
        ), "Could NOT found backword op in prog"
J
JZ-LIANG 已提交
6630 6631 6632

        # fetch second to last checkpoint at the beginning of BW
        fetched_checkpoint_varname = self._record_fetch_op(
6633 6634
            self.bw_strart_op_idx
        )
J
JZ-LIANG 已提交
6635 6636
        last_last_fetch_checkpoint = None

6637
        for i, op in enumerate(self.block.ops[self.bw_strart_op_idx :]):
J
JZ-LIANG 已提交
6638 6639 6640 6641 6642 6643 6644 6645 6646
            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
6647 6648 6649
                            second_to_last_fetch_checkpoint = (
                                fetched_checkpoint_varname
                            )
6650
                            # there is NO fetch ahead the first checkpoint
J
JZ-LIANG 已提交
6651
                            if input_var != self.sorted_checkpoint_names[0]:
6652 6653 6654
                                fetched_checkpoint_varname = (
                                    self._record_fetch_op(idx)
                                )
J
JZ-LIANG 已提交
6655

6656
                        # should check the current used checkpoint is ths last fetch one
6657 6658 6659 6660 6661
                        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 已提交
6662 6663 6664
                        # rename
                        self.block.ops[idx]._rename_input(
                            input_var,
6665 6666
                            self.checkpoint_name2fetch_name[input_var],
                        )
J
JZ-LIANG 已提交
6667 6668 6669 6670
                        self.checkpoint_usage_count[input_var] += 1
                    else:
                        raise ValueError(
                            "use checkpoint [{}] before fetch in BW".format(
6671 6672 6673
                                input_var
                            )
                        )
J
JZ-LIANG 已提交
6674

6675 6676 6677 6678 6679
        assert (
            len(self.un_fetch_checkpoint_names) == 0
        ), "{} checkpoints have NOT been Recorded".format(
            self.un_fetch_checkpoint_names
        )
J
JZ-LIANG 已提交
6680 6681 6682 6683 6684 6685 6686 6687 6688 6689

    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)
6690
                    logging.debug(
6691 6692
                        "Insert [{}] fetch op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
6693 6694 6695 6696 6697
                    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()
6698 6699 6700 6701 6702
        assert (
            len(self.idx2insertions) == 0
        ), "{} checkpoints left un-Fecthed".format(
            [ele[1] for ele in self.idx2insertions.values()]
        )
J
JZ-LIANG 已提交
6703 6704 6705 6706

    def _parse_forward(self):

        self.idx2insertions = {}
6707
        # don't offload the last checkpoints, faster, less memory saving
J
JZ-LIANG 已提交
6708 6709 6710 6711 6712 6713 6714
        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,
6715
                'idx': -1,
J
JZ-LIANG 已提交
6716 6717 6718 6719 6720 6721 6722 6723 6724
            }
        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(
6725 6726
            self.block.ops
        ), "Could NOT found Forward op in prog"
J
JZ-LIANG 已提交
6727 6728
        last_offload_checkpoint = None

6729
        for i, op in enumerate(
6730 6731
            self.block.ops[self.fw_strart_op_idx : self.bw_strart_op_idx]
        ):
J
JZ-LIANG 已提交
6732 6733 6734 6735 6736 6737 6738

            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:
6739 6740 6741 6742 6743
                    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 已提交
6744 6745 6746

                    if output_var in self.un_offload_checkpoint_names:
                        # insert sync op if last checkpoint has not been sync
6747
                        if last_offload_checkpoint is not None:
6748 6749 6750 6751 6752 6753 6754 6755 6756
                            if (
                                self.checkpoint_usage_count_and_idx[
                                    last_offload_checkpoint
                                ]['count']
                                == 0
                            ):
                                self._record_sync_op(
                                    idx, last_offload_checkpoint
                                )
J
JZ-LIANG 已提交
6757
                            else:
6758 6759 6760 6761 6762 6763 6764 6765 6766 6767 6768 6769 6770
                                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 已提交
6771 6772 6773 6774 6775
                        # 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(
6776 6777 6778 6779
                            "There should be just ONE op that output checkpoint [{}]".format(
                                output_var
                            )
                        )
J
JZ-LIANG 已提交
6780 6781
                # need to sync the last need to offload checkpoint before the last checkpoint as output op
                if output_var == last_checkpoint:
6782 6783 6784 6785 6786 6787 6788 6789 6790 6791 6792 6793 6794
                    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 已提交
6795
                    # sync if last checkpoint has not been sync
6796 6797 6798 6799 6800 6801
                    if (
                        self.checkpoint_usage_count_and_idx[
                            last_offload_checkpoint
                        ]['idx']
                        == 0
                    ):
J
JZ-LIANG 已提交
6802 6803 6804
                        self._record_sync_op(idx, last_offload_checkpoint)
                    else:
                        last_usage_idx = self.checkpoint_usage_count_and_idx[
6805 6806 6807 6808 6809 6810 6811 6812 6813 6814
                            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
                        )
6815
            # record checkpoint usage
J
JZ-LIANG 已提交
6816 6817
            for input_var in input_vars:
                if input_var in need_offload_checkpoint_names:
6818 6819 6820
                    assert (
                        input_var not in self.synced_checkpoints
                    ), "checkpoint [{}] used after sync".format(input_var)
J
JZ-LIANG 已提交
6821 6822 6823
                    self.checkpoint_usage_count_and_idx[input_var]['count'] += 1
                    self.checkpoint_usage_count_and_idx[input_var]['idx'] = idx

6824 6825 6826 6827 6828
        assert (
            len(self.un_offload_checkpoint_names) == 0
        ), "{} checkpoints have NOT been Recorded".format(
            self.un_fetch_checkpoint_names
        )
J
JZ-LIANG 已提交
6829 6830 6831
        assert len(self.synced_checkpoints) == len(
            need_offload_checkpoint_names
        ), "{} checkpoints have NOT been Recorded".format(
6832 6833
            set(need_offload_checkpoint_names) - set(self.synced_checkpoints)
        )
J
JZ-LIANG 已提交
6834 6835 6836 6837 6838

    def _update_forward(self):
        if len(self.idx2insertions) == 0:
            return
        for op_idx in reversed(
6839 6840
            range(self.fw_strart_op_idx, self.bw_strart_op_idx)
        ):
J
JZ-LIANG 已提交
6841 6842 6843 6844
            if op_idx in self.idx2insertions:
                operation, checkpoint_name = self.idx2insertions[op_idx]
                if operation == "offload":
                    self._insert_offload_op(op_idx, checkpoint_name)
6845
                    logging.debug(
6846 6847
                        "Insert [{}] offload op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
6848 6849 6850
                    del self.idx2insertions[op_idx]
                elif operation == "sync":
                    self._insert_sync_op(op_idx, checkpoint_name)
6851
                    logging.debug(
6852 6853
                        "Insert [{}] offload_sync op.".format(checkpoint_name)
                    )
J
JZ-LIANG 已提交
6854 6855 6856
                    del self.idx2insertions[op_idx]

        self.block._sync_with_cpp()
6857 6858 6859 6860 6861
        assert (
            len(self.idx2insertions) == 0
        ), "{} checkpoints left un-Offloaded".format(
            [ele[1] for ele in self.idx2insertions.values()]
        )
J
JZ-LIANG 已提交
6862 6863 6864 6865 6866 6867 6868 6869

    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
6870
        1. create pinned vars and temp vars
J
JZ-LIANG 已提交
6871 6872 6873 6874 6875 6876
        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
6877
        if startup_program is None:
J
JZ-LIANG 已提交
6878
            startup_program = paddle.static.default_startup_program()
J
JZ-LIANG 已提交
6879 6880

        with program_guard(self._main_program, startup_program):
6881 6882 6883 6884 6885 6886 6887 6888 6889 6890
            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 已提交
6891 6892 6893 6894
            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(
6895 6896
                    checkpoint_varname
                )
J
JZ-LIANG 已提交
6897
                self.checkpoint_name2pinned_name[
6898 6899
                    checkpoint_varname
                ] = pinned_var_name
J
JZ-LIANG 已提交
6900
                self.checkpoint_name2fetch_name[
6901 6902
                    checkpoint_varname
                ] = fetch_var_name
J
JZ-LIANG 已提交
6903 6904 6905 6906 6907 6908 6909 6910 6911 6912 6913 6914 6915
            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

6916 6917 6918 6919 6920 6921 6922 6923
    def backward(
        self,
        loss,
        startup_program=None,
        parameter_list=None,
        no_grad_set=None,
        callbacks=None,
    ):
M
mapingshuo 已提交
6924 6925 6926 6927 6928 6929 6930
        """
        call append_backward with checkpoints.

        Args:
            loss (Variable): loss variable to run optimizations.
            startup_program (Program): startup_program for initializing parameters
                in `parameter_list`.
6931 6932
            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 已提交
6933 6934 6935 6936 6937 6938 6939
            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

6940
                import paddle
M
mapingshuo 已提交
6941
                import paddle.fluid as fluid
6942

6943 6944
                paddle.enable_static()

M
mapingshuo 已提交
6945 6946 6947
                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')
6948 6949 6950 6951 6952
                    cost = paddle.nn.functional.cross_entropy(
                        input=prediction, label=input_y,
                        reduction='none', use_softmax=False
                    )
                    sum_cost = paddle.mean(cost)
M
mapingshuo 已提交
6953
                    return sum_cost, fc_1, prediction
6954 6955


M
mapingshuo 已提交
6956 6957 6958 6959
                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")
6960

M
mapingshuo 已提交
6961 6962
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
6963
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
6964 6965 6966 6967
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
6968
                    no_grad_set=None)
M
mapingshuo 已提交
6969 6970
                print("Finished backward")
        """
6971 6972 6973
        assert (
            self._checkpoints is not None
        ), "You should call _set_checkpoints first"
M
mapingshuo 已提交
6974

姜永久 已提交
6975
        if in_dygraph_mode():
M
mapingshuo 已提交
6976
            raise NotImplementedError(
6977 6978
                "DyGraph current does not support recompute"
            )
M
mapingshuo 已提交
6979 6980 6981 6982

        self._dtype = loss.dtype
        program = loss.block.program
        with program_guard(program, startup_program):
6983 6984 6985 6986 6987 6988 6989
            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 已提交
6990 6991 6992 6993 6994 6995
            # 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,
6996 6997
                    checkpoints=checkpoint_vars,
                )
J
JZ-LIANG 已提交
6998
            else:
6999 7000 7001 7002 7003 7004
                params_grads = append_backward(
                    loss,
                    parameter_list,
                    no_grad_set,
                    checkpoints=checkpoint_vars,
                )
J
JZ-LIANG 已提交
7005 7006 7007 7008 7009

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

M
mapingshuo 已提交
7010 7011 7012 7013 7014 7015 7016 7017 7018 7019 7020 7021
        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
7022
                import paddle
M
mapingshuo 已提交
7023
                import paddle.fluid as fluid
7024

7025 7026
                paddle.enable_static()

M
mapingshuo 已提交
7027 7028 7029
                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')
7030 7031 7032 7033 7034
                    cost = paddle.nn.functional.cross_entropy(
                        input=prediction, label=input_y,
                        reduction='none', use_softmax=False
                    )
                    sum_cost = paddle.mean(cost)
7035 7036
                    return sum_cost, fc_1, prediction

M
mapingshuo 已提交
7037 7038 7039 7040
                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")
7041

M
mapingshuo 已提交
7042 7043
                sgd = fluid.optimizer.Adam(learning_rate=0.01)
                sgd = fluid.optimizer.RecomputeOptimizer(sgd)
7044
                sgd._set_checkpoints([fc_1, pred])
M
mapingshuo 已提交
7045 7046 7047 7048
                params_grads = sgd.backward(
                    cost,
                    startup_program=None,
                    parameter_list=None,
7049
                    no_grad_set=None)
7050

M
mapingshuo 已提交
7051 7052
                optimize_ops = sgd.apply_optimize(
                    cost, startup_program=None, params_grads=params_grads)
7053

M
mapingshuo 已提交
7054 7055 7056
                print("Finished apply_optimize")
        """

7057 7058 7059 7060 7061 7062 7063 7064 7065 7066 7067 7068
        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
    ):
7069
        assert isinstance(loss, Variable), "The loss should be an Variable."
7070 7071 7072
        assert (
            self._checkpoints is not None
        ), "You should call _set_checkpoints first"
姜永久 已提交
7073
        if in_dygraph_mode():
M
mapingshuo 已提交
7074
            raise NotImplementedError(
7075 7076 7077 7078 7079 7080 7081 7082
                "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 已提交
7083

7084 7085 7086
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
M
mapingshuo 已提交
7087 7088 7089 7090

        return optimize_ops, params_grads


7091
class LookaheadOptimizer:
7092
    r"""
7093
        :api_attr: Static Graph
S
swtkiwi 已提交
7094

M
mapingshuo 已提交
7095 7096 7097 7098
    This implements the Lookahead optimizer of the
    paper : https://arxiv.org/abs/1907.08610.

    Lookahead keeps two sets of params: the fast_params and
7099 7100
    the slow_params. inner_optimizer update fast_params every
    training step. Lookahead updates the slow_params and fast_params
M
mapingshuo 已提交
7101 7102 7103
    every k training steps as follows:

    .. math::
7104

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

7107
        fast\_param_t &=  slow\_param_t
M
mapingshuo 已提交
7108 7109

    Args:
7110
        inner_optimizer (Optimizer): The optimizer that update fast params step by step.
M
mapingshuo 已提交
7111 7112 7113 7114 7115 7116 7117 7118 7119
        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
7120
            import numpy.random as random
M
mapingshuo 已提交
7121

7122
            paddle.enable_static()
7123

7124 7125 7126
            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")
7127 7128 7129 7130
            loss = paddle.nn.functional.cross_entropy(
                input=y, label=label,
                reduction='none', use_softmax=False
            )
7131
            loss = paddle.mean(x=loss)
7132 7133 7134 7135 7136 7137 7138 7139 7140
            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 已提交
7141

7142 7143 7144
            def train_reader(limit=5):
                for i in range(limit):
                    yield random.random([2]).astype('float32'), random.random([1]).astype('int64')
7145

7146 7147
            feeder = fluid.DataFeeder(feed_list=[x, label], place=place)
            reader = paddle.batch(paddle.reader.shuffle(train_reader, buf_size=50000),batch_size=1)
7148

7149 7150 7151
            for batch_data in reader():
                exe.run(fluid.default_main_program(),
                feed=feeder.feed(batch_data))
M
mapingshuo 已提交
7152 7153 7154 7155 7156

    """

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

姜永久 已提交
7157
        if in_dygraph_mode():
Z
zhongpu 已提交
7158
            raise Exception("In dygraph, don't support LookaheadOptimizer.")
7159
        assert inner_optimizer is not None, "inner optimizer can not be None"
M
mapingshuo 已提交
7160 7161 7162
        assert (
            0.0 <= alpha <= 1.0
        ), "alpha should be larger or equal to 0.0, and less or equal than 1.0"
7163
        assert isinstance(k, int) and k > 0, "k should be a positive integer"
M
mapingshuo 已提交
7164 7165 7166 7167 7168 7169 7170 7171 7172 7173

        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(
7174 7175
            loss, startup_program=startup_program
        )
M
mapingshuo 已提交
7176 7177 7178 7179 7180 7181 7182 7183 7184 7185 7186

        # 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)
7187 7188 7189 7190 7191 7192 7193
            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 已提交
7194 7195 7196 7197 7198 7199
            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)
7200 7201 7202 7203 7204 7205 7206
            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 已提交
7207

7208 7209 7210
            startup_block.append_op(
                type="assign", inputs={"X": fast_var}, outputs={"Out": slow_var}
            )
M
mapingshuo 已提交
7211

7212 7213
        with framework.program_guard(main_block.program, startup_program):
            # Add Var k to main prog and startup prog
7214
            k = paddle.static.create_global_var(
7215 7216 7217 7218 7219 7220
                name="lookahead_k",
                shape=[1],
                value=int(self.k),
                dtype='int32',
                persistable=True,
            )
M
mapingshuo 已提交
7221

7222
            # Add Var alpha to main prog and startup prog
7223
            alpha = paddle.static.create_global_var(
7224 7225 7226 7227 7228 7229
                name="lookahead_alpha",
                shape=[1],
                value=float(self.alpha),
                dtype='float32',
                persistable=True,
            )
M
mapingshuo 已提交
7230

7231
            # Add Var step
7232
            step = paddle.static.create_global_var(
7233 7234 7235 7236 7237 7238
                name="lookahead_step",
                shape=[1],
                value=int(0),
                dtype='int32',
                persistable=True,
            )
7239
            paddle.increment(x=step, value=1.0)
7240 7241

            # lookahead
7242 7243 7244
            zero_var = layers.fill_constant(
                shape=[1], dtype='float32', value=0.0
            )
7245

7246 7247 7248
            one_var = layers.fill_constant(
                shape=[1], dtype='float32', value=1.0
            )
7249

7250
            mod = paddle.remainder(step, k)
7251
            with layers.control_flow.Switch() as switch:
7252 7253 7254 7255 7256
                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)
7257 7258 7259 7260
                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]
7261 7262 7263 7264
                        tmp_var = paddle.add(
                            paddle.multiply(fast_var, alpha),
                            paddle.multiply(
                                slow_var, paddle.subtract(one_var, alpha)
7265 7266
                            ),
                        )
7267 7268 7269 7270
                        layers.assign(input=tmp_var, output=slow_var)
                        layers.assign(input=tmp_var, output=fast_var)
                with switch.default():
                    pass
M
mapingshuo 已提交
7271
        return mini_out
7272 7273


7274
class GradientMergeOptimizer:
7275 7276 7277 7278 7279 7280 7281 7282 7283 7284 7285 7286 7287 7288 7289 7290 7291 7292 7293 7294 7295 7296
    """
    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

7297
        import paddle
7298 7299 7300 7301 7302 7303 7304 7305 7306 7307
        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')
7308 7309 7310 7311 7312
            cost = paddle.nn.functional.cross_entropy(
                input=prediction, label=input_y,
                reduction='none', use_softmax=False
            )
            sum_cost = paddle.mean(cost)
7313 7314 7315 7316 7317 7318 7319 7320 7321 7322 7323 7324 7325 7326 7327 7328 7329 7330 7331 7332
            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]))
    """

7333 7334
    GRAD_MERGE_COND_NAME = "grad_merge_cond_name"

7335
    def __init__(self, inner_optimizer, k_steps=1, avg=True):
姜永久 已提交
7336
        if in_dygraph_mode():
7337 7338 7339
            raise Exception(
                "In dygraph, we don't support GradientMergeOptimizer."
                "You can do Gradient merge by yourself with k-times forward + backward, "
7340 7341
                "and one-time optimizer.minimize()"
            )
7342

7343 7344 7345 7346
        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"
7347 7348 7349 7350 7351

        self.inner_optimizer = inner_optimizer
        self.k_steps = k_steps
        self.type = "gradient_merge"
        self.avg = avg
7352
        self._optimize_ops = None
7353

7354 7355 7356 7357 7358 7359
    def _set_k_steps(self, k_steps):
        self.k_steps = k_steps

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

7360 7361 7362 7363 7364 7365 7366 7367
    def backward(
        self,
        loss,
        startup_program=None,
        parameter_list=None,
        no_grad_set=None,
        callbacks=None,
    ):
7368 7369 7370 7371 7372 7373 7374 7375 7376
        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(
7377 7378
            loss, startup_program=startup_program
        )
7379 7380 7381 7382 7383 7384 7385 7386 7387 7388 7389
        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
7390 7391 7392
        if op_maker.kOpRoleVarAttrName() in op.attr_names and int(
            op.all_attrs()[op_maker.kOpRoleAttrName()]
        ) == int(backward):
7393 7394 7395 7396 7397 7398
            return True
        return False

    def _remove_op_role_var(self, param, grad):
        op_maker = core.op_proto_and_checker_maker
        op = grad.op
7399 7400 7401 7402 7403
        assert self._is_the_backward_op(
            op
        ), 'grad.op={} is not the backward op which produces the grad={}'.format(
            op, grad.name
        )
7404 7405 7406

        block = grad.block
        var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
7407 7408 7409 7410 7411 7412 7413 7414 7415 7416
        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
        )
7417 7418 7419 7420 7421 7422 7423 7424 7425 7426 7427 7428 7429 7430 7431 7432 7433 7434 7435 7436 7437 7438 7439 7440 7441 7442

        # 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
7443
        k_step_var = paddle.static.create_global_var(
7444 7445 7446 7447 7448 7449 7450 7451
            name="gradient_merge_k",
            shape=[1],
            value=int(self.k_steps),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )

7452
        zero_var = paddle.static.create_global_var(
7453 7454 7455 7456 7457 7458 7459
            name="gradient_merge_zero",
            shape=[1],
            value=int(0),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )
7460 7461

        # Add step var & cond var
7462
        step_var = paddle.static.create_global_var(
7463 7464 7465 7466 7467 7468 7469
            name="gradient_merge_step",
            shape=[1],
            value=int(0),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )
7470

7471 7472 7473
        cond_var = main_block.create_var(
            name="gradient_merge_cond", shape=[1], dtype='bool'
        )
7474 7475 7476

        with device_guard("cpu"):
            # step_var = (step_var + 1) % k_step
7477
            paddle.increment(x=step_var, value=1.0)
7478 7479 7480 7481 7482 7483
            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},
            )
7484 7485

            # cond_var = (step_var == 0)
7486 7487 7488 7489 7490
            main_block.append_op(
                type='equal',
                inputs={'X': step_var, 'Y': zero_var},
                outputs={'Out': cond_var},
            )
7491 7492 7493 7494 7495 7496 7497 7498 7499 7500

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

7502
        # TODO(mapingshuo) support sparse embedding
7503 7504
        # step1: remove grad.op's op_role_var
        for param, grad in params_grads:
7505
            assert (
7506
                param.type != core.VarDesc.VarType.SELECTED_ROWS
7507 7508
            ), "SELECTED_ROWS is not supported in GradientMergeOptimizer for now"

7509
            self._remove_op_role_var(param, grad)
7510

7511
        param_to_grad = {k.name: v for (k, v) in params_grads}
7512 7513 7514
        param_names = param_to_grad.keys()
        param_to_gradient_merge = {}

7515 7516 7517 7518 7519
        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
7520
            param_var = main_block.var(param_name)
7521 7522 7523 7524 7525 7526 7527
            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,
            )
7528
            param_to_gradient_merge[param_name] = gradient_merge_var
7529

7530 7531 7532 7533
            startup_gradient_merge_var = startup_block.create_var(
                name=param_name + "@GRAD@GradientMerge",
                shape=param_var.shape,
                dtype=param_var.dtype,
7534 7535 7536 7537 7538 7539 7540 7541 7542 7543 7544
                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),
                },
            )
7545

7546 7547 7548
            # grad_merge += grad
            new_grad_op = main_block.append_op(
                type="elementwise_add",
7549
                inputs={'X': grad, 'Y': gradient_merge_var},
7550
                outputs={'Out': gradient_merge_var},
7551 7552 7553 7554 7555
                attrs={'axis': -1, 'use_mkldnn': False},
            )
            self._add_gm_op_role_var(
                new_grad_op, param, gradient_merge_var, cond
            )
7556 7557 7558 7559 7560 7561 7562 7563
            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)
7564
            op_maker = core.op_proto_and_checker_maker
7565 7566 7567 7568

            if self.avg:
                for param, new_grad in new_params_grads:
                    # grad /= k_steps
7569 7570 7571 7572 7573 7574 7575 7576 7577 7578 7579 7580 7581
                    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
                    )
7582

7583 7584 7585 7586 7587 7588
            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
7589

7590
            self._optimize_ops = self.inner_optimizer.apply_gradients(
7591 7592
                new_params_grads
            )
7593

7594 7595
            # clear gradient_merge_vars
            for param, new_grad in new_params_grads:
7596 7597 7598 7599 7600 7601 7602 7603 7604
                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
                )
7605 7606

        # step3. apply gradient
7607
        paddle.static.nn.cond(cond, true_fn=true_apply_gradient, false_fn=None)
7608 7609 7610

        return self._optimize_ops

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

7616 7617 7618 7619 7620 7621
        params_grads = self.backward(
            loss,
            startup_program=startup_program,
            parameter_list=parameter_list,
            no_grad_set=no_grad_set,
        )
7622

7623 7624 7625
        optimize_ops = self.apply_optimize(
            loss, startup_program=startup_program, params_grads=params_grads
        )
7626 7627

        return optimize_ops, params_grads