"vscode:/vscode.git/clone" did not exist on "59414b332eccf5caa4af50d65fe9ed7a88a4a5c6"
optimization_openai.py 5.2 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
# coding=utf-8
thomwolf's avatar
thomwolf committed
2
# Copyright 2018 The Open AI 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 OpenAI GPT model."""

thomwolf's avatar
thomwolf committed
17
18
19
import math
import torch
from torch.optim import Optimizer
thomwolf's avatar
thomwolf committed
20
from torch.optim.optimizer import required
thomwolf's avatar
thomwolf committed
21
from torch.nn.utils import clip_grad_norm_
lukovnikov's avatar
lukovnikov committed
22
import logging
lukovnikov's avatar
lukovnikov committed
23
from .optimization import *
lukovnikov's avatar
lukovnikov committed
24
25

logger = logging.getLogger(__name__)
thomwolf's avatar
thomwolf committed
26
27
28
29
30


class OpenAIAdam(Optimizer):
    """Implements Open AI version of Adam algorithm with weight decay fix.
    """
thomwolf's avatar
thomwolf committed
31
32
    def __init__(self, params, lr=required, schedule='warmup_linear', warmup=-1, t_total=-1,
                 b1=0.9, b2=0.999, e=1e-8, weight_decay=0,
thomwolf's avatar
thomwolf committed
33
                 vector_l2=False, max_grad_norm=-1, **kwargs):
thomwolf's avatar
thomwolf committed
34
35
        if lr is not required and lr < 0.0:
            raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
lukovnikov's avatar
lukovnikov committed
36
        if not isinstance(schedule, LRSchedule) and schedule not in SCHEDULES:
thomwolf's avatar
thomwolf committed
37
38
            raise ValueError("Invalid schedule parameter: {}".format(schedule))
        if not 0.0 <= b1 < 1.0:
lukovnikov's avatar
lukovnikov committed
39
            raise ValueError("Invalid b1 parameter: {} - should be in [0.0, 1.0[".format(b1))
thomwolf's avatar
thomwolf committed
40
        if not 0.0 <= b2 < 1.0:
lukovnikov's avatar
lukovnikov committed
41
            raise ValueError("Invalid b2 parameter: {} - should be in [0.0, 1.0[".format(b2))
thomwolf's avatar
thomwolf committed
42
        if not e >= 0.0:
lukovnikov's avatar
lukovnikov committed
43
44
45
46
47
48
49
50
51
52
            raise ValueError("Invalid epsilon value: {} - should be >= 0.0".format(e))
        # initialize schedule object
        if not isinstance(schedule, LRSchedule):
            schedule_type = SCHEDULES[schedule]
            schedule = schedule_type(warmup=warmup, t_total=t_total)
        else:
            if warmup != -1 or t_total != -1:
                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.")
        defaults = dict(lr=lr, schedule=schedule,
thomwolf's avatar
thomwolf committed
53
                        b1=b1, b2=b2, e=e, weight_decay=weight_decay, vector_l2=vector_l2,
thomwolf's avatar
thomwolf committed
54
55
56
                        max_grad_norm=max_grad_norm)
        super(OpenAIAdam, self).__init__(params, defaults)

thomwolf's avatar
thomwolf committed
57
58
59
60
61
62
63
    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
64
65
                lr_scheduled = group['lr']
                lr_scheduled *= group['schedule'].get_lr(state['step'])
thomwolf's avatar
thomwolf committed
66
67
68
                lr.append(lr_scheduled)
        return lr

thomwolf's avatar
thomwolf committed
69
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
    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
                    state['exp_avg'] = torch.zeros_like(p.data)
                    # Exponential moving average of squared gradient values
                    state['exp_avg_sq'] = torch.zeros_like(p.data)

                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                beta1, beta2 = group['b1'], group['b2']

                state['step'] += 1

                # 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
                exp_avg.mul_(beta1).add_(1 - beta1, grad)
                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
                denom = exp_avg_sq.sqrt().add_(group['e'])

                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']

lukovnikov's avatar
lukovnikov committed
115
116
                lr_scheduled = group['lr']
                lr_scheduled *= group['schedule'].get_lr(state['step'])
thomwolf's avatar
thomwolf committed
117

thomwolf's avatar
thomwolf committed
118
119
120
121
122
                step_size = lr_scheduled * math.sqrt(bias_correction2) / bias_correction1

                p.data.addcdiv_(-step_size, exp_avg, denom)

                # Add weight decay at the end (fixed version)
thomwolf's avatar
thomwolf committed
123
124
                if (len(p.size()) > 1 or group['vector_l2']) and group['weight_decay'] > 0:
                    p.data.add_(-lr_scheduled * group['weight_decay'], p.data)
thomwolf's avatar
thomwolf committed
125
126

        return loss