2019 the tensorflow s the hugging face team 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 orglicenseslicense2 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 functions and classes related to optimization weight updates import re from typing import callable list optional union import tensorflow as tf try from tensorflow keras optimizers legacy import adam except importerror from tensorflow keras optimizers import adam class warmuptf keras optimizers schedules learningrateschedule def init self initiallearningrate float decayschedulefn callable warmupsteps int power float 1 0 name str none super init self initiallearningrate initiallearningrate self warmupsteps warmupsteps self power power self decayschedulefn decayschedulefn self name name def callself step with tf namescopeself name or warmup as name implements polynomial warmup i e if globalstep warmupsteps the learning rate will be globalstepnumwarmupsteps initlr globalstepfloat tf caststep tf float32 warmupstepsfloat tf castself warmupsteps tf float32 warmuppercentdone globalstepfloat warmupstepsfloat warmuplearningrate self initiallearningrate tf math powwarmuppercentdone self power return tf cond globalstepfloat warmupstepsfloat lambda warmuplearningrate lambda self decayschedulefnstep self warmupsteps namename def getconfigself return initiallearningrate self initiallearningrate decayschedulefn self decayschedulefn warmupsteps self warmupsteps power self power name self name def createoptimizer initlr float numtrainsteps int numwarmupsteps int minlrratio float 0 0 adambeta1 float 0 9 adambeta2 float 0 999 adamepsilon float 1e8 adamclipnorm optionalfloat none adamglobalclipnorm optionalfloat none weightdecayrate float 0 0 power float 1 0 includeinweightdecay optionalliststr none implements linear decay of the learning rate lrschedule tf keras optimizers schedules polynomialdecay initiallearningrateinitlr decaystepsnumtrainsteps numwarmupsteps endlearningrateinitlr minlrratio powerpower if numwarmupsteps lrschedule warmup initiallearningrateinitlr decayschedulefnlrschedule warmupstepsnumwarmupsteps if weightdecayrate 0 0 optimizer adamweightdecay learningratelrschedule weightdecayrateweightdecayrate beta1adambeta1 beta2adambeta2 epsilonadamepsilon clipnormadamclipnorm globalclipnormadamglobalclipnorm excludefromweightdecaylayernorm layernorm bias includeinweightdecayincludeinweightdecay else optimizer tf keras optimizers adam learningratelrschedule beta1adambeta1 beta2adambeta2 epsilonadamepsilon clipnormadamclipnorm globalclipnormadamglobalclipnorm we return the optimizer and the lr scheduler in order to better track the evolution of the lr independently of the optimizer return optimizer lrschedule class adamweightdecayadam def init self learningrate unionfloat tf keras optimizers schedules learningrateschedule 0 001 beta1 float 0 9 beta2 float 0 999 epsilon float 1e7 amsgrad bool false weightdecayrate float 0 0 includeinweightdecay optionalliststr none excludefromweightdecay optionalliststr none name str adamweightdecay kwargs super initlearningrate beta1 beta2 epsilon amsgrad name kwargs self weightdecayrate weightdecayrate self includeinweightdecay includeinweightdecay self excludefromweightdecay excludefromweightdecay classmethod def fromconfigcls config retrieves the learning rate with the given state if applystate is none return self decayedlrtvardtype applystate applystate or coefficients applystate getvardevice vardtype if coefficients is none coefficients self fallbackapplystatevardevice vardtype applystatevardevice vardtype coefficients return coefficientslrt applystate applystate def resourceapplydenseself grad var applystatenone lrt kwargs self getlrvar device var dtype basedtype applystate decay self decayweightsopvar lrt applystate with tf controldependenciesdecay return superadamweightdecay self resourceapplydensegrad var kwargs def resourceapplysparseself grad var indices applystatenone lrt kwargs self getlrvar device var dtype basedtype applystate decay self decayweightsopvar lrt applystate with tf controldependenciesdecay return superadamweightdecay self resourceapplysparsegrad var indices kwargs def getconfigself config super getconfig config updateweightdecayrate self weightdecayrate return config def douseweightdecayself paramname extracted from https github comopennmtopennmttfblobmasteropennmtoptimizersutils py gradient accumulation utility when used with a distribution strategy the accumulator should be called in a replica context gradients will be accumulated locally on each replica and without synchronization users should then call gradients scale the gradients if required and pass the result to applygradients we use the onread synchronization policy so that no synchronization is performed on assignment to get the value we call value which returns the value on the current replica without synchronization initializes the accumulator self gradients self accumsteps none property def stepself the accumulated gradients on the current replica if not self gradients raise valueerrorthe accumulator should be called first to initialize the gradients return gradient value if gradient is not none else gradient for gradient in self gradients def callself gradients resets the accumulated gradients on the current replica if not self gradients return self accumsteps assign0 for gradient in self gradients if gradient is not none gradient assigntf zeroslikegradient 2019 the tensorflow s the hugging face team 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 functions and classes related to optimization weight updates applies a warmup schedule on a given learning rate decay schedule args initial_learning_rate float the initial learning rate for the schedule after the warmup so this will be the learning rate at the end of the warmup decay_schedule_fn callable the schedule function to apply after the warmup for the rest of training warmup_steps int the number of steps for the warmup part of training power float optional defaults to 1 0 the power to use for the polynomial warmup defaults is a linear warmup name str optional optional name prefix for the returned tensors during the schedule implements polynomial warmup i e if global_step warmup_steps the learning rate will be global_step num_warmup_steps init_lr creates an optimizer with a learning rate schedule using a warmup phase followed by a linear decay args init_lr float the desired learning rate at the end of the warmup phase num_train_steps int the total number of training steps num_warmup_steps int the number of warmup steps min_lr_ratio float optional defaults to 0 the final learning rate at the end of the linear decay will be init_lr min_lr_ratio adam_beta1 float optional defaults to 0 9 the beta1 to use in adam adam_beta2 float optional defaults to 0 999 the beta2 to use in adam adam_epsilon float optional defaults to 1e 8 the epsilon to use in adam adam_clipnorm float optional defaults to none if not none clip the gradient norm for each weight tensor to this value adam_global_clipnorm float optional defaults to none if not none clip gradient norm to this value when using this argument the norm is computed over all weight tensors as if they were concatenated into a single vector weight_decay_rate float optional defaults to 0 the weight decay to use power float optional defaults to 1 0 the power to use for polynomialdecay include_in_weight_decay list str optional list of the parameter names or re patterns to apply weight decay to if none is passed weight decay is applied to all parameters except bias and layer norm parameters implements linear decay of the learning rate we return the optimizer and the lr scheduler in order to better track the evolution of the lr independently of the optimizer adam enables l2 weight decay and clip_by_global_norm on gradients 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 as shown in decoupled weight decay regularization https arxiv org abs 1711 05101 instead we want to 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 args learning_rate union float tf keras optimizers schedules learningrateschedule optional defaults to 0 001 the learning rate to use or a schedule beta_1 float optional defaults to 0 9 the beta1 parameter in adam which is the exponential decay rate for the 1st momentum estimates beta_2 float optional defaults to 0 999 the beta2 parameter in adam which is the exponential decay rate for the 2nd momentum estimates epsilon float optional defaults to 1e 07 the epsilon parameter in adam which is a small constant for numerical stability amsgrad bool optional defaults to false whether to apply amsgrad variant of this algorithm or not see on the convergence of adam and beyond https arxiv org abs 1904 09237 weight_decay_rate float optional defaults to 0 0 the weight decay to apply include_in_weight_decay list str optional list of the parameter names or re patterns to apply weight decay to if none is passed weight decay is applied to all parameters by default unless they are in exclude_from_weight_decay exclude_from_weight_decay list str optional list of the parameter names or re patterns to exclude from applying weight decay to if a include_in_weight_decay is passed the names in it will supersede this list name str optional defaults to adamweightdecay optional name for the operations created when applying gradients kwargs dict str any optional keyword arguments allowed to be clipnorm clipvalue lr decay clipnorm is clip gradients by norm clipvalue is clip gradients by value decay is included for backward compatibility to allow time inverse decay of learning rate lr is included for backward compatibility recommended to use learning_rate instead creates an optimizer from its config with warmup custom object retrieves the learning rate with the given state whether to use l2 weight decay for param_name extracted from https github com opennmt opennmt tf blob master opennmt optimizers utils py gradient accumulation utility when used with a distribution strategy the accumulator should be called in a replica context gradients will be accumulated locally on each replica and without synchronization users should then call gradients scale the gradients if required and pass the result to apply_gradients we use the on_read synchronization policy so that no synchronization is performed on assignment to get the value we call value which returns the value on the current replica without synchronization initializes the accumulator number of accumulated steps the accumulated gradients on the current replica accumulates gradients on the current replica create the step variable resets the accumulated gradients on the current replica
import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
    from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
    from tensorflow.keras.optimizers import Adam


class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):
    

    def __init__(
        self,
        initial_learning_rate: float,
        decay_schedule_fn: Callable,
        warmup_steps: int,
        power: float = 1.0,
        name: str = None,
    ):
        super().__init__()
        self.initial_learning_rate = initial_learning_rate
        self.warmup_steps = warmup_steps
        self.power = power
        self.decay_schedule_fn = decay_schedule_fn
        self.name = name

    def __call__(self, step):
        with tf.name_scope(self.name or "WarmUp") as name:
            
            
            global_step_float = tf.cast(step, tf.float32)
            warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)
            warmup_percent_done = global_step_float / warmup_steps_float
            warmup_learning_rate = self.initial_learning_rate * tf.math.pow(warmup_percent_done, self.power)
            return tf.cond(
                global_step_float < warmup_steps_float,
                lambda: warmup_learning_rate,
                lambda: self.decay_schedule_fn(step - self.warmup_steps),
                name=name,
            )

    def get_config(self):
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "decay_schedule_fn": self.decay_schedule_fn,
            "warmup_steps": self.warmup_steps,
            "power": self.power,
            "name": self.name,
        }


def create_optimizer(
    init_lr: float,
    num_train_steps: int,
    num_warmup_steps: int,
    min_lr_ratio: float = 0.0,
    adam_beta1: float = 0.9,
    adam_beta2: float = 0.999,
    adam_epsilon: float = 1e-8,
    adam_clipnorm: Optional[float] = None,
    adam_global_clipnorm: Optional[float] = None,
    weight_decay_rate: float = 0.0,
    power: float = 1.0,
    include_in_weight_decay: Optional[List[str]] = None,
):
    
    
    lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
        initial_learning_rate=init_lr,
        decay_steps=num_train_steps - num_warmup_steps,
        end_learning_rate=init_lr * min_lr_ratio,
        power=power,
    )
    if num_warmup_steps:
        lr_schedule = WarmUp(
            initial_learning_rate=init_lr,
            decay_schedule_fn=lr_schedule,
            warmup_steps=num_warmup_steps,
        )
    if weight_decay_rate > 0.0:
        optimizer = AdamWeightDecay(
            learning_rate=lr_schedule,
            weight_decay_rate=weight_decay_rate,
            beta_1=adam_beta1,
            beta_2=adam_beta2,
            epsilon=adam_epsilon,
            clipnorm=adam_clipnorm,
            global_clipnorm=adam_global_clipnorm,
            exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"],
            include_in_weight_decay=include_in_weight_decay,
        )
    else:
        optimizer = tf.keras.optimizers.Adam(
            learning_rate=lr_schedule,
            beta_1=adam_beta1,
            beta_2=adam_beta2,
            epsilon=adam_epsilon,
            clipnorm=adam_clipnorm,
            global_clipnorm=adam_global_clipnorm,
        )
    
    
    return optimizer, lr_schedule


class AdamWeightDecay(Adam):
    

    def __init__(
        self,
        learning_rate: Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.001,
        beta_1: float = 0.9,
        beta_2: float = 0.999,
        epsilon: float = 1e-7,
        amsgrad: bool = False,
        weight_decay_rate: float = 0.0,
        include_in_weight_decay: Optional[List[str]] = None,
        exclude_from_weight_decay: Optional[List[str]] = None,
        name: str = "AdamWeightDecay",
        **kwargs,
    ):
        super().__init__(learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs)
        self.weight_decay_rate = weight_decay_rate
        self._include_in_weight_decay = include_in_weight_decay
        self._exclude_from_weight_decay = exclude_from_weight_decay

    @classmethod
    def from_config(cls, config):
        
        custom_objects = {"WarmUp": WarmUp}
        return super(AdamWeightDecay, cls).from_config(config, custom_objects=custom_objects)

    def _prepare_local(self, var_device, var_dtype, apply_state):
        super(AdamWeightDecay, self)._prepare_local(var_device, var_dtype, apply_state)
        apply_state[(var_device, var_dtype)]["weight_decay_rate"] = tf.constant(
            self.weight_decay_rate, name="adam_weight_decay_rate"
        )

    def _decay_weights_op(self, var, learning_rate, apply_state):
        do_decay = self._do_use_weight_decay(var.name)
        if do_decay:
            return var.assign_sub(
                learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]["weight_decay_rate"],
                use_locking=self._use_locking,
            )
        return tf.no_op()

    def apply_gradients(self, grads_and_vars, name=None, **kwargs):
        grads, tvars = list(zip(*grads_and_vars))
        return super(AdamWeightDecay, self).apply_gradients(zip(grads, tvars), name=name, **kwargs)

    def _get_lr(self, var_device, var_dtype, apply_state):
        
        if apply_state is None:
            return self._decayed_lr_t[var_dtype], {}

        apply_state = apply_state or {}
        coefficients = apply_state.get((var_device, var_dtype))
        if coefficients is None:
            coefficients = self._fallback_apply_state(var_device, var_dtype)
            apply_state[(var_device, var_dtype)] = coefficients

        return coefficients["lr_t"], {"apply_state": apply_state}

    def _resource_apply_dense(self, grad, var, apply_state=None):
        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)
        decay = self._decay_weights_op(var, lr_t, apply_state)
        with tf.control_dependencies([decay]):
            return super(AdamWeightDecay, self)._resource_apply_dense(grad, var, **kwargs)

    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)
        decay = self._decay_weights_op(var, lr_t, apply_state)
        with tf.control_dependencies([decay]):
            return super(AdamWeightDecay, self)._resource_apply_sparse(grad, var, indices, **kwargs)

    def get_config(self):
        config = super().get_config()
        config.update({"weight_decay_rate": self.weight_decay_rate})
        return config

    def _do_use_weight_decay(self, param_name):
        
        if self.weight_decay_rate == 0:
            return False

        if self._include_in_weight_decay:
            for r in self._include_in_weight_decay:
                if re.search(r, param_name) is not None:
                    return True

        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



class GradientAccumulator(object):
    

    
    
    

    def __init__(self):
        
        self._gradients = []
        self._accum_steps = None

    @property
    def step(self):
        
        if self._accum_steps is None:
            self._accum_steps = tf.Variable(
                tf.constant(0, dtype=tf.int64),
                trainable=False,
                synchronization=tf.VariableSynchronization.ON_READ,
                aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
            )

        return self._accum_steps.value()

    @property
    def gradients(self):
        
        if not self._gradients:
            raise ValueError("The accumulator should be called first to initialize the gradients")
        return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]

    def __call__(self, gradients):
        
        if not self._gradients:
            _ = self.step  
            self._gradients.extend(
                [
                    tf.Variable(
                        tf.zeros_like(gradient),
                        trainable=False,
                        synchronization=tf.VariableSynchronization.ON_READ,
                        aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
                    )
                    if gradient is not None
                    else gradient
                    for gradient in gradients
                ]
            )
        if len(gradients) != len(self._gradients):
            raise ValueError(f"Expected {len(self._gradients)} gradients, but got {len(gradients)}")

        for accum_gradient, gradient in zip(self._gradients, gradients):
            if accum_gradient is not None and gradient is not None:
                accum_gradient.assign_add(gradient)

        self._accum_steps.assign_add(1)

    def reset(self):
        
        if not self._gradients:
            return
        self._accum_steps.assign(0)
        for gradient in self._gradients:
            if gradient is not None:
                gradient.assign(tf.zeros_like(gradient))