optimization.py 11.9 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
# coding=utf-8
thomwolf's avatar
thomwolf committed
2
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
thomwolf's avatar
thomwolf committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch optimization for BERT model."""

17
18
19
import math
import torch
from torch.optim import Optimizer
Li Li's avatar
Li Li committed
20
from torch.optim.optimizer import required
21
from torch.nn.utils import clip_grad_norm_
lukovnikov's avatar
lukovnikov committed
22
import logging
lukovnikov's avatar
lukovnikov committed
23
from abc import ABC, abstractmethod
lukovnikov's avatar
lukovnikov committed
24
25

logger = logging.getLogger(__name__)
26

lukovnikov's avatar
lukovnikov committed
27

lukovnikov's avatar
lukovnikov committed
28
class _LRSchedule(ABC):
lukovnikov's avatar
lukovnikov committed
29
30
    """ Parent of all LRSchedules here. """
    warn_t_total = False        # is set to True for schedules where progressing beyond t_total steps doesn't make sense
lukovnikov's avatar
lukovnikov committed
31
    def __init__(self, warmup=0.002, t_total=-1, **kw):
lukovnikov's avatar
lukovnikov committed
32
33
34
35
36
        """
        :param warmup:  what fraction of t_total steps will be used for linear warmup
        :param t_total: how many training steps (updates) are planned
        :param kw:
        """
lukovnikov's avatar
lukovnikov committed
37
        super(_LRSchedule, self).__init__(**kw)
lukovnikov's avatar
lukovnikov committed
38
        if t_total < 0:
lukovnikov's avatar
lukovnikov committed
39
40
41
            logger.warning("t_total value of {} results in schedule not being applied".format(t_total))
        if not 0.0 <= warmup < 1.0 and not warmup == -1:
            raise ValueError("Invalid warmup: {} - should be in [0.0, 1.0[ or -1".format(warmup))
lukovnikov's avatar
lukovnikov committed
42
43
        warmup = max(warmup, 0.)
        self.warmup, self.t_total = float(warmup), float(t_total)
lukovnikov's avatar
lukovnikov committed
44
45
46
        self.warned_for_t_total_at_progress = -1

    def get_lr(self, step, nowarn=False):
lukovnikov's avatar
lukovnikov committed
47
48
49
50
51
        """
        :param step:    which of t_total steps we're on
        :param nowarn:  set to True to suppress warning regarding training beyond specified 't_total' steps
        :return:        learning rate multiplier for current update
        """
lukovnikov's avatar
lukovnikov committed
52
53
        if self.t_total < 0:
            return 1.
lukovnikov's avatar
lukovnikov committed
54
        progress = float(step) / self.t_total
lukovnikov's avatar
lukovnikov committed
55
56
57
58
59
60
61
62
63
64
        ret = self.get_lr_(progress)
        # warning for exceeding t_total (only active with warmup_linear
        if not nowarn and self.warn_t_total and progress > 1. and progress > self.warned_for_t_total_at_progress:
            logger.warning(
                "Training beyond specified 't_total'. Learning rate multiplier set to {}. Please set 't_total' of {} correctly."
                    .format(ret, self.__class__.__name__))
            self.warned_for_t_total_at_progress = progress
        # end warning
        return ret

lukovnikov's avatar
lukovnikov committed
65
    @abstractmethod
lukovnikov's avatar
lukovnikov committed
66
67
68
69
70
    def get_lr_(self, progress):
        """
        :param progress:    value between 0 and 1 (unless going beyond t_total steps) specifying training progress
        :return:            learning rate multiplier for current update
        """
lukovnikov's avatar
lukovnikov committed
71
72
73
        return 1.


lukovnikov's avatar
lukovnikov committed
74
75
76
77
78
79
class ConstantLR(_LRSchedule):
    def get_lr_(self, progress):
        return 1.


class WarmupCosineSchedule(_LRSchedule):
lukovnikov's avatar
lukovnikov committed
80
81
82
    """
    Cosine learning rate schedule with linear warmup. Cosine after warmup is without restarts.
    """
lukovnikov's avatar
lukovnikov committed
83
84
    warn_t_total = True
    def __init__(self, warmup=0.002, t_total=-1, cycles=.5, **kw):
lukovnikov's avatar
lukovnikov committed
85
86
87
88
89
90
        """
        :param warmup:      see LRSchedule
        :param t_total:     see LRSchedule
        :param cycles:      number of cycles. Default: 0.5, corresponding to cosine decay from 1. at progress==warmup and 0 at progress==1.
        :param kw:
        """
lukovnikov's avatar
lukovnikov committed
91
92
93
94
95
96
97
98
        super(WarmupCosineSchedule, self).__init__(warmup=warmup, t_total=t_total, **kw)
        self.cycles = cycles

    def get_lr_(self, progress):
        if progress < self.warmup:
            return progress / self.warmup
        else:
            progress = (progress - self.warmup) / (1 - self.warmup)   # progress after warmup
lukovnikov's avatar
lukovnikov committed
99
            return 0.5 * (1. + math.cos(math.pi * self.cycles * 2 * progress))
lukovnikov's avatar
lukovnikov committed
100
101


lukovnikov's avatar
lukovnikov committed
102
103
104
105
class WarmupCosineWithHardRestartsSchedule(WarmupCosineSchedule):
    """
    Cosine learning rate schedule with linear warmup and hard restarts (if cycles > 1).
    """
lukovnikov's avatar
lukovnikov committed
106
    def __init__(self, warmup=0.002, t_total=-1, cycles=1., **kw):
lukovnikov's avatar
lukovnikov committed
107
        super(WarmupCosineWithHardRestartsSchedule, self).__init__(warmup=warmup, t_total=t_total, cycles=cycles, **kw)
lukovnikov's avatar
lukovnikov committed
108
        assert(cycles >= 1.)
lukovnikov's avatar
lukovnikov committed
109
110
111
112
113
114

    def get_lr_(self, progress):
        if progress < self.warmup:
            return progress / self.warmup
        else:
            progress = (progress - self.warmup) / (1 - self.warmup)     # progress after warmup
lukovnikov's avatar
lukovnikov committed
115
            ret = 0.5 * (1. + math.cos(math.pi * ((self.cycles * progress) % 1)))
lukovnikov's avatar
lukovnikov committed
116
            return ret
lukovnikov's avatar
lukovnikov committed
117
118


lukovnikov's avatar
lukovnikov committed
119
120
121
122
123
124
125
126
class WarmupCosineWithWarmupRestartsSchedule(WarmupCosineWithHardRestartsSchedule):
    """
    Cosine learning rate schedule with linear warmups and linear warmup restarts.
    The same warmup rate is used for warmup restarts as for initial warmup.
    The total effective fraction of warmup steps over all cycles is warmup * cycles!
    """
    def __init__(self, warmup=0.002, t_total=-1, cycles=1., **kw):
        assert(warmup * cycles < 1.)
lukovnikov's avatar
lukovnikov committed
127
128
        warmup = warmup * cycles if warmup >= 0 else warmup
        super(WarmupCosineWithWarmupRestartsSchedule, self).__init__(warmup=warmup, t_total=t_total, cycles=cycles, **kw)
lukovnikov's avatar
lukovnikov committed
129

lukovnikov's avatar
lukovnikov committed
130
131
132
133
134
135
136
137
138
139
    def get_lr_(self, progress):
        progress = progress * self.cycles % 1.
        if progress < self.warmup:
            return progress / self.warmup
        else:
            progress = (progress - self.warmup) / (1 - self.warmup)     # progress after warmup
            ret = 0.5 * (1. + math.cos(math.pi * progress))
            return ret


lukovnikov's avatar
lukovnikov committed
140
class WarmupConstantSchedule(_LRSchedule):
lukovnikov's avatar
lukovnikov committed
141
142
143
    """
    Applies linear warmup. After warmup always returns 1..
    """
lukovnikov's avatar
lukovnikov committed
144
145
146
147
148
149
    def get_lr_(self, progress):
        if progress < self.warmup:
            return progress / self.warmup
        return 1.


lukovnikov's avatar
lukovnikov committed
150
class WarmupLinearSchedule(_LRSchedule):
lukovnikov's avatar
lukovnikov committed
151
152
153
    """
    Linear warmup. Linear decay after warmup.
    """
lukovnikov's avatar
lukovnikov committed
154
155
156
157
    warn_t_total = True
    def get_lr_(self, progress):
        if progress < self.warmup:
            return progress / self.warmup
lukovnikov's avatar
lukovnikov committed
158
        return max((progress - 1.) / (self.warmup - 1.), 0.)
lukovnikov's avatar
lukovnikov committed
159

160
161

SCHEDULES = {
lukovnikov's avatar
lukovnikov committed
162
163
    None:       ConstantLR,
    "none":     ConstantLR,
lukovnikov's avatar
lukovnikov committed
164
165
166
    "warmup_cosine": WarmupCosineSchedule,
    "warmup_constant": WarmupConstantSchedule,
    "warmup_linear": WarmupLinearSchedule
167
168
169
}


thomwolf's avatar
thomwolf committed
170
class BertAdam(Optimizer):
thomwolf's avatar
thomwolf committed
171
    """Implements BERT version of Adam algorithm with weight decay fix.
thomwolf's avatar
thomwolf committed
172
    Params:
thomwolf's avatar
thomwolf committed
173
174
175
        lr: learning rate
        warmup: portion of t_total for the warmup, -1  means no warmup. Default: -1
        t_total: total number of training steps for the learning
lukovnikov's avatar
lukovnikov committed
176
            rate schedule, -1  means constant learning rate of 1. (no warmup regardless of warmup setting). Default: -1
lukovnikov's avatar
lukovnikov committed
177
178
179
        schedule: schedule to use for the warmup (see above).
            Can be 'warmup_linear', 'warmup_constant', 'warmup_cosine', or a LRSchedule object.
            Default: 'warmup_linear'
thomwolf's avatar
thomwolf committed
180
181
182
        b1: Adams b1. Default: 0.9
        b2: Adams b2. Default: 0.999
        e: Adams epsilon. Default: 1e-6
183
        weight_decay: Weight decay. Default: 0.01
thomwolf's avatar
thomwolf committed
184
        max_grad_norm: Maximum norm for the gradients (-1 means no clipping). Default: 1.0
185
    """
Li Li's avatar
Li Li committed
186
    def __init__(self, params, lr=required, warmup=-1, t_total=-1, schedule='warmup_linear',
lukovnikov's avatar
lukovnikov committed
187
                 b1=0.9, b2=0.999, e=1e-6, weight_decay=0.01, max_grad_norm=1.0, **kwargs):
Li Li's avatar
Li Li committed
188
        if lr is not required and lr < 0.0:
thomwolf's avatar
thomwolf committed
189
            raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
lukovnikov's avatar
lukovnikov committed
190
        if not isinstance(schedule, _LRSchedule) and schedule not in SCHEDULES:
191
192
            raise ValueError("Invalid schedule parameter: {}".format(schedule))
        if not 0.0 <= b1 < 1.0:
thomwolf's avatar
thomwolf committed
193
            raise ValueError("Invalid b1 parameter: {} - should be in [0.0, 1.0[".format(b1))
194
        if not 0.0 <= b2 < 1.0:
thomwolf's avatar
thomwolf committed
195
196
197
            raise ValueError("Invalid b2 parameter: {} - should be in [0.0, 1.0[".format(b2))
        if not e >= 0.0:
            raise ValueError("Invalid epsilon value: {} - should be >= 0.0".format(e))
lukovnikov's avatar
lukovnikov committed
198
        # initialize schedule object
lukovnikov's avatar
lukovnikov committed
199
        if not isinstance(schedule, _LRSchedule):
lukovnikov's avatar
lukovnikov committed
200
201
202
203
            schedule_type = SCHEDULES[schedule]
            schedule = schedule_type(warmup=warmup, t_total=t_total)
        else:
            if warmup != -1 or t_total != -1:
lukovnikov's avatar
lukovnikov committed
204
205
                logger.warning("Non-default warmup and t_total are ineffective when LRSchedule object is provided. "
                               "Please specify custom warmup and t_total in LRSchedule object.")
lukovnikov's avatar
lukovnikov committed
206
        defaults = dict(lr=lr, schedule=schedule,
lukovnikov's avatar
lukovnikov committed
207
                        b1=b1, b2=b2, e=e, weight_decay=weight_decay,
208
                        max_grad_norm=max_grad_norm)
thomwolf's avatar
thomwolf committed
209
        super(BertAdam, self).__init__(params, defaults)
210
211
212
213
214
215
216
217

    def get_lr(self):
        lr = []
        for group in self.param_groups:
            for p in group['params']:
                state = self.state[p]
                if len(state) == 0:
                    return [0]
lukovnikov's avatar
lukovnikov committed
218
                lr_scheduled = group['lr']
lukovnikov's avatar
lukovnikov committed
219
                lr_scheduled *= group['schedule'].get_lr(state['step'])
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
                lr.append(lr_scheduled)
        return lr

    def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state['step'] = 0
                    # Exponential moving average of gradient values
thomwolf's avatar
thomwolf committed
248
                    state['next_m'] = torch.zeros_like(p.data)
249
                    # Exponential moving average of squared gradient values
thomwolf's avatar
thomwolf committed
250
                    state['next_v'] = torch.zeros_like(p.data)
251

thomwolf's avatar
thomwolf committed
252
                next_m, next_v = state['next_m'], state['next_v']
253
254
255
256
257
258
259
                beta1, beta2 = group['b1'], group['b2']

                # Add grad clipping
                if group['max_grad_norm'] > 0:
                    clip_grad_norm_(p, group['max_grad_norm'])

                # Decay the first and second moment running average coefficient
thomwolf's avatar
thomwolf committed
260
261
262
263
                # In-place operations to update the averages at the same time
                next_m.mul_(beta1).add_(1 - beta1, grad)
                next_v.mul_(beta2).addcmul_(1 - beta2, grad, grad)
                update = next_m / (next_v.sqrt() + group['e'])
264
265
266
267
268

                # Just adding the square of the weights to the loss function is *not*
                # the correct way of using L2 regularization/weight decay with Adam,
                # since that will interact with the m and v parameters in strange ways.
                #
thomwolf's avatar
thomwolf committed
269
                # Instead we want to decay the weights in a manner that doesn't interact
270
271
                # with the m/v parameters. This is equivalent to adding the square
                # of the weights to the loss with plain (non-momentum) SGD.
272
273
                if group['weight_decay'] > 0.0:
                    update += group['weight_decay'] * p.data
thomwolf's avatar
thomwolf committed
274

lukovnikov's avatar
lukovnikov committed
275
                lr_scheduled = group['lr']
lukovnikov's avatar
lukovnikov committed
276
                lr_scheduled *= group['schedule'].get_lr(state['step'])
thomwolf's avatar
thomwolf committed
277
278
279
280
281
282
283

                update_with_lr = lr_scheduled * update
                p.data.add_(-update_with_lr)

                state['step'] += 1

                # step_size = lr_scheduled * math.sqrt(bias_correction2) / bias_correction1
thomwolf's avatar
thomwolf committed
284
                # No bias correction
thomwolf's avatar
thomwolf committed
285
286
                # bias_correction1 = 1 - beta1 ** state['step']
                # bias_correction2 = 1 - beta2 ** state['step']
287
288

        return loss