Spaces:
No application file
No application file
| # coding=utf-8 | |
| # Copyright 2018 The Google AI Team Authors. | |
| # | |
| # 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. | |
| # Lint as: python2, python3 | |
| """Functions and classes related to optimization (weight updates).""" | |
| from __future__ import absolute_import | |
| from __future__ import division | |
| from __future__ import print_function | |
| import re | |
| import six | |
| import tensorflow.compat.v1 as tf | |
| # pylint: disable=g-direct-tensorflow-import | |
| from tensorflow.python.ops import array_ops | |
| from tensorflow.python.ops import linalg_ops | |
| from tensorflow.python.ops import math_ops | |
| # pylint: enable=g-direct-tensorflow-import | |
| class LAMBOptimizer(tf.train.Optimizer): | |
| """LAMB (Layer-wise Adaptive Moments optimizer for Batch training).""" | |
| # A new optimizer that includes correct L2 weight decay, adaptive | |
| # element-wise updating, and layer-wise justification. The LAMB optimizer | |
| # was proposed by Yang You, Jing Li, Jonathan Hseu, Xiaodan Song, | |
| # James Demmel, and Cho-Jui Hsieh in a paper titled as Reducing BERT | |
| # Pre-Training Time from 3 Days to 76 Minutes (arxiv.org/abs/1904.00962) | |
| def __init__(self, | |
| learning_rate, | |
| weight_decay_rate=0.0, | |
| beta_1=0.9, | |
| beta_2=0.999, | |
| epsilon=1e-6, | |
| exclude_from_weight_decay=None, | |
| exclude_from_layer_adaptation=None, | |
| name="LAMBOptimizer"): | |
| """Constructs a LAMBOptimizer.""" | |
| super(LAMBOptimizer, self).__init__(False, name) | |
| self.learning_rate = learning_rate | |
| self.weight_decay_rate = weight_decay_rate | |
| self.beta_1 = beta_1 | |
| self.beta_2 = beta_2 | |
| self.epsilon = epsilon | |
| self.exclude_from_weight_decay = exclude_from_weight_decay | |
| # exclude_from_layer_adaptation is set to exclude_from_weight_decay if the | |
| # arg is None. | |
| # TODO(jingli): validate if exclude_from_layer_adaptation is necessary. | |
| if exclude_from_layer_adaptation: | |
| self.exclude_from_layer_adaptation = exclude_from_layer_adaptation | |
| else: | |
| self.exclude_from_layer_adaptation = exclude_from_weight_decay | |
| def apply_gradients(self, grads_and_vars, global_step=None, name=None): | |
| """See base class.""" | |
| assignments = [] | |
| for (grad, param) in grads_and_vars: | |
| if grad is None or param is None: | |
| continue | |
| param_name = self._get_variable_name(param.name) | |
| m = tf.get_variable( | |
| name=six.ensure_str(param_name) + "/adam_m", | |
| shape=param.shape.as_list(), | |
| dtype=tf.float32, | |
| trainable=False, | |
| initializer=tf.zeros_initializer()) | |
| v = tf.get_variable( | |
| name=six.ensure_str(param_name) + "/adam_v", | |
| shape=param.shape.as_list(), | |
| dtype=tf.float32, | |
| trainable=False, | |
| initializer=tf.zeros_initializer()) | |
| # Standard Adam update. | |
| next_m = ( | |
| tf.multiply(self.beta_1, m) + tf.multiply(1.0 - self.beta_1, grad)) | |
| next_v = ( | |
| tf.multiply(self.beta_2, v) + tf.multiply(1.0 - self.beta_2, | |
| tf.square(grad))) | |
| update = next_m / (tf.sqrt(next_v) + self.epsilon) | |
| # 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. | |
| # | |
| # Instead we want ot decay the weights in a manner that doesn't interact | |
| # with the m/v parameters. This is equivalent to adding the square | |
| # of the weights to the loss with plain (non-momentum) SGD. | |
| if self._do_use_weight_decay(param_name): | |
| update += self.weight_decay_rate * param | |
| ratio = 1.0 | |
| if self._do_layer_adaptation(param_name): | |
| w_norm = linalg_ops.norm(param, ord=2) | |
| g_norm = linalg_ops.norm(update, ord=2) | |
| ratio = array_ops.where(math_ops.greater(w_norm, 0), array_ops.where( | |
| math_ops.greater(g_norm, 0), (w_norm / g_norm), 1.0), 1.0) | |
| update_with_lr = ratio * self.learning_rate * update | |
| next_param = param - update_with_lr | |
| assignments.extend( | |
| [param.assign(next_param), | |
| m.assign(next_m), | |
| v.assign(next_v)]) | |
| return tf.group(*assignments, name=name) | |
| def _do_use_weight_decay(self, param_name): | |
| """Whether to use L2 weight decay for `param_name`.""" | |
| if not self.weight_decay_rate: | |
| return False | |
| if self.exclude_from_weight_decay: | |
| for r in self.exclude_from_weight_decay: | |
| if re.search(r, param_name) is not None: | |
| return False | |
| return True | |
| def _do_layer_adaptation(self, param_name): | |
| """Whether to do layer-wise learning rate adaptation for `param_name`.""" | |
| if self.exclude_from_layer_adaptation: | |
| for r in self.exclude_from_layer_adaptation: | |
| if re.search(r, param_name) is not None: | |
| return False | |
| return True | |
| def _get_variable_name(self, param_name): | |
| """Get the variable name from the tensor name.""" | |
| m = re.match("^(.*):\\d+$", six.ensure_str(param_name)) | |
| if m is not None: | |
| param_name = m.group(1) | |
| return param_name | |