optimizer_factory.py 8.11 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Hongkun Yu's avatar
Hongkun Yu committed
14

Abdullah Rashwan's avatar
Abdullah Rashwan committed
15
"""Optimizer factory class."""
Frederick Liu's avatar
Frederick Liu committed
16
from typing import Callable, Optional, Union, List, Tuple
Abdullah Rashwan's avatar
Abdullah Rashwan committed
17

Le Hou's avatar
Le Hou committed
18
import gin
Abdullah Rashwan's avatar
Abdullah Rashwan committed
19
import tensorflow as tf
Abdullah Rashwan's avatar
Abdullah Rashwan committed
20
import tensorflow_addons.optimizers as tfa_optimizers
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
21
from official.modeling.optimization import slide_optimizer
22
from official.modeling.optimization import adafactor_optimizer
Abdullah Rashwan's avatar
Abdullah Rashwan committed
23
from official.modeling.optimization import ema_optimizer
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
24
from official.modeling.optimization import lars_optimizer
Abdullah Rashwan's avatar
Abdullah Rashwan committed
25
26
27
28
29
30
from official.modeling.optimization import lr_schedule
from official.modeling.optimization.configs import optimization_config as opt_cfg
from official.nlp import optimization as nlp_optimization

OPTIMIZERS_CLS = {
    'sgd': tf.keras.optimizers.SGD,
31
    # TODO(chenmoneygithub): experimental.SGD
Abdullah Rashwan's avatar
Abdullah Rashwan committed
32
    'adam': tf.keras.optimizers.Adam,
33
    # TODO(chenmoneygithub): experimental.Adam
Abdullah Rashwan's avatar
Abdullah Rashwan committed
34
    'adamw': nlp_optimization.AdamWeightDecay,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
35
    'lamb': tfa_optimizers.LAMB,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
36
37
    'rmsprop': tf.keras.optimizers.RMSprop,
    'lars': lars_optimizer.LARS,
Hao Wu's avatar
Hao Wu committed
38
    'adagrad': tf.keras.optimizers.Adagrad,
39
40
    'slide': slide_optimizer.SLIDE,
    'adafactor': adafactor_optimizer.Adafactor,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
41
42
43
}

LR_CLS = {
44
45
46
47
    'stepwise': lr_schedule.PiecewiseConstantDecayWithOffset,
    'polynomial': lr_schedule.PolynomialDecayWithOffset,
    'exponential': lr_schedule.ExponentialDecayWithOffset,
    'cosine': lr_schedule.CosineDecayWithOffset,
48
    'power': lr_schedule.DirectPowerDecay,
Le Hou's avatar
Le Hou committed
49
    'power_linear': lr_schedule.PowerAndLinearDecay,
Le Hou's avatar
Le Hou committed
50
    'power_with_offset': lr_schedule.PowerDecayWithOffset,
51
    'step_cosine_with_offset': lr_schedule.StepCosineDecayWithOffset,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
52
53
54
}

WARMUP_CLS = {
Abdullah Rashwan's avatar
Abdullah Rashwan committed
55
56
    'linear': lr_schedule.LinearWarmup,
    'polynomial': lr_schedule.PolynomialWarmUp
Abdullah Rashwan's avatar
Abdullah Rashwan committed
57
58
59
}


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
60
61
def register_optimizer_cls(key: str,
                           optimizer_config_cls: tf.keras.optimizers.Optimizer):
62
63
64
65
66
67
68
69
70
71
72
73
74
75
  """Register customize optimizer cls.

  The user will still need to subclass data classes in
  configs.optimization_config to be used with OptimizerFactory.

  Args:
    key: A string to that the optimizer_config_cls is registered with.
    optimizer_config_cls: A class which inherits tf.keras.optimizers.Optimizer.
  """
  if key in OPTIMIZERS_CLS:
    raise ValueError('%s already registered in OPTIMIZER_CLS.' % key)
  OPTIMIZERS_CLS[key] = optimizer_config_cls


Hongkun Yu's avatar
Hongkun Yu committed
76
class OptimizerFactory:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
  """Optimizer factory class.

  This class builds learning rate and optimizer based on an optimization config.
  To use this class, you need to do the following:
  (1) Define optimization config, this includes optimizer, and learning rate
      schedule.
  (2) Initialize the class using the optimization config.
  (3) Build learning rate.
  (4) Build optimizer.

  This is a typical example for using this class:
  params = {
        'optimizer': {
            'type': 'sgd',
Abdullah Rashwan's avatar
Abdullah Rashwan committed
91
            'sgd': {'momentum': 0.9}
Abdullah Rashwan's avatar
Abdullah Rashwan committed
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        },
        'learning_rate': {
            'type': 'stepwise',
            'stepwise': {'boundaries': [10000, 20000],
                         'values': [0.1, 0.01, 0.001]}
        },
        'warmup': {
            'type': 'linear',
            'linear': {'warmup_steps': 500, 'warmup_learning_rate': 0.01}
        }
    }
  opt_config = OptimizationConfig(params)
  opt_factory = OptimizerFactory(opt_config)
  lr = opt_factory.build_learning_rate()
  optimizer = opt_factory.build_optimizer(lr)
  """

  def __init__(self, config: opt_cfg.OptimizationConfig):
    """Initializing OptimizerFactory.

    Args:
      config: OptimizationConfig instance contain optimization config.
    """
    self._config = config
    self._optimizer_config = config.optimizer.get()
    self._optimizer_type = config.optimizer.type

Abdullah Rashwan's avatar
Abdullah Rashwan committed
119
120
121
122
    self._use_ema = config.ema is not None
    self._ema_config = config.ema

    if self._optimizer_config is None:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
123
124
125
126
127
      raise ValueError('Optimizer type must be specified')

    self._lr_config = config.learning_rate.get()
    self._lr_type = config.learning_rate.type

Abdullah Rashwan's avatar
Abdullah Rashwan committed
128
129
130
    if self._lr_type is None:
      raise ValueError('Learning rate type must be specified')

Abdullah Rashwan's avatar
Abdullah Rashwan committed
131
132
133
134
135
136
137
    self._warmup_config = config.warmup.get()
    self._warmup_type = config.warmup.type

  def build_learning_rate(self):
    """Build learning rate.

    Builds learning rate from config. Learning rate schedule is built according
Abdullah Rashwan's avatar
Abdullah Rashwan committed
138
139
    to the learning rate config. If learning rate type is consant,
    lr_config.learning_rate is returned.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
140
141

    Returns:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
142
143
      tf.keras.optimizers.schedules.LearningRateSchedule instance. If
      learning rate type is consant, lr_config.learning_rate is returned.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
144
    """
Abdullah Rashwan's avatar
Abdullah Rashwan committed
145
146
    if self._lr_type == 'constant':
      lr = self._lr_config.learning_rate
Abdullah Rashwan's avatar
Abdullah Rashwan committed
147
148
149
150
151
152
153
154
    else:
      lr = LR_CLS[self._lr_type](**self._lr_config.as_dict())

    if self._warmup_config:
      lr = WARMUP_CLS[self._warmup_type](lr, **self._warmup_config.as_dict())

    return lr

Le Hou's avatar
Le Hou committed
155
  @gin.configurable
Abdullah Rashwan's avatar
Abdullah Rashwan committed
156
  def build_optimizer(
Le Hou's avatar
Le Hou committed
157
158
      self,
      lr: Union[tf.keras.optimizers.schedules.LearningRateSchedule, float],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
159
160
161
      gradient_aggregator: Optional[Callable[
          [List[Tuple[tf.Tensor, tf.Tensor]]], List[Tuple[tf.Tensor,
                                                          tf.Tensor]]]] = None,
Frederick Liu's avatar
Frederick Liu committed
162
      gradient_transformers: Optional[List[Callable[
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
163
164
          [List[Tuple[tf.Tensor, tf.Tensor]]], List[Tuple[tf.Tensor,
                                                          tf.Tensor]]]]] = None,
Rebecca Chen's avatar
Rebecca Chen committed
165
166
      postprocessor: Optional[Callable[[tf.keras.optimizers.Optimizer],
                                       tf.keras.optimizers.Optimizer]] = None):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
167
168
169
170
171
172
173
    """Build optimizer.

    Builds optimizer from config. It takes learning rate as input, and builds
    the optimizer according to the optimizer config. Typically, the learning
    rate built using self.build_lr() is passed as an argument to this method.

    Args:
Hongkun Yu's avatar
Hongkun Yu committed
174
175
      lr: A floating point value, or a
        tf.keras.optimizers.schedules.LearningRateSchedule instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
176
      gradient_aggregator: Optional function to overwrite gradient aggregation.
Frederick Liu's avatar
Frederick Liu committed
177
178
179
180
181
      gradient_transformers: Optional list of functions to use to transform
        gradients before applying updates to Variables. The functions are
        applied after gradient_aggregator. The functions should accept and
        return a list of (gradient, variable) tuples. clipvalue, clipnorm,
        global_clipnorm should not be set when gradient_transformers is passed.
Le Hou's avatar
Le Hou committed
182
183
      postprocessor: An optional function for postprocessing the optimizer. It
        takes an optimizer and returns an optimizer.
Hongkun Yu's avatar
Hongkun Yu committed
184

Abdullah Rashwan's avatar
Abdullah Rashwan committed
185
    Returns:
Chen Qian's avatar
Chen Qian committed
186
187
      `tf.keras.optimizers.Optimizer` or
      `tf.keras.optimizers.experimental.Optimizer` instance.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
188
189
190
    """

    optimizer_dict = self._optimizer_config.as_dict()
Frederick Liu's avatar
Frederick Liu committed
191
    ## Delete clipnorm, clipvalue, global_clipnorm if None
Abdullah Rashwan's avatar
Abdullah Rashwan committed
192
193
194
195
    if optimizer_dict['clipnorm'] is None:
      del optimizer_dict['clipnorm']
    if optimizer_dict['clipvalue'] is None:
      del optimizer_dict['clipvalue']
Frederick Liu's avatar
Frederick Liu committed
196
197
    if optimizer_dict['global_clipnorm'] is None:
      del optimizer_dict['global_clipnorm']
Abdullah Rashwan's avatar
Abdullah Rashwan committed
198

Abdullah Rashwan's avatar
Abdullah Rashwan committed
199
    optimizer_dict['learning_rate'] = lr
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
200
201
    if gradient_aggregator is not None:
      optimizer_dict['gradient_aggregator'] = gradient_aggregator
Frederick Liu's avatar
Frederick Liu committed
202
203
    if gradient_transformers is not None:
      optimizer_dict['gradient_transformers'] = gradient_transformers
Abdullah Rashwan's avatar
Abdullah Rashwan committed
204
205

    optimizer = OPTIMIZERS_CLS[self._optimizer_type](**optimizer_dict)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
206
207
208
209

    if self._use_ema:
      optimizer = ema_optimizer.ExponentialMovingAverage(
          optimizer, **self._ema_config.as_dict())
Le Hou's avatar
Le Hou committed
210
211
    if postprocessor:
      optimizer = postprocessor(optimizer)
Chen Qian's avatar
Chen Qian committed
212
213
214
215
    assert isinstance(
        optimizer, (tf.keras.optimizers.Optimizer,
                    tf.keras.optimizers.experimental.Optimizer)
    ), ('OptimizerFactory.build_optimizer returning a non-optimizer object: '
Le Hou's avatar
Le Hou committed
216
        '{}'.format(optimizer))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
217

Abdullah Rashwan's avatar
Abdullah Rashwan committed
218
    return optimizer