optimization.py 7.56 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HugginFace Inc. team.
#
# 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
23
24
import logging

logger = logging.getLogger(__name__)
25
26

def warmup_cosine(x, warmup=0.002):
thomwolf's avatar
thomwolf committed
27
28
29
    if x < warmup:
        return x/warmup
    return 0.5 * (1.0 + torch.cos(math.pi * x))
30
31

def warmup_constant(x, warmup=0.002):
32
33
    """ Linearly increases learning rate over `warmup`*`t_total` (as provided to BertAdam) training steps.
        Learning rate is 1. afterwards. """
thomwolf's avatar
thomwolf committed
34
35
36
    if x < warmup:
        return x/warmup
    return 1.0
37

lukovnikov's avatar
lukovnikov committed
38
39
40
41
42
43
44
45
46
47
48
class Warmup_Linear_with_Warning(object):
    def __init__(self, **kw):
        super(Warmup_Linear_with_Warning, self).__init__()
        self.warned_at_x = -1

    def __call__(self, x, warmup=0.002):
        if x > 1 and x > self.warned_at_x:
            logger.warning("Training beyond specified 't_total' steps. Learning rate set to zero. Please set 't_total' of BertAdam correctly.")
            self.warned_at_x = x
        return warmup_linear(x, warmup=warmup)

49
def warmup_linear(x, warmup=0.002):
50
51
    """ Specifies a triangular learning rate schedule where peak is reached at `warmup`*`t_total`-th (as provided to BertAdam) training step.
        After `t_total`-th training step, learning rate is zero. """
thomwolf's avatar
thomwolf committed
52
53
    if x < warmup:
        return x/warmup
54
    return max((x-1.)/(warmup-1.), 0)
55
56
57
58

SCHEDULES = {
    'warmup_cosine':warmup_cosine,
    'warmup_constant':warmup_constant,
lukovnikov's avatar
lukovnikov committed
59
    'warmup_linear': Warmup_Linear_with_Warning(), #warmup_linear,
60
61
62
}


thomwolf's avatar
thomwolf committed
63
class BertAdam(Optimizer):
thomwolf's avatar
thomwolf committed
64
    """Implements BERT version of Adam algorithm with weight decay fix.
thomwolf's avatar
thomwolf committed
65
    Params:
thomwolf's avatar
thomwolf committed
66
67
68
69
70
71
72
73
        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
            rate schedule, -1  means constant learning rate. Default: -1
        schedule: schedule to use for the warmup (see above). Default: 'warmup_linear'
        b1: Adams b1. Default: 0.9
        b2: Adams b2. Default: 0.999
        e: Adams epsilon. Default: 1e-6
74
        weight_decay: Weight decay. Default: 0.01
thomwolf's avatar
thomwolf committed
75
        max_grad_norm: Maximum norm for the gradients (-1 means no clipping). Default: 1.0
76
    """
Li Li's avatar
Li Li committed
77
    def __init__(self, params, lr=required, warmup=-1, t_total=-1, schedule='warmup_linear',
78
                 b1=0.9, b2=0.999, e=1e-6, weight_decay=0.01,
thomwolf's avatar
thomwolf committed
79
                 max_grad_norm=1.0):
Li Li's avatar
Li Li committed
80
        if lr is not required and lr < 0.0:
thomwolf's avatar
thomwolf committed
81
            raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
82
83
        if schedule not in SCHEDULES:
            raise ValueError("Invalid schedule parameter: {}".format(schedule))
thomwolf's avatar
thomwolf committed
84
85
        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))
86
        if not 0.0 <= b1 < 1.0:
thomwolf's avatar
thomwolf committed
87
            raise ValueError("Invalid b1 parameter: {} - should be in [0.0, 1.0[".format(b1))
88
        if not 0.0 <= b2 < 1.0:
thomwolf's avatar
thomwolf committed
89
90
91
            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))
92
        defaults = dict(lr=lr, schedule=schedule, warmup=warmup, t_total=t_total,
93
                        b1=b1, b2=b2, e=e, weight_decay=weight_decay,
94
                        max_grad_norm=max_grad_norm)
thomwolf's avatar
thomwolf committed
95
        super(BertAdam, self).__init__(params, defaults)
96
97
98
99
100
101
102
103

    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]
thomwolf's avatar
thomwolf committed
104
105
106
107
108
                if group['t_total'] != -1:
                    schedule_fct = SCHEDULES[group['schedule']]
                    lr_scheduled = group['lr'] * schedule_fct(state['step']/group['t_total'], group['warmup'])
                else:
                    lr_scheduled = group['lr']
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
                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
137
                    state['next_m'] = torch.zeros_like(p.data)
138
                    # Exponential moving average of squared gradient values
thomwolf's avatar
thomwolf committed
139
                    state['next_v'] = torch.zeros_like(p.data)
140

thomwolf's avatar
thomwolf committed
141
                next_m, next_v = state['next_m'], state['next_v']
142
143
144
145
146
147
148
                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
149
150
151
152
                # 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'])
153
154
155
156
157

                # 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
158
                # Instead we want to decay the weights in a manner that doesn't interact
159
160
                # with the m/v parameters. This is equivalent to adding the square
                # of the weights to the loss with plain (non-momentum) SGD.
161
162
                if group['weight_decay'] > 0.0:
                    update += group['weight_decay'] * p.data
thomwolf's avatar
thomwolf committed
163
164
165
166
167
168
169
170
171
172
173
174
175

                if group['t_total'] != -1:
                    schedule_fct = SCHEDULES[group['schedule']]
                    lr_scheduled = group['lr'] * schedule_fct(state['step']/group['t_total'], group['warmup'])
                else:
                    lr_scheduled = group['lr']

                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
176
                # No bias correction
thomwolf's avatar
thomwolf committed
177
178
                # bias_correction1 = 1 - beta1 ** state['step']
                # bias_correction2 = 1 - beta2 ** state['step']
179
180

        return loss