Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/keras/src/optimizers/legacy/adamax.py: 25%
55 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2018 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"""Adamax optimizer implementation."""
17import tensorflow.compat.v2 as tf
19from keras.src import backend_config
20from keras.src.optimizers.legacy import optimizer_v2
22# isort: off
23from tensorflow.python.util.tf_export import keras_export
26@keras_export(
27 "keras.optimizers.legacy.Adamax",
28 v1=["keras.optimizers.Adamax", "keras.optimizers.legacy.Adamax"],
29)
30class Adamax(optimizer_v2.OptimizerV2):
31 """Optimizer that implements the Adamax algorithm.
33 It is a variant of Adam based on the infinity norm.
34 Default parameters follow those provided in the paper.
35 Adamax is sometimes superior to adam, specially in models with embeddings.
37 Initialization:
39 ```python
40 m = 0 # Initialize initial 1st moment vector
41 v = 0 # Initialize the exponentially weighted infinity norm
42 t = 0 # Initialize timestep
43 ```
45 The update rule for parameter `w` with gradient `g` is
46 described at the end of section 7.1 of the paper:
48 ```python
49 t += 1
50 m = beta1 * m + (1 - beta) * g
51 v = max(beta2 * v, abs(g))
52 current_lr = learning_rate / (1 - beta1 ** t)
53 w = w - current_lr * m / (v + epsilon)
54 ```
56 Similarly to `Adam`, the epsilon is added for numerical stability
57 (especially to get rid of division by zero when `v_t == 0`).
59 In contrast to `Adam`, the sparse implementation of this algorithm
60 (used when the gradient is an IndexedSlices object, typically because of
61 `tf.gather` or an embedding lookup in the forward pass) only updates
62 variable slices and corresponding `m_t`, `v_t` terms when that part of
63 the variable was used in the forward pass. This means that the sparse
64 behavior is contrast to the dense behavior (similar to some momentum
65 implementations which ignore momentum unless a variable slice was actually
66 used).
68 Args:
69 learning_rate: A `Tensor`, floating point value, or a schedule that is a
70 `tf.keras.optimizers.schedules.LearningRateSchedule`. The learning rate.
71 beta_1: A float value or a constant float tensor. The exponential decay
72 rate for the 1st moment estimates.
73 beta_2: A float value or a constant float tensor. The exponential decay
74 rate for the exponentially weighted infinity norm.
75 epsilon: A small constant for numerical stability.
76 name: Optional name for the operations created when applying gradients.
77 Defaults to `"Adamax"`.
78 **kwargs: keyword arguments. Allowed arguments are `clipvalue`,
79 `clipnorm`, `global_clipnorm`.
80 If `clipvalue` (float) is set, the gradient of each weight
81 is clipped to be no higher than this value.
82 If `clipnorm` (float) is set, the gradient of each weight
83 is individually clipped so that its norm is no higher than this value.
84 If `global_clipnorm` (float) is set the gradient of all weights is
85 clipped so that their global norm is no higher than this value.
87 Reference:
88 - [Kingma et al., 2014](http://arxiv.org/abs/1412.6980)
89 """
91 _HAS_AGGREGATE_GRAD = True
93 def __init__(
94 self,
95 learning_rate=0.001,
96 beta_1=0.9,
97 beta_2=0.999,
98 epsilon=1e-7,
99 name="Adamax",
100 **kwargs
101 ):
102 super().__init__(name, **kwargs)
103 self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
104 self._set_hyper("decay", self._initial_decay)
105 self._set_hyper("beta_1", beta_1)
106 self._set_hyper("beta_2", beta_2)
107 self.epsilon = epsilon or backend_config.epsilon()
109 def _create_slots(self, var_list):
110 # Separate for-loops to respect the ordering of slot variables from v1.
111 for var in var_list:
112 self.add_slot(var, "m") # Create slots for the first moments.
113 for var in var_list:
114 self.add_slot(var, "v") # Create slots for the second moments.
116 def _prepare_local(self, var_device, var_dtype, apply_state):
117 super()._prepare_local(var_device, var_dtype, apply_state)
119 local_step = tf.cast(self.iterations + 1, var_dtype)
120 beta_1_t = tf.identity(self._get_hyper("beta_1", var_dtype))
121 beta_2_t = tf.identity(self._get_hyper("beta_2", var_dtype))
122 beta_1_power = tf.pow(beta_1_t, local_step)
123 lr_t = apply_state[(var_device, var_dtype)]["lr_t"]
125 apply_state[(var_device, var_dtype)].update(
126 dict(
127 neg_scaled_lr=-lr_t / (1 - beta_1_power),
128 epsilon=tf.convert_to_tensor(self.epsilon, var_dtype),
129 beta_1_t=beta_1_t,
130 beta_1_power=beta_1_power,
131 one_minus_beta_1_t=1 - beta_1_t,
132 beta_2_t=beta_2_t,
133 zero=tf.zeros((), dtype=tf.int64),
134 )
135 )
137 def _resource_apply_dense(self, grad, var, apply_state=None):
138 var_device, var_dtype = var.device, var.dtype.base_dtype
139 coefficients = (apply_state or {}).get(
140 (var_device, var_dtype)
141 ) or self._fallback_apply_state(var_device, var_dtype)
143 m = self.get_slot(var, "m")
144 v = self.get_slot(var, "v")
145 return tf.raw_ops.ResourceApplyAdaMax(
146 var=var.handle,
147 m=m.handle,
148 v=v.handle,
149 beta1_power=coefficients["beta_1_power"],
150 lr=coefficients["lr_t"],
151 beta1=coefficients["beta_1_t"],
152 beta2=coefficients["beta_2_t"],
153 epsilon=coefficients["epsilon"],
154 grad=grad,
155 use_locking=self._use_locking,
156 )
158 def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
159 var_device, var_dtype = var.device, var.dtype.base_dtype
160 coefficients = (apply_state or {}).get(
161 (var_device, var_dtype)
162 ) or self._fallback_apply_state(var_device, var_dtype)
164 # m_t = beta1 * m + (1 - beta1) * g_t
165 m = self.get_slot(var, "m")
166 m_slice = tf.gather(m, indices, axis=coefficients["zero"])
167 m_t_slice = (
168 m_slice * coefficients["beta_1_t"]
169 + grad * coefficients["one_minus_beta_1_t"]
170 )
171 with tf.control_dependencies([m_t_slice]):
172 m_t = self._resource_scatter_update(m, indices, m_t_slice)
174 # u_t = max(beta2 * u, abs(g_t))
175 v = self.get_slot(var, "v")
176 v_slice = tf.gather(v, indices, axis=coefficients["zero"])
177 v_t_slice = tf.maximum(v_slice * coefficients["beta_2_t"], tf.abs(grad))
178 with tf.control_dependencies([v_t_slice]):
179 v_t = self._resource_scatter_update(v, indices, v_t_slice)
180 # theta_t = theta - lr / (1 - beta1^t) * m_t / u_t
181 var_slice = coefficients["neg_scaled_lr"] * (
182 m_t_slice / (v_t_slice + coefficients["epsilon"])
183 )
184 with tf.control_dependencies([var_slice]):
185 var_update = self._resource_scatter_add(var, indices, var_slice)
186 return tf.group(*[var_update, m_t, v_t])
188 def get_config(self):
189 config = super().get_config()
190 config.update(
191 {
192 "learning_rate": self._serialize_hyperparameter(
193 "learning_rate"
194 ),
195 "decay": self._initial_decay,
196 "beta_1": self._serialize_hyperparameter("beta_1"),
197 "beta_2": self._serialize_hyperparameter("beta_2"),
198 "epsilon": self.epsilon,
199 }
200 )
201 return config