optimization.py 8.14 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."""

thomwolf's avatar
thomwolf committed
17
import logging
18
import math
thomwolf's avatar
thomwolf committed
19

20
21
import torch
from torch.optim import Optimizer
thomwolf's avatar
thomwolf committed
22
from torch.optim.lr_scheduler import LambdaLR
lukovnikov's avatar
lukovnikov committed
23
24

logger = logging.getLogger(__name__)
25

thomwolf's avatar
thomwolf committed
26
class ConstantLRSchedule(LambdaLR):
thomwolf's avatar
thomwolf committed
27
28
    """ Constant learning rate schedule.
    """
thomwolf's avatar
thomwolf committed
29
    def __init__(self, optimizer, last_epoch=-1):
thomwolf's avatar
thomwolf committed
30
        super(ConstantLRSchedule, self).__init__(optimizer, lambda _: 1.0, last_epoch=last_epoch)
lukovnikov's avatar
lukovnikov committed
31

thomwolf's avatar
thomwolf committed
32
33
34
35
36

class WarmupConstantSchedule(LambdaLR):
    """ Linear warmup and then constant.
        Linearly increases learning rate schedule from 0 to 1 over `warmup_steps` training steps.
        Keeps learning rate schedule equal to 1. after warmup_steps.
lukovnikov's avatar
lukovnikov committed
37
    """
thomwolf's avatar
thomwolf committed
38
    def __init__(self, optimizer, warmup_steps, last_epoch=-1):
lukovnikov's avatar
lukovnikov committed
39

thomwolf's avatar
thomwolf committed
40
41
        def lr_lambda(step):
            if step < warmup_steps:
thomwolf's avatar
thomwolf committed
42
                return float(step) / float(max(1.0, warmup_steps))
thomwolf's avatar
thomwolf committed
43
            return 1.
lukovnikov's avatar
lukovnikov committed
44

thomwolf's avatar
thomwolf committed
45
        super(WarmupConstantSchedule, self).__init__(optimizer, lr_lambda, last_epoch=last_epoch)
lukovnikov's avatar
lukovnikov committed
46

thomwolf's avatar
thomwolf committed
47
48
49
50
51

class WarmupLinearSchedule(LambdaLR):
    """ Linear warmup and then linear decay.
        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
        Linearly decreases learning rate from 1. to 0. over remaining `t_total - warmup_steps` steps.
lukovnikov's avatar
lukovnikov committed
52
    """
thomwolf's avatar
thomwolf committed
53
    def __init__(self, optimizer, warmup_steps, t_total, last_epoch=-1):
lukovnikov's avatar
lukovnikov committed
54

thomwolf's avatar
thomwolf committed
55
56
        def lr_lambda(step):
            if step < warmup_steps:
thomwolf's avatar
thomwolf committed
57
                return float(step) / float(max(1, warmup_steps))
thomwolf's avatar
thomwolf committed
58
            return max(0.0, float(t_total - step) / float(max(1.0, t_total - warmup_steps)))
lukovnikov's avatar
lukovnikov committed
59

thomwolf's avatar
thomwolf committed
60
        super(WarmupLinearSchedule, self).__init__(optimizer, lr_lambda, last_epoch=last_epoch)
lukovnikov's avatar
lukovnikov committed
61

thomwolf's avatar
thomwolf committed
62

thomwolf's avatar
thomwolf committed
63
64
65
66
67
class WarmupCosineSchedule(LambdaLR):
    """ Linear warmup and then cosine decay.
        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
        Decreases learning rate from 1. to 0. over remaining `t_total - warmup_steps` steps following a cosine curve.
        If `cycles` (default=0.5) is different from default, learning rate follows cosine function after warmup.
lukovnikov's avatar
lukovnikov committed
68
    """
thomwolf's avatar
thomwolf committed
69
70
    warn_t_total = True
    def __init__(self, optimizer, warmup_steps, t_total, cycles=.5, last_epoch=-1):
thomwolf's avatar
thomwolf committed
71
72
73

        def lr_lambda(step):
            if step < warmup_steps:
thomwolf's avatar
thomwolf committed
74
                return float(step) / float(max(1.0, warmup_steps))
thomwolf's avatar
thomwolf committed
75
76
77
            else:
                progress = float(step - warmup_steps) / float(max(1, t_total - warmup_steps))   # progress after warmup
                return max(0.0, 0.5 * (1. + math.cos(math.pi * float(cycles) * 2.0 * progress)))
lukovnikov's avatar
lukovnikov committed
78

thomwolf's avatar
thomwolf committed
79
        super(WarmupCosineSchedule, self).__init__(optimizer, lr_lambda, last_epoch=last_epoch)
thomwolf's avatar
thomwolf committed
80

thomwolf's avatar
thomwolf committed
81
82
83
84
85
class WarmupCosineWithHardRestartsSchedule(LambdaLR):
    """ Linear warmup and then cosine cycles with hard restarts.
        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
        If `cycles` (default=1.) is different from default, learning rate follows `cycles` times a cosine decaying
        learning rate (with hard restarts).
lukovnikov's avatar
lukovnikov committed
86
    """
thomwolf's avatar
thomwolf committed
87
    def __init__(self, optimizer, warmup_steps, t_total, cycles=1., last_epoch=-1):
lukovnikov's avatar
lukovnikov committed
88

thomwolf's avatar
thomwolf committed
89
90
        def lr_lambda(step):
            if step < warmup_steps:
thomwolf's avatar
thomwolf committed
91
                return float(step) / float(max(1, warmup_steps))
thomwolf's avatar
thomwolf committed
92
93
94
95
96
            else:
                progress = float(step - warmup_steps) / float(max(1, t_total - warmup_steps))   # progress after warmup
                if progress >= 1.0:
                    return 0.0
                return max(0.0, 0.5 * (1. + math.cos(math.pi * ((float(cycles) * progress) % 1.0))))
97

thomwolf's avatar
thomwolf committed
98
        super(WarmupCosineWithHardRestartsSchedule, self).__init__(optimizer, lr_lambda, last_epoch=last_epoch)
99
100


thomwolf's avatar
thomwolf committed
101
102
class AdamW(Optimizer):
    """ Implements Adam algorithm with weight decay fix.
103
104

    Parameters:
thomwolf's avatar
thomwolf committed
105
106
107
108
109
        lr (float): learning rate. Default 1e-3.
        betas (tuple of 2 floats): Adams beta parameters (b1, b2). Default: (0.9, 0.999)
        eps (float): Adams epsilon. Default: 1e-6
        weight_decay (float): Weight decay. Default: 0.0
        correct_bias (bool): can be set to False to avoid correcting bias in Adam (e.g. like in Bert TF repository). Default True.
110
    """
thomwolf's avatar
thomwolf committed
111
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.0, correct_bias=True):
thomwolf's avatar
thomwolf committed
112
        if lr < 0.0:
thomwolf's avatar
thomwolf committed
113
            raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
thomwolf's avatar
thomwolf committed
114
115
116
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter: {} - should be in [0.0, 1.0[".format(betas[0]))
        if not 0.0 <= betas[1]  < 1.0:
thomwolf's avatar
thomwolf committed
117
            raise ValueError("Invalid beta parameter: {} - should be in [0.0, 1.0[".format(betas[1]))
thomwolf's avatar
thomwolf committed
118
        if not 0.0 <= eps:
thomwolf's avatar
thomwolf committed
119
            raise ValueError("Invalid epsilon value: {} - should be >= 0.0".format(eps))
thomwolf's avatar
thomwolf committed
120
121
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
                        correct_bias=correct_bias)
thomwolf's avatar
thomwolf committed
122
        super(AdamW, self).__init__(params, defaults)
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

    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
149
                    state['exp_avg'] = torch.zeros_like(p.data)
150
                    # Exponential moving average of squared gradient values
thomwolf's avatar
thomwolf committed
151
                    state['exp_avg_sq'] = torch.zeros_like(p.data)
152

thomwolf's avatar
thomwolf committed
153
154
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                beta1, beta2 = group['betas']
155

thomwolf's avatar
thomwolf committed
156
                state['step'] += 1
157
158

                # Decay the first and second moment running average coefficient
thomwolf's avatar
thomwolf committed
159
                # In-place operations to update the averages at the same time
thomwolf's avatar
thomwolf committed
160
161
                exp_avg.mul_(beta1).add_(1.0 - beta1, grad)
                exp_avg_sq.mul_(beta2).addcmul_(1.0 - beta2, grad, grad)
thomwolf's avatar
thomwolf committed
162
163
164
165
                denom = exp_avg_sq.sqrt().add_(group['eps'])

                step_size = group['lr']
                if group['correct_bias']:  # No bias correction for Bert
thomwolf's avatar
thomwolf committed
166
167
                    bias_correction1 = 1.0 - beta1 ** state['step']
                    bias_correction2 = 1.0 - beta2 ** state['step']
thomwolf's avatar
thomwolf committed
168
169
170
                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1

                p.data.addcdiv_(-step_size, exp_avg, denom)
171
172
173
174
175

                # 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
176
                # Instead we want to decay the weights in a manner that doesn't interact
177
178
                # with the m/v parameters. This is equivalent to adding the square
                # of the weights to the loss with plain (non-momentum) SGD.
thomwolf's avatar
thomwolf committed
179
                # Add weight decay at the end (fixed version)
thomwolf's avatar
thomwolf committed
180
                if group['weight_decay'] > 0.0:
thomwolf's avatar
thomwolf committed
181
                    p.data.add_(-group['lr'] * group['weight_decay'], p.data)
182
183

        return loss