Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/adamax.py: 33%

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.""" 

16# pylint: disable=g-classes-have-attributes 

17 

18from tensorflow.python.framework import dtypes 

19from tensorflow.python.framework import ops 

20from tensorflow.python.framework import tensor_conversion 

21from tensorflow.python.keras import backend_config 

22from tensorflow.python.keras.optimizer_v2 import optimizer_v2 

23from tensorflow.python.ops import array_ops 

24from tensorflow.python.ops import control_flow_ops 

25from tensorflow.python.ops import math_ops 

26from tensorflow.python.training import gen_training_ops 

27from tensorflow.python.util.tf_export import keras_export 

28 

29 

30@keras_export('keras.optimizers.Adamax') 

31class Adamax(optimizer_v2.OptimizerV2): 

32 """Optimizer that implements the Adamax algorithm. 

33 

34 It is a variant of Adam based on the infinity norm. 

35 Default parameters follow those provided in the paper. 

36 Adamax is sometimes superior to adam, specially in models with embeddings. 

37 

38 Initialization: 

39 

40 ```python 

41 m = 0 # Initialize initial 1st moment vector 

42 v = 0 # Initialize the exponentially weighted infinity norm 

43 t = 0 # Initialize timestep 

44 ``` 

45 

46 The update rule for parameter `w` with gradient `g` is 

47 described at the end of section 7.1 of the paper: 

48 

49 ```python 

50 t += 1 

51 m = beta1 * m + (1 - beta) * g 

52 v = max(beta2 * v, abs(g)) 

53 current_lr = learning_rate / (1 - beta1 ** t) 

54 w = w - current_lr * m / (v + epsilon) 

55 ``` 

56 

57 Similarly to `Adam`, the epsilon is added for numerical stability 

58 (especially to get rid of division by zero when `v_t == 0`). 

59 

60 In contrast to `Adam`, the sparse implementation of this algorithm 

61 (used when the gradient is an IndexedSlices object, typically because of 

62 `tf.gather` or an embedding lookup in the forward pass) only updates 

63 variable slices and corresponding `m_t`, `v_t` terms when that part of 

64 the variable was used in the forward pass. This means that the sparse 

65 behavior is contrast to the dense behavior (similar to some momentum 

66 implementations which ignore momentum unless a variable slice was actually 

67 used). 

68 

69 Args: 

70 learning_rate: A `Tensor`, floating point value, or a schedule that is a 

71 `tf.keras.optimizers.schedules.LearningRateSchedule`. The learning rate. 

72 beta_1: A float value or a constant float tensor. The exponential decay 

73 rate for the 1st moment estimates. 

74 beta_2: A float value or a constant float tensor. The exponential decay 

75 rate for the exponentially weighted infinity norm. 

76 epsilon: A small constant for numerical stability. 

77 name: Optional name for the operations created when applying gradients. 

78 Defaults to `"Adamax"`. 

79 **kwargs: Keyword arguments. Allowed to be one of 

80 `"clipnorm"` or `"clipvalue"`. 

81 `"clipnorm"` (float) clips gradients by norm; `"clipvalue"` (float) clips 

82 gradients by value. 

83 

84 Reference: 

85 - [Kingma et al., 2014](http://arxiv.org/abs/1412.6980) 

86 """ 

87 

88 _HAS_AGGREGATE_GRAD = True 

89 

90 def __init__(self, 

91 learning_rate=0.001, 

92 beta_1=0.9, 

93 beta_2=0.999, 

94 epsilon=1e-7, 

95 name='Adamax', 

96 **kwargs): 

97 super(Adamax, self).__init__(name, **kwargs) 

98 self._set_hyper('learning_rate', kwargs.get('lr', learning_rate)) 

99 self._set_hyper('decay', self._initial_decay) 

100 self._set_hyper('beta_1', beta_1) 

101 self._set_hyper('beta_2', beta_2) 

102 self.epsilon = epsilon or backend_config.epsilon() 

103 

104 def _create_slots(self, var_list): 

105 # Separate for-loops to respect the ordering of slot variables from v1. 

106 for var in var_list: 

107 self.add_slot(var, 'm') # Create slots for the first moments. 

108 for var in var_list: 

109 self.add_slot(var, 'v') # Create slots for the second moments. 

110 

111 def _prepare_local(self, var_device, var_dtype, apply_state): 

112 super(Adamax, self)._prepare_local(var_device, var_dtype, apply_state) 

113 

114 local_step = math_ops.cast(self.iterations + 1, var_dtype) 

115 beta_1_t = array_ops.identity(self._get_hyper('beta_1', var_dtype)) 

116 beta_2_t = array_ops.identity(self._get_hyper('beta_2', var_dtype)) 

117 beta_1_power = math_ops.pow(beta_1_t, local_step) 

118 lr_t = apply_state[(var_device, var_dtype)]['lr_t'] 

119 

120 apply_state[(var_device, var_dtype)].update( 

121 dict( 

122 neg_scaled_lr=-lr_t / (1 - beta_1_power), 

123 epsilon=tensor_conversion.convert_to_tensor_v2_with_dispatch( 

124 self.epsilon, var_dtype 

125 ), 

126 beta_1_t=beta_1_t, 

127 beta_1_power=beta_1_power, 

128 one_minus_beta_1_t=1 - beta_1_t, 

129 beta_2_t=beta_2_t, 

130 zero=array_ops.zeros((), dtype=dtypes.int64), 

131 ) 

132 ) 

133 

134 def _resource_apply_dense(self, grad, var, apply_state=None): 

135 var_device, var_dtype = var.device, var.dtype.base_dtype 

136 coefficients = ((apply_state or {}).get((var_device, var_dtype)) 

137 or self._fallback_apply_state(var_device, var_dtype)) 

138 

139 m = self.get_slot(var, 'm') 

140 v = self.get_slot(var, 'v') 

141 return gen_training_ops.ResourceApplyAdaMax( 

142 var=var.handle, 

143 m=m.handle, 

144 v=v.handle, 

145 beta1_power=coefficients['beta_1_power'], 

146 lr=coefficients['lr_t'], 

147 beta1=coefficients['beta_1_t'], 

148 beta2=coefficients['beta_2_t'], 

149 epsilon=coefficients['epsilon'], 

150 grad=grad, 

151 use_locking=self._use_locking) 

152 

153 def _resource_apply_sparse(self, grad, var, indices, apply_state=None): 

154 var_device, var_dtype = var.device, var.dtype.base_dtype 

155 coefficients = ((apply_state or {}).get((var_device, var_dtype)) 

156 or self._fallback_apply_state(var_device, var_dtype)) 

157 

158 # m_t = beta1 * m + (1 - beta1) * g_t 

159 m = self.get_slot(var, 'm') 

160 m_slice = array_ops.gather(m, indices, axis=coefficients['zero']) 

161 m_t_slice = (m_slice * coefficients['beta_1_t'] + 

162 grad * coefficients['one_minus_beta_1_t']) 

163 with ops.control_dependencies([m_t_slice]): 

164 m_t = self._resource_scatter_update(m, indices, m_t_slice) 

165 

166 # u_t = max(beta2 * u, abs(g_t)) 

167 v = self.get_slot(var, 'v') 

168 v_slice = array_ops.gather(v, indices, axis=coefficients['zero']) 

169 v_t_slice = math_ops.maximum(v_slice * coefficients['beta_2_t'], 

170 math_ops.abs(grad)) 

171 with ops.control_dependencies([v_t_slice]): 

172 v_t = self._resource_scatter_update(v, indices, v_t_slice) 

173 # theta_t = theta - lr / (1 - beta1^t) * m_t / u_t 

174 var_slice = coefficients['neg_scaled_lr'] * ( 

175 m_t_slice / (v_t_slice + coefficients['epsilon'])) 

176 with ops.control_dependencies([var_slice]): 

177 var_update = self._resource_scatter_add(var, indices, var_slice) 

178 return control_flow_ops.group(*[var_update, m_t, v_t]) 

179 

180 def get_config(self): 

181 config = super(Adamax, self).get_config() 

182 config.update({ 

183 'learning_rate': self._serialize_hyperparameter('learning_rate'), 

184 'decay': self._initial_decay, 

185 'beta_1': self._serialize_hyperparameter('beta_1'), 

186 'beta_2': self._serialize_hyperparameter('beta_2'), 

187 'epsilon': self.epsilon, 

188 }) 

189 return config