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))