|
|
from tensorflow.python.keras.optimizers import Optimizer |
|
|
from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2 |
|
|
from tensorflow.python import ops, math_ops, state_ops, control_flow_ops |
|
|
from tensorflow.python.keras import backend as K |
|
|
from tensorflow.python.keras import backend_config |
|
|
import tensorflow as tf |
|
|
|
|
|
|
|
|
class AccumOptimizer(Optimizer): |
|
|
"""Optimizer |
|
|
Inheriting Optimizer class, wrapping the original optimizer |
|
|
to achieve a new corresponding optimizer of gradient accumulation. |
|
|
# Arguments |
|
|
optimizer: an instance of keras optimizer (supporting |
|
|
all keras optimizers currently available); |
|
|
steps_per_update: the steps of gradient accumulation |
|
|
# Returns |
|
|
a new keras optimizer. |
|
|
""" |
|
|
def __init__(self, optimizer, steps_per_update=1, **kwargs): |
|
|
super(AccumOptimizer, self).__init__(name='AccumOptimizer', **kwargs) |
|
|
self.optimizer = optimizer |
|
|
with K.name_scope(self.__class__.__name__): |
|
|
self.steps_per_update = steps_per_update |
|
|
self.iterations = K.variable(0, dtype='int64', name='iterations') |
|
|
self.cond = K.equal(self.iterations % self.steps_per_update, 0) |
|
|
self.lr = self.optimizer.lr |
|
|
self.optimizer.lr = K.switch(self.cond, self.optimizer.lr, 0.) |
|
|
for attr in ['momentum', 'rho', 'beta_1', 'beta_2']: |
|
|
if hasattr(self.optimizer, attr): |
|
|
value = getattr(self.optimizer, attr) |
|
|
setattr(self, attr, value) |
|
|
setattr(self.optimizer, attr, K.switch(self.cond, value, 1 - 1e-7)) |
|
|
for attr in self.optimizer.get_config(): |
|
|
if not hasattr(self, attr): |
|
|
value = getattr(self.optimizer, attr) |
|
|
setattr(self, attr, value) |
|
|
|
|
|
def get_gradients(loss, params): |
|
|
return [ag / self.steps_per_update for ag in self.accum_grads] |
|
|
self.optimizer.get_gradients = get_gradients |
|
|
def get_updates(self, loss, params): |
|
|
self.updates = [ |
|
|
K.update_add(self.iterations, 1), |
|
|
K.update_add(self.optimizer.iterations, K.cast(self.cond, 'int64')), |
|
|
] |
|
|
|
|
|
self.accum_grads = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params] |
|
|
grads = self.get_gradients(loss, params) |
|
|
for g, ag in zip(grads, self.accum_grads): |
|
|
self.updates.append(K.update(ag, K.switch(self.cond, g, ag + g))) |
|
|
|
|
|
self.updates.extend(self.optimizer.get_updates(loss, params)[1:]) |
|
|
self.weights.extend(self.optimizer.weights) |
|
|
return self.updates |
|
|
def get_config(self): |
|
|
iterations = K.eval(self.iterations) |
|
|
K.set_value(self.iterations, 0) |
|
|
config = self.optimizer.get_config() |
|
|
K.set_value(self.iterations, iterations) |
|
|
return config |
|
|
|
|
|
|
|
|
__all__ = ['AdamAccumulated'] |
|
|
|
|
|
|
|
|
|
|
|
class AdamAccumulated(OptimizerV2): |
|
|
"""Optimizer that implements the Adam algorithm with gradient accumulation.""" |
|
|
|
|
|
def __init__(self, |
|
|
accumulation_steps, |
|
|
learning_rate=0.001, |
|
|
beta_1=0.9, |
|
|
beta_2=0.999, |
|
|
epsilon=1e-7, |
|
|
amsgrad=False, |
|
|
name='Adam', |
|
|
**kwargs): |
|
|
r"""Construct a new Adam optimizer. |
|
|
Args: |
|
|
accumulation_steps: An integer. Update gradient in every accumulation steps. |
|
|
learning_rate: A Tensor or a floating point value. The learning rate. |
|
|
beta_1: A float value or a constant float tensor. The exponential decay |
|
|
rate for the 1st moment estimates. |
|
|
beta_2: A float value or a constant float tensor. The exponential decay |
|
|
rate for the 2nd moment estimates. |
|
|
epsilon: A small constant for numerical stability. This epsilon is |
|
|
"epsilon hat" in the Kingma and Ba paper (in the formula just before |
|
|
Section 2.1), not the epsilon in Algorithm 1 of the paper. |
|
|
amsgrad: boolean. Whether to apply AMSGrad variant of this algorithm from |
|
|
the paper "On the Convergence of Adam and beyond". |
|
|
name: Optional name for the operations created when applying gradients. |
|
|
Defaults to "Adam". @compatibility(eager) When eager execution is |
|
|
enabled, `learning_rate`, `beta_1`, `beta_2`, and `epsilon` can each be |
|
|
a callable that takes no arguments and returns the actual value to use. |
|
|
This can be useful for changing these values across different |
|
|
invocations of optimizer functions. @end_compatibility |
|
|
**kwargs: 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. |
|
|
""" |
|
|
|
|
|
super(AdamAccumulated, self).__init__(name, **kwargs) |
|
|
self._set_hyper('accumulation_steps', accumulation_steps) |
|
|
self._set_hyper('learning_rate', kwargs.get('lr', learning_rate)) |
|
|
self._set_hyper('decay', self._initial_decay) |
|
|
self._set_hyper('beta_1', beta_1) |
|
|
self._set_hyper('beta_2', beta_2) |
|
|
self.epsilon = epsilon or backend_config.epsilon() |
|
|
self.amsgrad = amsgrad |
|
|
|
|
|
def _create_slots(self, var_list): |
|
|
for var in var_list: |
|
|
self.add_slot(var, 'g') |
|
|
for var in var_list: |
|
|
self.add_slot(var, 'm') |
|
|
for var in var_list: |
|
|
self.add_slot(var, 'v') |
|
|
if self.amsgrad: |
|
|
for var in var_list: |
|
|
self.add_slot(var, 'vhat') |
|
|
|
|
|
def set_weights(self, weights): |
|
|
params = self.weights |
|
|
num_vars = int((len(params) - 1) / 2) |
|
|
if len(weights) == 3 * num_vars + 1: |
|
|
weights = weights[:len(params)] |
|
|
super(AdamAccumulated, self).set_weights(weights) |
|
|
|
|
|
def _resource_apply_dense(self, grad, var): |
|
|
var_dtype = var.dtype.base_dtype |
|
|
lr_t = self._decayed_lr(var_dtype) |
|
|
beta_1_t = self._get_hyper('beta_1', var_dtype) |
|
|
beta_2_t = self._get_hyper('beta_2', var_dtype) |
|
|
accumulation_steps = self._get_hyper('accumulation_steps', 'int64') |
|
|
update_cond = tf.equal((self.iterations + 1) % accumulation_steps, 0) |
|
|
sub_step = self.iterations % accumulation_steps + 1 |
|
|
local_step = math_ops.cast(self.iterations // accumulation_steps + 1, var_dtype) |
|
|
beta_1_power = math_ops.pow(beta_1_t, local_step) |
|
|
beta_2_power = math_ops.pow(beta_2_t, local_step) |
|
|
epsilon_t = ops.convert_to_tensor(self.epsilon, var_dtype) |
|
|
lr = (lr_t * math_ops.sqrt(1 - beta_2_power) / (1 - beta_1_power)) |
|
|
lr = tf.where(update_cond, lr, 0.0) |
|
|
|
|
|
g = self.get_slot(var, 'g') |
|
|
g_a = grad / math_ops.cast(accumulation_steps, var_dtype) |
|
|
g_t = tf.where(tf.equal(sub_step, 1), |
|
|
g_a, |
|
|
g + (g_a - g) / math_ops.cast(sub_step, var_dtype)) |
|
|
g_t = state_ops.assign(g, g_t, use_locking=self._use_locking) |
|
|
|
|
|
m = self.get_slot(var, 'm') |
|
|
m_t = tf.where(update_cond, m * beta_1_t + g_t * (1 - beta_1_t), m) |
|
|
m_t = state_ops.assign(m, m_t, use_locking=self._use_locking) |
|
|
|
|
|
v = self.get_slot(var, 'v') |
|
|
v_t = tf.where(update_cond, v * beta_2_t + (g_t * g_t) * (1 - beta_2_t), v) |
|
|
v_t = state_ops.assign(v, v_t, use_locking=self._use_locking) |
|
|
|
|
|
if not self.amsgrad: |
|
|
v_sqrt = math_ops.sqrt(v_t) |
|
|
var_update = state_ops.assign_sub( |
|
|
var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking) |
|
|
return control_flow_ops.group(*[var_update, m_t, v_t]) |
|
|
else: |
|
|
v_hat = self.get_slot(var, 'vhat') |
|
|
v_hat_t = tf.where(update_cond, math_ops.maximum(v_hat, v_t), v_hat) |
|
|
with ops.control_dependencies([v_hat_t]): |
|
|
v_hat_t = state_ops.assign( |
|
|
v_hat, v_hat_t, use_locking=self._use_locking) |
|
|
v_hat_sqrt = math_ops.sqrt(v_hat_t) |
|
|
var_update = state_ops.assign_sub( |
|
|
var, |
|
|
lr * m_t / (v_hat_sqrt + epsilon_t), |
|
|
use_locking=self._use_locking) |
|
|
return control_flow_ops.group(*[var_update, m_t, v_t, v_hat_t]) |
|
|
|
|
|
def get_config(self): |
|
|
config = super(AdamAccumulated, self).get_config() |
|
|
config.update({ |
|
|
'accumulation_steps': self._serialize_hyperparameter('accumulation_steps'), |
|
|
'learning_rate': self._serialize_hyperparameter('learning_rate'), |
|
|
'decay': self._serialize_hyperparameter('decay'), |
|
|
'beta_1': self._serialize_hyperparameter('beta_1'), |
|
|
'beta_2': self._serialize_hyperparameter('beta_2'), |
|
|
'epsilon': self.epsilon, |
|
|
'amsgrad': self.amsgrad, |
|
|
}) |
|
|
return config |
|
|
|
