def get_updates(self, loss, params):
# 更新判据
cond = K.equal(self.iterations % self.grad_accum_steps, 0)
cond = K.cast(cond, K.floatx())
# 获取梯度
grads = self.get_gradients(loss, params)
self.accum_grads = [
K.zeros(K.int_shape(p),
dtype=K.dtype(p),
name='accum_grad_%s' % i) for i, p in enumerate(params)
]
old_update = K.update
def new_update(x, new_x):
new_x = cond * new_x + (1 - cond) * x
return old_update(x, new_x)
K.update = new_update
updates = super(NewOptimizer, self).get_updates(loss, params)
K.update = old_update
# 累积梯度
with tf.control_dependencies(updates):
accum_updates = [
K.update(ag, g + (1 - cond) * ag)
for g, ag in zip(grads, self.accum_grads)
]
return accum_updates
def get_updates(self, loss, params):
grads = self.get_gradients(loss, params)
self.updates = [K.update_add(self.iterations, 1)]
self.weights = [self.iterations]
lr = self.learning_rate
for i, (p, g) in enumerate(zip(params, grads)):
g2 = K.square(g) + self.epsilon1
shape, dtype = K.int_shape(p), K.dtype(p)
factored_shape = self.factored_shape(shape)
if factored_shape is None:
# 定义参数
v = K.zeros(shape, dtype=dtype, name='v_' + str(i))
self.weights.append(v)
# 定义更新
v_t = self.beta2 * v + (1.0 - self.beta2) * g2
self.updates.append(K.update(v, v_t))
else:
# 定义参数
shape1, axis1, shape2, axis2 = factored_shape
vr = K.zeros(shape1, dtype=dtype, name='vr_' + str(i))
vc = K.zeros(shape2, dtype=dtype, name='vc_' + str(i))
self.weights.extend([vr, vc])
# 定义更新
vr_t = self.beta2 * vr + K.mean(g2, axis=axis1, keepdims=True)
vc_t = self.beta2 * vc + K.mean(g2, axis=axis2, keepdims=True)
self.updates.extend([K.update(vr, vr_t), K.update(vc, vc_t)])
# 合成矩阵
v_t = vr_t * vc_t / K.mean(vr_t, axis=axis2, keepdims=True)
# 增量主体
u = g / K.sqrt(v_t)
# 增量裁剪
if self.clipping_threshold is not None:
u_rms = K.mean(K.sum(K.square(u)))
d = self.clipping_threshold
u = u / K.maximum(1.0, u_rms / d)
# 增量滑动
if self.beta1 > 0.0:
# 定义参数
m = K.zeros(shape, dtype=dtype, name='m_' + str(i))
self.weights.append(m)
# 定义更新
m_t = self.beta1 * m + (1.0 - self.beta1) * u
self.updates.append(K.update(m, m_t))
u = m_t
# 增量调整
if self.multiply_by_parameter_scale:
u = u * K.maximum(K.mean(K.sum(K.square(p))), self.epsilon2)
# 更新参数
self.updates.append(K.update(p, p - lr * u))
return self.updates
def call(self, inputs):
"""如果custom_position_ids，那么第二个输入为自定义的位置id
"""
if self.custom_position_ids:
inputs, position_ids = inputs
if K.dtype(position_ids) != 'int32':
position_ids = K.cast(position_ids, 'int32')
pos_embeddings = K.gather(self.embeddings, position_ids)
else:
input_shape = K.shape(inputs)
batch_size, seq_len = input_shape[0], input_shape[1]
pos_embeddings = self.embeddings[:seq_len]
pos_embeddings = K.expand_dims(pos_embeddings, 0)
if self.merge_mode != 'add':
pos_embeddings = K.tile(pos_embeddings, [batch_size, 1, 1])
if self.merge_mode == 'add':
return inputs + pos_embeddings
else:
return K.concatenate([inputs, pos_embeddings])
def get_updates(self, loss, params):
updates = super(NewOptimizer, self).get_updates(loss, params)
k, alpha = self.steps_per_slow_update, self.slow_step_size
cond = K.equal(self.iterations % k, 0)
slow_vars = [
K.zeros(K.int_shape(p),
dtype=K.dtype(p),
name='slow_var_%s' % i) for i, p in enumerate(params)
]
with tf.control_dependencies(updates):
slow_updates = [
K.update(q, K.switch(cond, q + alpha * (p - q), q))
for p, q in zip(params, slow_vars)
]
with tf.control_dependencies(slow_updates):
copy_updates = [
K.update(p, K.switch(cond, q, p))
for p, q in zip(params, slow_vars)
]
return copy_updates
