optimizer_factory.py 8.12 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
21

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
22
from official.modeling.optimization import slide_optimizer
23
from official.modeling.optimization import adafactor_optimizer
Abdullah Rashwan's avatar
Abdullah Rashwan committed
24
from official.modeling.optimization import ema_optimizer
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
25
from official.modeling.optimization import lars_optimizer
26
from official.modeling.optimization import legacy_adamw
Abdullah Rashwan's avatar
Abdullah Rashwan committed
27
28
29
30
31
from official.modeling.optimization import lr_schedule
from official.modeling.optimization.configs import optimization_config as opt_cfg

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

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

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


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
61
62
def register_optimizer_cls(key: str,
                           optimizer_config_cls: tf.keras.optimizers.Optimizer):
63
64
65
66
67
68
69
70
71
72
73
74
75
76
  """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
77
class OptimizerFactory:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
78
79
80
81
82
83
84
85
86
87
88
  """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:
Mark Daoust's avatar
Mark Daoust committed
89
90

  ```
Abdullah Rashwan's avatar
Abdullah Rashwan committed
91
92
93
  params = {
        'optimizer': {
            'type': 'sgd',
Abdullah Rashwan's avatar
Abdullah Rashwan committed
94
            'sgd': {'momentum': 0.9}
Abdullah Rashwan's avatar
Abdullah Rashwan committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
        },
        '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)
Mark Daoust's avatar
Mark Daoust committed
110
  ```
Abdullah Rashwan's avatar
Abdullah Rashwan committed
111
112
113
114
115
116
117
118
119
120
121
122
  """

  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
123
124
125
126
    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
127
128
129
130
131
      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
132
133
134
    if self._lr_type is None:
      raise ValueError('Learning rate type must be specified')

Abdullah Rashwan's avatar
Abdullah Rashwan committed
135
136
137
138
139
140
141
    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
142
143
    to the learning rate config. If learning rate type is consant,
    lr_config.learning_rate is returned.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
144
145

    Returns:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
146
147
      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
148
    """
Abdullah Rashwan's avatar
Abdullah Rashwan committed
149
150
    if self._lr_type == 'constant':
      lr = self._lr_config.learning_rate
Abdullah Rashwan's avatar
Abdullah Rashwan committed
151
152
153
154
155
156
157
158
    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
159
  @gin.configurable
Abdullah Rashwan's avatar
Abdullah Rashwan committed
160
  def build_optimizer(
Le Hou's avatar
Le Hou committed
161
162
      self,
      lr: Union[tf.keras.optimizers.schedules.LearningRateSchedule, float],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
163
164
165
      gradient_aggregator: Optional[Callable[
          [List[Tuple[tf.Tensor, tf.Tensor]]], List[Tuple[tf.Tensor,
                                                          tf.Tensor]]]] = None,
Frederick Liu's avatar
Frederick Liu committed
166
      gradient_transformers: Optional[List[Callable[
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
167
168
          [List[Tuple[tf.Tensor, tf.Tensor]]], List[Tuple[tf.Tensor,
                                                          tf.Tensor]]]]] = None,
Rebecca Chen's avatar
Rebecca Chen committed
169
170
      postprocessor: Optional[Callable[[tf.keras.optimizers.Optimizer],
                                       tf.keras.optimizers.Optimizer]] = None):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
171
172
173
174
175
176
177
    """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
178
179
      lr: A floating point value, or a
        tf.keras.optimizers.schedules.LearningRateSchedule instance.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
180
      gradient_aggregator: Optional function to overwrite gradient aggregation.
Frederick Liu's avatar
Frederick Liu committed
181
182
183
184
185
      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
186
187
      postprocessor: An optional function for postprocessing the optimizer. It
        takes an optimizer and returns an optimizer.
Hongkun Yu's avatar
Hongkun Yu committed
188

Abdullah Rashwan's avatar
Abdullah Rashwan committed
189
    Returns:
Chen Qian's avatar
Chen Qian committed
190
191
      `tf.keras.optimizers.Optimizer` or
      `tf.keras.optimizers.experimental.Optimizer` instance.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
192
193
194
    """

    optimizer_dict = self._optimizer_config.as_dict()
Frederick Liu's avatar
Frederick Liu committed
195
    ## Delete clipnorm, clipvalue, global_clipnorm if None
Abdullah Rashwan's avatar
Abdullah Rashwan committed
196
197
198
199
    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
200
201
    if optimizer_dict['global_clipnorm'] is None:
      del optimizer_dict['global_clipnorm']
Abdullah Rashwan's avatar
Abdullah Rashwan committed
202

Abdullah Rashwan's avatar
Abdullah Rashwan committed
203
    optimizer_dict['learning_rate'] = lr
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
204
205
    if gradient_aggregator is not None:
      optimizer_dict['gradient_aggregator'] = gradient_aggregator
Frederick Liu's avatar
Frederick Liu committed
206
207
    if gradient_transformers is not None:
      optimizer_dict['gradient_transformers'] = gradient_transformers
Abdullah Rashwan's avatar
Abdullah Rashwan committed
208
209

    optimizer = OPTIMIZERS_CLS[self._optimizer_type](**optimizer_dict)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
210
211
212
213

    if self._use_ema:
      optimizer = ema_optimizer.ExponentialMovingAverage(
          optimizer, **self._ema_config.as_dict())
Le Hou's avatar
Le Hou committed
214
215
    if postprocessor:
      optimizer = postprocessor(optimizer)
Chen Qian's avatar
Chen Qian committed
216
217
218
219
    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
220
        '{}'.format(optimizer))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
221

Abdullah Rashwan's avatar
Abdullah Rashwan committed
222
    return optimizer