We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
LayerNorm not excluded from weight decay in TF
# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...):
bert-base-cased

Language I am using the model on (English, Chinese ...):
English

The problem arises when using:
* [X] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [X] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Add a print statement to `_do_use_weight_decay` in [AdamWeightDecay](https://github.com/huggingface/transformers/blob/master/src/transformers/optimization_tf.py) to see which parameters are actually excluded:
```python
    def _do_use_weight_decay(self, param_name):
        """Whether to use L2 weight decay for `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:
                    print(f"Found: {param_name}")
                    return False
        return True
```
2. run `python examples/text-classification/run_tf_glue.py --model_name_or_path bert-base-cased --task_name mrpc --output_dir temp --logging_dir temp --do_train --overwrite_output_dir --optimizer_name adamw`. 
3. Observe that no weights related to layer norms are printed. 

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
The weights of the layer norms (and the biases) should be printed. 

See for example: https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py. 

Based on the fact that no layer norm weights are printed with "layer_norm" simply switching "layer_norm" to "LayerNorm" seems like the easiest change. 

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 2.9.0
- Platform: Darwin-19.4.0-x86_64-i386-64bit
- Python version: 3.7.7
- PyTorch version (GPU?): 1.5.0 (False)
- Tensorflow version (GPU?): 2.2.0 (False)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/transformers/optimization_tf.py`
Content:
```
1 # Copyright 2019 The TensorFlow Authors. 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 at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # 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 and
13 # limitations under the License.
14 # ==============================================================================
15 """Functions and classes related to optimization (weight updates)."""
16 
17 
18 import re
19 
20 import tensorflow as tf
21 
22 
23 class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):
24     """Applies a warmup schedule on a given learning rate decay schedule."""
25 
26     def __init__(
27         self, initial_learning_rate, decay_schedule_fn, warmup_steps, power=1.0, name=None,
28     ):
29         super().__init__()
30         self.initial_learning_rate = initial_learning_rate
31         self.warmup_steps = warmup_steps
32         self.power = power
33         self.decay_schedule_fn = decay_schedule_fn
34         self.name = name
35 
36     def __call__(self, step):
37         with tf.name_scope(self.name or "WarmUp") as name:
38             # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
39             # learning rate will be `global_step/num_warmup_steps * init_lr`.
40             global_step_float = tf.cast(step, tf.float32)
41             warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)
42             warmup_percent_done = global_step_float / warmup_steps_float
43             warmup_learning_rate = self.initial_learning_rate * tf.math.pow(warmup_percent_done, self.power)
44             return tf.cond(
45                 global_step_float < warmup_steps_float,
46                 lambda: warmup_learning_rate,
47                 lambda: self.decay_schedule_fn(step),
48                 name=name,
49             )
50 
51     def get_config(self):
52         return {
53             "initial_learning_rate": self.initial_learning_rate,
54             "decay_schedule_fn": self.decay_schedule_fn,
55             "warmup_steps": self.warmup_steps,
56             "power": self.power,
57             "name": self.name,
58         }
59 
60 
61 def create_optimizer(init_lr, num_train_steps, num_warmup_steps, end_lr=0.0, optimizer_type="adamw"):
62     """Creates an optimizer with learning rate schedule."""
63     # Implements linear decay of the learning rate.
64     lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
65         initial_learning_rate=init_lr, decay_steps=num_train_steps, end_learning_rate=end_lr,
66     )
67     if num_warmup_steps:
68         lr_schedule = WarmUp(
69             initial_learning_rate=init_lr, decay_schedule_fn=lr_schedule, warmup_steps=num_warmup_steps,
70         )
71 
72     optimizer = AdamWeightDecay(
73         learning_rate=lr_schedule,
74         weight_decay_rate=0.01,
75         beta_1=0.9,
76         beta_2=0.999,
77         epsilon=1e-6,
78         exclude_from_weight_decay=["layer_norm", "bias"],
79     )
80 
81     return optimizer
82 
83 
84 class AdamWeightDecay(tf.keras.optimizers.Adam):
85     """Adam enables L2 weight decay and clip_by_global_norm on gradients.
86   Just adding the square of the weights to the loss function is *not* the
87   correct way of using L2 regularization/weight decay with Adam, since that will
88   interact with the m and v parameters in strange ways.
89   Instead we want ot decay the weights in a manner that doesn't interact with
90   the m/v parameters. This is equivalent to adding the square of the weights to
91   the loss with plain (non-momentum) SGD.
92   """
93 
94     def __init__(
95         self,
96         learning_rate=0.001,
97         beta_1=0.9,
98         beta_2=0.999,
99         epsilon=1e-7,
100         amsgrad=False,
101         weight_decay_rate=0.0,
102         include_in_weight_decay=None,
103         exclude_from_weight_decay=None,
104         name="AdamWeightDecay",
105         **kwargs
106     ):
107         super().__init__(learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs)
108         self.weight_decay_rate = weight_decay_rate
109         self._include_in_weight_decay = include_in_weight_decay
110         self._exclude_from_weight_decay = exclude_from_weight_decay
111 
112     @classmethod
113     def from_config(cls, config):
114         """Creates an optimizer from its config with WarmUp custom object."""
115         custom_objects = {"WarmUp": WarmUp}
116         return super(AdamWeightDecay, cls).from_config(config, custom_objects=custom_objects)
117 
118     def _prepare_local(self, var_device, var_dtype, apply_state):
119         super(AdamWeightDecay, self)._prepare_local(var_device, var_dtype, apply_state)
120         apply_state[(var_device, var_dtype)]["weight_decay_rate"] = tf.constant(
121             self.weight_decay_rate, name="adam_weight_decay_rate"
122         )
123 
124     def _decay_weights_op(self, var, learning_rate, apply_state):
125         do_decay = self._do_use_weight_decay(var.name)
126         if do_decay:
127             return var.assign_sub(
128                 learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]["weight_decay_rate"],
129                 use_locking=self._use_locking,
130             )
131         return tf.no_op()
132 
133     def apply_gradients(self, grads_and_vars, name=None):
134         grads, tvars = list(zip(*grads_and_vars))
135         return super(AdamWeightDecay, self).apply_gradients(zip(grads, tvars), name=name,)
136 
137     def _get_lr(self, var_device, var_dtype, apply_state):
138         """Retrieves the learning rate with the given state."""
139         if apply_state is None:
140             return self._decayed_lr_t[var_dtype], {}
141 
142         apply_state = apply_state or {}
143         coefficients = apply_state.get((var_device, var_dtype))
144         if coefficients is None:
145             coefficients = self._fallback_apply_state(var_device, var_dtype)
146             apply_state[(var_device, var_dtype)] = coefficients
147 
148         return coefficients["lr_t"], dict(apply_state=apply_state)
149 
150     def _resource_apply_dense(self, grad, var, apply_state=None):
151         lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)
152         decay = self._decay_weights_op(var, lr_t, apply_state)
153         with tf.control_dependencies([decay]):
154             return super(AdamWeightDecay, self)._resource_apply_dense(grad, var, **kwargs)
155 
156     def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
157         lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)
158         decay = self._decay_weights_op(var, lr_t, apply_state)
159         with tf.control_dependencies([decay]):
160             return super(AdamWeightDecay, self)._resource_apply_sparse(grad, var, indices, **kwargs)
161 
162     def get_config(self):
163         config = super().get_config()
164         config.update({"weight_decay_rate": self.weight_decay_rate})
165         return config
166 
167     def _do_use_weight_decay(self, param_name):
168         """Whether to use L2 weight decay for `param_name`."""
169         if self.weight_decay_rate == 0:
170             return False
171 
172         if self._include_in_weight_decay:
173             for r in self._include_in_weight_decay:
174                 if re.search(r, param_name) is not None:
175                     return True
176 
177         if self._exclude_from_weight_decay:
178             for r in self._exclude_from_weight_decay:
179                 if re.search(r, param_name) is not None:
180                     return False
181         return True
182 
183 
184 # Extracted from https://github.com/OpenNMT/OpenNMT-tf/blob/master/opennmt/optimizers/utils.py
185 class GradientAccumulator(object):
186     """Gradient accumulation utility.
187   When used with a distribution strategy, the accumulator should be called in a
188   replica context. Gradients will be accumulated locally on each replica and
189   without synchronization. Users should then call ``.gradients``, scale the
190   gradients if required, and pass the result to ``apply_gradients``.
191   """
192 
193     # We use the ON_READ synchronization policy so that no synchronization is
194     # performed on assignment. To get the value, we call .value() which returns the
195     # value on the current replica without synchronization.
196 
197     def __init__(self):
198         """Initializes the accumulator."""
199         self._gradients = []
200         self._accum_steps = None
201 
202     @property
203     def step(self):
204         """Number of accumulated steps."""
205         if self._accum_steps is None:
206             self._accum_steps = tf.Variable(
207                 tf.constant(0, dtype=tf.int64),
208                 trainable=False,
209                 synchronization=tf.VariableSynchronization.ON_READ,
210                 aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
211             )
212 
213         return self._accum_steps.value()
214 
215     @property
216     def gradients(self):
217         """The accumulated gradients on the current replica."""
218         if not self._gradients:
219             raise ValueError("The accumulator should be called first to initialize the gradients")
220         return list(gradient.value() if gradient is not None else gradient for gradient in self._gradients)
221 
222     def __call__(self, gradients):
223         """Accumulates :obj:`gradients` on the current replica."""
224         if not self._gradients:
225             _ = self.step  # Create the step variable.
226             self._gradients.extend(
227                 [
228                     tf.Variable(
229                         tf.zeros_like(gradient),
230                         trainable=False,
231                         synchronization=tf.VariableSynchronization.ON_READ,
232                         aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
233                     )
234                     if gradient is not None
235                     else gradient
236                     for gradient in gradients
237                 ]
238             )
239         if len(gradients) != len(self._gradients):
240             raise ValueError("Expected %s gradients, but got %d" % (len(self._gradients), len(gradients)))
241 
242         for accum_gradient, gradient in zip(self._gradients, gradients):
243             if accum_gradient is not None and gradient is not None:
244                 accum_gradient.assign_add(gradient)
245 
246         self._accum_steps.assign_add(1)
247 
248     def reset(self):
249         """Resets the accumulated gradients on the current replica."""
250         if not self._gradients:
251             return
252         self._accum_steps.assign(0)
253         for gradient in self._gradients:
254             if gradient is not None:
255                 gradient.assign(tf.zeros_like(gradient))
256
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/transformers/optimization_tf.py b/src/transformers/optimization_tf.py
--- a/src/transformers/optimization_tf.py
+++ b/src/transformers/optimization_tf.py
@@ -75,7 +75,7 @@
         beta_1=0.9,
         beta_2=0.999,
         epsilon=1e-6,
-        exclude_from_weight_decay=["layer_norm", "bias"],
+        exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"],
     )
 
     return optimizer

{"golden_diff": "diff --git a/src/transformers/optimization_tf.py b/src/transformers/optimization_tf.py\n--- a/src/transformers/optimization_tf.py\n+++ b/src/transformers/optimization_tf.py\n@@ -75,7 +75,7 @@\n         beta_1=0.9,\n         beta_2=0.999,\n         epsilon=1e-6,\n-        exclude_from_weight_decay=[\"layer_norm\", \"bias\"],\n+        exclude_from_weight_decay=[\"LayerNorm\", \"layer_norm\", \"bias\"],\n     )\n \n     return optimizer\n", "issue": "LayerNorm not excluded from weight decay in TF\n# \ud83d\udc1b Bug\r\n\r\n## Information\r\n\r\nModel I am using (Bert, XLNet ...):\r\nbert-base-cased\r\n\r\nLanguage I am using the model on (English, Chinese ...):\r\nEnglish\r\n\r\nThe problem arises when using:\r\n* [X] the official example scripts: (give details below)\r\n* [ ] my own modified scripts: (give details below)\r\n\r\nThe tasks I am working on is:\r\n* [X] an official GLUE/SQUaD task: (give the name)\r\n* [ ] my own task or dataset: (give details below)\r\n\r\n## To reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\n1. Add a print statement to `_do_use_weight_decay` in [AdamWeightDecay](https://github.com/huggingface/transformers/blob/master/src/transformers/optimization_tf.py) to see which parameters are actually excluded:\r\n```python\r\n    def _do_use_weight_decay(self, param_name):\r\n        \"\"\"Whether to use L2 weight decay for `param_name`.\"\"\"\r\n        if self.weight_decay_rate == 0:\r\n            return False\r\n\r\n        if self._include_in_weight_decay:\r\n            for r in self._include_in_weight_decay:\r\n                if re.search(r, param_name) is not None:\r\n                    return True\r\n\r\n        if self._exclude_from_weight_decay:\r\n            for r in self._exclude_from_weight_decay:\r\n                if re.search(r, param_name) is not None:\r\n                    print(f\"Found: {param_name}\")\r\n                    return False\r\n        return True\r\n```\r\n2. run `python examples/text-classification/run_tf_glue.py --model_name_or_path bert-base-cased --task_name mrpc --output_dir temp --logging_dir temp --do_train --overwrite_output_dir --optimizer_name adamw`. \r\n3. Observe that no weights related to layer norms are printed. \r\n\r\n<!-- If you have code snippets, error messages, stack traces please provide them here as well.\r\n     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting\r\n     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->\r\n\r\n## Expected behavior\r\n\r\n<!-- A clear and concise description of what you would expect to happen. -->\r\nThe weights of the layer norms (and the biases) should be printed. \r\n\r\nSee for example: https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py. \r\n\r\nBased on the fact that no layer norm weights are printed with \"layer_norm\" simply switching \"layer_norm\" to \"LayerNorm\" seems like the easiest change. \r\n\r\n## Environment info\r\n<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.\r\n     Don't forget to fill out the missing fields in that output! -->\r\n     \r\n- `transformers` version: 2.9.0\r\n- Platform: Darwin-19.4.0-x86_64-i386-64bit\r\n- Python version: 3.7.7\r\n- PyTorch version (GPU?): 1.5.0 (False)\r\n- Tensorflow version (GPU?): 2.2.0 (False)\r\n- Using GPU in script?: No\r\n- Using distributed or parallel set-up in script?: No\r\n\n", "before_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Functions and classes related to optimization (weight updates).\"\"\"\n\n\nimport re\n\nimport tensorflow as tf\n\n\nclass WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):\n    \"\"\"Applies a warmup schedule on a given learning rate decay schedule.\"\"\"\n\n    def __init__(\n        self, initial_learning_rate, decay_schedule_fn, warmup_steps, power=1.0, name=None,\n    ):\n        super().__init__()\n        self.initial_learning_rate = initial_learning_rate\n        self.warmup_steps = warmup_steps\n        self.power = power\n        self.decay_schedule_fn = decay_schedule_fn\n        self.name = name\n\n    def __call__(self, step):\n        with tf.name_scope(self.name or \"WarmUp\") as name:\n            # Implements polynomial warmup. i.e., if global_step < warmup_steps, the\n            # learning rate will be `global_step/num_warmup_steps * init_lr`.\n            global_step_float = tf.cast(step, tf.float32)\n            warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)\n            warmup_percent_done = global_step_float / warmup_steps_float\n            warmup_learning_rate = self.initial_learning_rate * tf.math.pow(warmup_percent_done, self.power)\n            return tf.cond(\n                global_step_float < warmup_steps_float,\n                lambda: warmup_learning_rate,\n                lambda: self.decay_schedule_fn(step),\n                name=name,\n            )\n\n    def get_config(self):\n        return {\n            \"initial_learning_rate\": self.initial_learning_rate,\n            \"decay_schedule_fn\": self.decay_schedule_fn,\n            \"warmup_steps\": self.warmup_steps,\n            \"power\": self.power,\n            \"name\": self.name,\n        }\n\n\ndef create_optimizer(init_lr, num_train_steps, num_warmup_steps, end_lr=0.0, optimizer_type=\"adamw\"):\n    \"\"\"Creates an optimizer with learning rate schedule.\"\"\"\n    # Implements linear decay of the learning rate.\n    lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(\n        initial_learning_rate=init_lr, decay_steps=num_train_steps, end_learning_rate=end_lr,\n    )\n    if num_warmup_steps:\n        lr_schedule = WarmUp(\n            initial_learning_rate=init_lr, decay_schedule_fn=lr_schedule, warmup_steps=num_warmup_steps,\n        )\n\n    optimizer = AdamWeightDecay(\n        learning_rate=lr_schedule,\n        weight_decay_rate=0.01,\n        beta_1=0.9,\n        beta_2=0.999,\n        epsilon=1e-6,\n        exclude_from_weight_decay=[\"layer_norm\", \"bias\"],\n    )\n\n    return optimizer\n\n\nclass AdamWeightDecay(tf.keras.optimizers.Adam):\n    \"\"\"Adam enables L2 weight decay and clip_by_global_norm on gradients.\n  Just adding the square of the weights to the loss function is *not* the\n  correct way of using L2 regularization/weight decay with Adam, since that will\n  interact with the m and v parameters in strange ways.\n  Instead we want ot decay the weights in a manner that doesn't interact with\n  the m/v parameters. This is equivalent to adding the square of the weights to\n  the loss with plain (non-momentum) SGD.\n  \"\"\"\n\n    def __init__(\n        self,\n        learning_rate=0.001,\n        beta_1=0.9,\n        beta_2=0.999,\n        epsilon=1e-7,\n        amsgrad=False,\n        weight_decay_rate=0.0,\n        include_in_weight_decay=None,\n        exclude_from_weight_decay=None,\n        name=\"AdamWeightDecay\",\n        **kwargs\n    ):\n        super().__init__(learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs)\n        self.weight_decay_rate = weight_decay_rate\n        self._include_in_weight_decay = include_in_weight_decay\n        self._exclude_from_weight_decay = exclude_from_weight_decay\n\n    @classmethod\n    def from_config(cls, config):\n        \"\"\"Creates an optimizer from its config with WarmUp custom object.\"\"\"\n        custom_objects = {\"WarmUp\": WarmUp}\n        return super(AdamWeightDecay, cls).from_config(config, custom_objects=custom_objects)\n\n    def _prepare_local(self, var_device, var_dtype, apply_state):\n        super(AdamWeightDecay, self)._prepare_local(var_device, var_dtype, apply_state)\n        apply_state[(var_device, var_dtype)][\"weight_decay_rate\"] = tf.constant(\n            self.weight_decay_rate, name=\"adam_weight_decay_rate\"\n        )\n\n    def _decay_weights_op(self, var, learning_rate, apply_state):\n        do_decay = self._do_use_weight_decay(var.name)\n        if do_decay:\n            return var.assign_sub(\n                learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)][\"weight_decay_rate\"],\n                use_locking=self._use_locking,\n            )\n        return tf.no_op()\n\n    def apply_gradients(self, grads_and_vars, name=None):\n        grads, tvars = list(zip(*grads_and_vars))\n        return super(AdamWeightDecay, self).apply_gradients(zip(grads, tvars), name=name,)\n\n    def _get_lr(self, var_device, var_dtype, apply_state):\n        \"\"\"Retrieves the learning rate with the given state.\"\"\"\n        if apply_state is None:\n            return self._decayed_lr_t[var_dtype], {}\n\n        apply_state = apply_state or {}\n        coefficients = apply_state.get((var_device, var_dtype))\n        if coefficients is None:\n            coefficients = self._fallback_apply_state(var_device, var_dtype)\n            apply_state[(var_device, var_dtype)] = coefficients\n\n        return coefficients[\"lr_t\"], dict(apply_state=apply_state)\n\n    def _resource_apply_dense(self, grad, var, apply_state=None):\n        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)\n        decay = self._decay_weights_op(var, lr_t, apply_state)\n        with tf.control_dependencies([decay]):\n            return super(AdamWeightDecay, self)._resource_apply_dense(grad, var, **kwargs)\n\n    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):\n        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)\n        decay = self._decay_weights_op(var, lr_t, apply_state)\n        with tf.control_dependencies([decay]):\n            return super(AdamWeightDecay, self)._resource_apply_sparse(grad, var, indices, **kwargs)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update({\"weight_decay_rate\": self.weight_decay_rate})\n        return config\n\n    def _do_use_weight_decay(self, param_name):\n        \"\"\"Whether to use L2 weight decay for `param_name`.\"\"\"\n        if self.weight_decay_rate == 0:\n            return False\n\n        if self._include_in_weight_decay:\n            for r in self._include_in_weight_decay:\n                if re.search(r, param_name) is not None:\n                    return True\n\n        if self._exclude_from_weight_decay:\n            for r in self._exclude_from_weight_decay:\n                if re.search(r, param_name) is not None:\n                    return False\n        return True\n\n\n# Extracted from https://github.com/OpenNMT/OpenNMT-tf/blob/master/opennmt/optimizers/utils.py\nclass GradientAccumulator(object):\n    \"\"\"Gradient accumulation utility.\n  When used with a distribution strategy, the accumulator should be called in a\n  replica context. Gradients will be accumulated locally on each replica and\n  without synchronization. Users should then call ``.gradients``, scale the\n  gradients if required, and pass the result to ``apply_gradients``.\n  \"\"\"\n\n    # We use the ON_READ synchronization policy so that no synchronization is\n    # performed on assignment. To get the value, we call .value() which returns the\n    # value on the current replica without synchronization.\n\n    def __init__(self):\n        \"\"\"Initializes the accumulator.\"\"\"\n        self._gradients = []\n        self._accum_steps = None\n\n    @property\n    def step(self):\n        \"\"\"Number of accumulated steps.\"\"\"\n        if self._accum_steps is None:\n            self._accum_steps = tf.Variable(\n                tf.constant(0, dtype=tf.int64),\n                trainable=False,\n                synchronization=tf.VariableSynchronization.ON_READ,\n                aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,\n            )\n\n        return self._accum_steps.value()\n\n    @property\n    def gradients(self):\n        \"\"\"The accumulated gradients on the current replica.\"\"\"\n        if not self._gradients:\n            raise ValueError(\"The accumulator should be called first to initialize the gradients\")\n        return list(gradient.value() if gradient is not None else gradient for gradient in self._gradients)\n\n    def __call__(self, gradients):\n        \"\"\"Accumulates :obj:`gradients` on the current replica.\"\"\"\n        if not self._gradients:\n            _ = self.step  # Create the step variable.\n            self._gradients.extend(\n                [\n                    tf.Variable(\n                        tf.zeros_like(gradient),\n                        trainable=False,\n                        synchronization=tf.VariableSynchronization.ON_READ,\n                        aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,\n                    )\n                    if gradient is not None\n                    else gradient\n                    for gradient in gradients\n                ]\n            )\n        if len(gradients) != len(self._gradients):\n            raise ValueError(\"Expected %s gradients, but got %d\" % (len(self._gradients), len(gradients)))\n\n        for accum_gradient, gradient in zip(self._gradients, gradients):\n            if accum_gradient is not None and gradient is not None:\n                accum_gradient.assign_add(gradient)\n\n        self._accum_steps.assign_add(1)\n\n    def reset(self):\n        \"\"\"Resets the accumulated gradients on the current replica.\"\"\"\n        if not self._gradients:\n            return\n        self._accum_steps.assign(0)\n        for gradient in self._gradients:\n            if gradient is not None:\n                gradient.assign(tf.zeros_like(gradient))\n", "path": "src/transformers/optimization_tf.py"}], "after_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Functions and classes related to optimization (weight updates).\"\"\"\n\n\nimport re\n\nimport tensorflow as tf\n\n\nclass WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):\n    \"\"\"Applies a warmup schedule on a given learning rate decay schedule.\"\"\"\n\n    def __init__(\n        self, initial_learning_rate, decay_schedule_fn, warmup_steps, power=1.0, name=None,\n    ):\n        super().__init__()\n        self.initial_learning_rate = initial_learning_rate\n        self.warmup_steps = warmup_steps\n        self.power = power\n        self.decay_schedule_fn = decay_schedule_fn\n        self.name = name\n\n    def __call__(self, step):\n        with tf.name_scope(self.name or \"WarmUp\") as name:\n            # Implements polynomial warmup. i.e., if global_step < warmup_steps, the\n            # learning rate will be `global_step/num_warmup_steps * init_lr`.\n            global_step_float = tf.cast(step, tf.float32)\n            warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)\n            warmup_percent_done = global_step_float / warmup_steps_float\n            warmup_learning_rate = self.initial_learning_rate * tf.math.pow(warmup_percent_done, self.power)\n            return tf.cond(\n                global_step_float < warmup_steps_float,\n                lambda: warmup_learning_rate,\n                lambda: self.decay_schedule_fn(step),\n                name=name,\n            )\n\n    def get_config(self):\n        return {\n            \"initial_learning_rate\": self.initial_learning_rate,\n            \"decay_schedule_fn\": self.decay_schedule_fn,\n            \"warmup_steps\": self.warmup_steps,\n            \"power\": self.power,\n            \"name\": self.name,\n        }\n\n\ndef create_optimizer(init_lr, num_train_steps, num_warmup_steps, end_lr=0.0, optimizer_type=\"adamw\"):\n    \"\"\"Creates an optimizer with learning rate schedule.\"\"\"\n    # Implements linear decay of the learning rate.\n    lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(\n        initial_learning_rate=init_lr, decay_steps=num_train_steps, end_learning_rate=end_lr,\n    )\n    if num_warmup_steps:\n        lr_schedule = WarmUp(\n            initial_learning_rate=init_lr, decay_schedule_fn=lr_schedule, warmup_steps=num_warmup_steps,\n        )\n\n    optimizer = AdamWeightDecay(\n        learning_rate=lr_schedule,\n        weight_decay_rate=0.01,\n        beta_1=0.9,\n        beta_2=0.999,\n        epsilon=1e-6,\n        exclude_from_weight_decay=[\"LayerNorm\", \"layer_norm\", \"bias\"],\n    )\n\n    return optimizer\n\n\nclass AdamWeightDecay(tf.keras.optimizers.Adam):\n    \"\"\"Adam enables L2 weight decay and clip_by_global_norm on gradients.\n  Just adding the square of the weights to the loss function is *not* the\n  correct way of using L2 regularization/weight decay with Adam, since that will\n  interact with the m and v parameters in strange ways.\n  Instead we want ot decay the weights in a manner that doesn't interact with\n  the m/v parameters. This is equivalent to adding the square of the weights to\n  the loss with plain (non-momentum) SGD.\n  \"\"\"\n\n    def __init__(\n        self,\n        learning_rate=0.001,\n        beta_1=0.9,\n        beta_2=0.999,\n        epsilon=1e-7,\n        amsgrad=False,\n        weight_decay_rate=0.0,\n        include_in_weight_decay=None,\n        exclude_from_weight_decay=None,\n        name=\"AdamWeightDecay\",\n        **kwargs\n    ):\n        super().__init__(learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs)\n        self.weight_decay_rate = weight_decay_rate\n        self._include_in_weight_decay = include_in_weight_decay\n        self._exclude_from_weight_decay = exclude_from_weight_decay\n\n    @classmethod\n    def from_config(cls, config):\n        \"\"\"Creates an optimizer from its config with WarmUp custom object.\"\"\"\n        custom_objects = {\"WarmUp\": WarmUp}\n        return super(AdamWeightDecay, cls).from_config(config, custom_objects=custom_objects)\n\n    def _prepare_local(self, var_device, var_dtype, apply_state):\n        super(AdamWeightDecay, self)._prepare_local(var_device, var_dtype, apply_state)\n        apply_state[(var_device, var_dtype)][\"weight_decay_rate\"] = tf.constant(\n            self.weight_decay_rate, name=\"adam_weight_decay_rate\"\n        )\n\n    def _decay_weights_op(self, var, learning_rate, apply_state):\n        do_decay = self._do_use_weight_decay(var.name)\n        if do_decay:\n            return var.assign_sub(\n                learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)][\"weight_decay_rate\"],\n                use_locking=self._use_locking,\n            )\n        return tf.no_op()\n\n    def apply_gradients(self, grads_and_vars, name=None):\n        grads, tvars = list(zip(*grads_and_vars))\n        return super(AdamWeightDecay, self).apply_gradients(zip(grads, tvars), name=name,)\n\n    def _get_lr(self, var_device, var_dtype, apply_state):\n        \"\"\"Retrieves the learning rate with the given state.\"\"\"\n        if apply_state is None:\n            return self._decayed_lr_t[var_dtype], {}\n\n        apply_state = apply_state or {}\n        coefficients = apply_state.get((var_device, var_dtype))\n        if coefficients is None:\n            coefficients = self._fallback_apply_state(var_device, var_dtype)\n            apply_state[(var_device, var_dtype)] = coefficients\n\n        return coefficients[\"lr_t\"], dict(apply_state=apply_state)\n\n    def _resource_apply_dense(self, grad, var, apply_state=None):\n        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)\n        decay = self._decay_weights_op(var, lr_t, apply_state)\n        with tf.control_dependencies([decay]):\n            return super(AdamWeightDecay, self)._resource_apply_dense(grad, var, **kwargs)\n\n    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):\n        lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)\n        decay = self._decay_weights_op(var, lr_t, apply_state)\n        with tf.control_dependencies([decay]):\n            return super(AdamWeightDecay, self)._resource_apply_sparse(grad, var, indices, **kwargs)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update({\"weight_decay_rate\": self.weight_decay_rate})\n        return config\n\n    def _do_use_weight_decay(self, param_name):\n        \"\"\"Whether to use L2 weight decay for `param_name`.\"\"\"\n        if self.weight_decay_rate == 0:\n            return False\n\n        if self._include_in_weight_decay:\n            for r in self._include_in_weight_decay:\n                if re.search(r, param_name) is not None:\n                    return True\n\n        if self._exclude_from_weight_decay:\n            for r in self._exclude_from_weight_decay:\n                if re.search(r, param_name) is not None:\n                    return False\n        return True\n\n\n# Extracted from https://github.com/OpenNMT/OpenNMT-tf/blob/master/opennmt/optimizers/utils.py\nclass GradientAccumulator(object):\n    \"\"\"Gradient accumulation utility.\n  When used with a distribution strategy, the accumulator should be called in a\n  replica context. Gradients will be accumulated locally on each replica and\n  without synchronization. Users should then call ``.gradients``, scale the\n  gradients if required, and pass the result to ``apply_gradients``.\n  \"\"\"\n\n    # We use the ON_READ synchronization policy so that no synchronization is\n    # performed on assignment. To get the value, we call .value() which returns the\n    # value on the current replica without synchronization.\n\n    def __init__(self):\n        \"\"\"Initializes the accumulator.\"\"\"\n        self._gradients = []\n        self._accum_steps = None\n\n    @property\n    def step(self):\n        \"\"\"Number of accumulated steps.\"\"\"\n        if self._accum_steps is None:\n            self._accum_steps = tf.Variable(\n                tf.constant(0, dtype=tf.int64),\n                trainable=False,\n                synchronization=tf.VariableSynchronization.ON_READ,\n                aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,\n            )\n\n        return self._accum_steps.value()\n\n    @property\n    def gradients(self):\n        \"\"\"The accumulated gradients on the current replica.\"\"\"\n        if not self._gradients:\n            raise ValueError(\"The accumulator should be called first to initialize the gradients\")\n        return list(gradient.value() for gradient in self._gradients)\n\n    def __call__(self, gradients):\n        \"\"\"Accumulates :obj:`gradients` on the current replica.\"\"\"\n        if not self._gradients:\n            _ = self.step  # Create the step variable.\n            self._gradients.extend(\n                [\n                    tf.Variable(\n                        tf.zeros_like(gradient),\n                        trainable=False,\n                        synchronization=tf.VariableSynchronization.ON_READ,\n                        aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,\n                    )\n                    for gradient in gradients\n                ]\n            )\n        if len(gradients) != len(self._gradients):\n            raise ValueError(\"Expected %s gradients, but got %d\" % (len(self._gradients), len(gradients)))\n\n        for accum_gradient, gradient in zip(self._gradients, gradients):\n            accum_gradient.assign_add(gradient)\n\n        self._accum_steps.assign_add(1)\n\n    def reset(self):\n        \"\"\"Resets the accumulated gradients on the current replica.\"\"\"\n        if not self._gradients:\n            return\n        self._accum_steps.assign(0)\n        for gradient in self._gradients:\n            gradient.assign(tf.zeros_like(gradient))\n", "path": "src/transformers/optimization_tf.py"}]}