def get_updates(self, loss, params):
updates = super(NewOptimizer, self).get_updates(loss, params)
self.model_weights = params
self.ema_weights = [K.zeros(K.shape(w)) for w in params]
self.old_weights = K.batch_get_value(params)
K.batch_set_value(zip(self.ema_weights, self.old_weights))
ema_updates, ema_momentum = [], self.ema_momentum
with tf.control_dependencies(updates):
for w1, w2 in zip(self.ema_weights, params):
new_w = ema_momentum * w1 + (1 - ema_momentum) * w2
ema_updates.append(K.update(w1, new_w))
return ema_updates
def reset_old_weights(self):
"""恢复模型到旧权重。
"""
K.batch_set_value(zip(self.model_weights, self.old_weights))
def apply_ema_weights(self):
"""备份原模型权重，然后将平均权重应用到模型上去。
"""
self.old_weights = K.batch_get_value(self.model_weights)
ema_weights = K.batch_get_value(self.ema_weights)
K.batch_set_value(zip(self.model_weights, ema_weights))
