Aluode/PerceptionLabPortable

1# Copyright 2019 The TensorFlow Authors, The Hugging Face Team. All Rights Reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14# ==============================================================================15"""Functions and classes related to optimization (weight updates)."""16 17from typing import Callable, Optional, Union18 19import tensorflow as tf20 21 22try:23    from tf_keras.optimizers.legacy import Adam24except (ImportError, ModuleNotFoundError):25    from tensorflow.keras.optimizers.legacy import Adam26 27from .modeling_tf_utils import keras28 29 30# This block because Keras loves randomly moving things to different places - this changed somewhere between 2.10 - 2.1531if hasattr(keras.optimizers.schedules, "learning_rate_schedule"):32    schedules = keras.optimizers.schedules.learning_rate_schedule33else:34    schedules = keras.optimizers.schedules35 36 37class WarmUp(schedules.LearningRateSchedule):38    """39    Applies a warmup schedule on a given learning rate decay schedule.40 41    Args:42        initial_learning_rate (`float`):43            The initial learning rate for the schedule after the warmup (so this will be the learning rate at the end44            of the warmup).45        decay_schedule_fn (`Callable`):46            The schedule function to apply after the warmup for the rest of training.47        warmup_steps (`int`):48            The number of steps for the warmup part of training.49        power (`float`, *optional*, defaults to 1.0):50            The power to use for the polynomial warmup (defaults is a linear warmup).51        name (`str`, *optional*):52            Optional name prefix for the returned tensors during the schedule.53    """54 55    def __init__(56        self,57        initial_learning_rate: float,58        decay_schedule_fn: Callable,59        warmup_steps: int,60        power: float = 1.0,61        name: Optional[str] = None,62    ):63        super().__init__()64        self.initial_learning_rate = initial_learning_rate65        self.warmup_steps = warmup_steps66        self.power = power67        self.decay_schedule_fn = decay_schedule_fn68        self.name = name69 70    def __call__(self, step):71        with tf.name_scope(self.name or "WarmUp") as name:72            # Implements polynomial warmup. i.e., if global_step < warmup_steps, the73            # learning rate will be `global_step/num_warmup_steps * init_lr`.74            global_step_float = tf.cast(step, tf.float32)75            warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)76            warmup_percent_done = global_step_float / warmup_steps_float77            warmup_learning_rate = self.initial_learning_rate * tf.math.pow(warmup_percent_done, self.power)78            return tf.cond(79                global_step_float < warmup_steps_float,80                lambda: warmup_learning_rate,81                lambda: self.decay_schedule_fn(step - self.warmup_steps),82                name=name,83            )84 85    def get_config(self):86        return {87            "initial_learning_rate": self.initial_learning_rate,88            "decay_schedule_fn": self.decay_schedule_fn,89            "warmup_steps": self.warmup_steps,90            "power": self.power,91            "name": self.name,92        }93 94 95def create_optimizer(96    init_lr: float,97    num_train_steps: int,98    num_warmup_steps: int,99    min_lr_ratio: float = 0.0,100    adam_beta1: float = 0.9,101    adam_beta2: float = 0.999,102    adam_epsilon: float = 1e-8,103    adam_clipnorm: Optional[float] = None,104    adam_global_clipnorm: Optional[float] = None,105    weight_decay_rate: float = 0.0,106    power: float = 1.0,107    include_in_weight_decay: Optional[list[str]] = None,108):109    """110    Creates an optimizer with a learning rate schedule using a warmup phase followed by a linear decay.111 112    Args:113        init_lr (`float`):114            The desired learning rate at the end of the warmup phase.115        num_train_steps (`int`):116            The total number of training steps.117        num_warmup_steps (`int`):118            The number of warmup steps.119        min_lr_ratio (`float`, *optional*, defaults to 0):120            The final learning rate at the end of the linear decay will be `init_lr * min_lr_ratio`.121        adam_beta1 (`float`, *optional*, defaults to 0.9):122            The beta1 to use in Adam.123        adam_beta2 (`float`, *optional*, defaults to 0.999):124            The beta2 to use in Adam.125        adam_epsilon (`float`, *optional*, defaults to 1e-8):126            The epsilon to use in Adam.127        adam_clipnorm (`float`, *optional*, defaults to `None`):128            If not `None`, clip the gradient norm for each weight tensor to this value.129        adam_global_clipnorm (`float`, *optional*, defaults to `None`)130            If not `None`, clip gradient norm to this value. When using this argument, the norm is computed over all131            weight tensors, as if they were concatenated into a single vector.132        weight_decay_rate (`float`, *optional*, defaults to 0):133            The weight decay to use.134        power (`float`, *optional*, defaults to 1.0):135            The power to use for PolynomialDecay.136        include_in_weight_decay (`list[str]`, *optional*):137            List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is138            applied to all parameters except bias and layer norm parameters.139    """140    # Implements linear decay of the learning rate.141    lr_schedule = schedules.PolynomialDecay(142        initial_learning_rate=init_lr,143        decay_steps=num_train_steps - num_warmup_steps,144        end_learning_rate=init_lr * min_lr_ratio,145        power=power,146    )147    if num_warmup_steps:148        lr_schedule = WarmUp(149            initial_learning_rate=init_lr,150            decay_schedule_fn=lr_schedule,151            warmup_steps=num_warmup_steps,152        )153    if weight_decay_rate > 0.0:154        optimizer = AdamWeightDecay(155            learning_rate=lr_schedule,156            weight_decay_rate=weight_decay_rate,157            beta_1=adam_beta1,158            beta_2=adam_beta2,159            epsilon=adam_epsilon,160            clipnorm=adam_clipnorm,161            global_clipnorm=adam_global_clipnorm,162            exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"],163            include_in_weight_decay=include_in_weight_decay,164        )165    else:166        optimizer = keras.optimizers.Adam(167            learning_rate=lr_schedule,168            beta_1=adam_beta1,169            beta_2=adam_beta2,170            epsilon=adam_epsilon,171            clipnorm=adam_clipnorm,172            global_clipnorm=adam_global_clipnorm,173        )174    # We return the optimizer and the LR scheduler in order to better track the175    # evolution of the LR independently of the optimizer.176    return optimizer, lr_schedule177 178 179class AdamWeightDecay(Adam):180    """181    Adam enables L2 weight decay and clip_by_global_norm on gradients. Just adding the square of the weights to the182    loss function is *not* the correct way of using L2 regularization/weight decay with Adam, since that will interact183    with the m and v parameters in strange ways as shown in [Decoupled Weight Decay184    Regularization](https://huggingface.co/papers/1711.05101).185 186    Instead we want to decay the weights in a manner that doesn't interact with the m/v parameters. This is equivalent187    to adding the square of the weights to the loss with plain (non-momentum) SGD.188 189    Args:190        learning_rate (`Union[float, LearningRateSchedule]`, *optional*, defaults to 0.001):191            The learning rate to use or a schedule.192        beta_1 (`float`, *optional*, defaults to 0.9):193            The beta1 parameter in Adam, which is the exponential decay rate for the 1st momentum estimates.194        beta_2 (`float`, *optional*, defaults to 0.999):195            The beta2 parameter in Adam, which is the exponential decay rate for the 2nd momentum estimates.196        epsilon (`float`, *optional*, defaults to 1e-07):197            The epsilon parameter in Adam, which is a small constant for numerical stability.198        amsgrad (`bool`, *optional*, defaults to `False`):199            Whether to apply AMSGrad variant of this algorithm or not, see [On the Convergence of Adam and200            Beyond](https://huggingface.co/papers/1904.09237).201        weight_decay_rate (`float`, *optional*, defaults to 0.0):202            The weight decay to apply.203        include_in_weight_decay (`list[str]`, *optional*):204            List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is205            applied to all parameters by default (unless they are in `exclude_from_weight_decay`).206        exclude_from_weight_decay (`list[str]`, *optional*):207            List of the parameter names (or re patterns) to exclude from applying weight decay to. If a208            `include_in_weight_decay` is passed, the names in it will supersede this list.209        name (`str`, *optional*, defaults to `"AdamWeightDecay"`):210            Optional name for the operations created when applying gradients.211        kwargs (`dict[str, Any]`, *optional*):212            Keyword arguments. Allowed to be {`clipnorm`, `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by213            norm; `clipvalue` is clip gradients by value, `decay` is included for backward compatibility to allow time214            inverse decay of learning rate. `lr` is included for backward compatibility, recommended to use215            `learning_rate` instead.216    """217 218    def __init__(219        self,220        learning_rate: Union[float, schedules.LearningRateSchedule] = 0.001,221        beta_1: float = 0.9,222        beta_2: float = 0.999,223        epsilon: float = 1e-7,224        amsgrad: bool = False,225        weight_decay_rate: float = 0.0,226        include_in_weight_decay: Optional[list[str]] = None,227        exclude_from_weight_decay: Optional[list[str]] = None,228        name: str = "AdamWeightDecay",229        **kwargs,230    ):231        super().__init__(learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs)232        self.weight_decay_rate = weight_decay_rate233        self._include_in_weight_decay = include_in_weight_decay234        self._exclude_from_weight_decay = exclude_from_weight_decay235 236    @classmethod237    def from_config(cls, config):238        """Creates an optimizer from its config with WarmUp custom object."""239        custom_objects = {"WarmUp": WarmUp}240        return super().from_config(config, custom_objects=custom_objects)241 242    def _prepare_local(self, var_device, var_dtype, apply_state):243        super()._prepare_local(var_device, var_dtype, apply_state)244        apply_state[(var_device, var_dtype)]["weight_decay_rate"] = tf.constant(245            self.weight_decay_rate, name="adam_weight_decay_rate"246        )247 248    def _decay_weights_op(self, var, learning_rate, apply_state):249        do_decay = self._do_use_weight_decay(var.name)250        if do_decay:251            return var.assign_sub(252                learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]["weight_decay_rate"],253                use_locking=self._use_locking,254            )255        return tf.no_op()256 257    def apply_gradients(self, grads_and_vars, name=None, **kwargs):258        grads, tvars = list(zip(*grads_and_vars))259        return super().apply_gradients(zip(grads, tvars), name=name, **kwargs)260 261    def _get_lr(self, var_device, var_dtype, apply_state):262        """Retrieves the learning rate with the given state."""263        if apply_state is None:264            return self._decayed_lr_t[var_dtype], {}265 266        apply_state = apply_state or {}267        coefficients = apply_state.get((var_device, var_dtype))268        if coefficients is None:269            coefficients = self._fallback_apply_state(var_device, var_dtype)270            apply_state[(var_device, var_dtype)] = coefficients271 272        return coefficients["lr_t"], {"apply_state": apply_state}273 274    def _resource_apply_dense(self, grad, var, apply_state=None):275        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)276        decay = self._decay_weights_op(var, lr_t, apply_state)277        with tf.control_dependencies([decay]):278            return super()._resource_apply_dense(grad, var, **kwargs)279 280    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):281        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)282        decay = self._decay_weights_op(var, lr_t, apply_state)283        with tf.control_dependencies([decay]):284            return super()._resource_apply_sparse(grad, var, indices, **kwargs)285 286    def get_config(self):287        config = super().get_config()288        config.update({"weight_decay_rate": self.weight_decay_rate})289        return config290 291    def _do_use_weight_decay(self, param_name):292        """Whether to use L2 weight decay for `param_name`."""293        if self.weight_decay_rate == 0:294            return False295 296        if self._include_in_weight_decay:297            for r in self._include_in_weight_decay:298                if r in param_name:299                    return True300 301        if self._exclude_from_weight_decay:302            for r in self._exclude_from_weight_decay:303                if r in param_name:304                    return False305        return True306 307 308# Extracted from https://github.com/OpenNMT/OpenNMT-tf/blob/master/opennmt/optimizers/utils.py309class GradientAccumulator:310    """311    Gradient accumulation utility. When used with a distribution strategy, the accumulator should be called in a312    replica context. Gradients will be accumulated locally on each replica and without synchronization. Users should313    then call `.gradients`, scale the gradients if required, and pass the result to `apply_gradients`.314    """315 316    # We use the ON_READ synchronization policy so that no synchronization is317    # performed on assignment. To get the value, we call .value() which returns the318    # value on the current replica without synchronization.319 320    def __init__(self):321        """Initializes the accumulator."""322        self._gradients = []323        self._accum_steps = None324 325    @property326    def step(self):327        """Number of accumulated steps."""328        if self._accum_steps is None:329            self._accum_steps = tf.Variable(330                tf.constant(0, dtype=tf.int64),331                trainable=False,332                synchronization=tf.VariableSynchronization.ON_READ,333                aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,334            )335 336        return self._accum_steps.value()337 338    @property339    def gradients(self):340        """The accumulated gradients on the current replica."""341        if not self._gradients:342            raise ValueError("The accumulator should be called first to initialize the gradients")343        return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]344 345    def __call__(self, gradients):346        """Accumulates `gradients` on the current replica."""347        if not self._gradients:348            _ = self.step  # Create the step variable.349            self._gradients.extend(350                [351                    tf.Variable(352                        tf.zeros_like(gradient),353                        trainable=False,354                        synchronization=tf.VariableSynchronization.ON_READ,355                        aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,356                    )357                    if gradient is not None358                    else gradient359                    for gradient in gradients360                ]361            )362        if len(gradients) != len(self._gradients):363            raise ValueError(f"Expected {len(self._gradients)} gradients, but got {len(gradients)}")364 365        for accum_gradient, gradient in zip(self._gradients, gradients):366            if accum_gradient is not None and gradient is not None:367                accum_gradient.assign_add(gradient)368 369        self._accum_steps.assign_add(1)370 371    def reset(self):372        """Resets the accumulated gradients on the current replica."""373        if not self._gradients:374            return375        self._accum_steps.assign(0)376        for gradient in self._gradients:377            if gradient is not None:378                gradient.assign(tf.zeros_like(gradient))379