def transfer_weights( model: Model, new_model: Model, optimizer: tf.keras.optimizers.Optimizer, new_optimizer: tf.keras.optimizers.Optimizer, ignore_weights: list = None, ) -> None: if type(model) is not type(new_model): raise ValueError( "Transferring weights to another model type is not supported") if ignore_weights is None: ignore_weights = list() ignore_weights_ref = set(weight.experimental_ref() for weight in ignore_weights) weights = model.weights new_weights = new_model.weights for weight, new_weight in zip(weights, new_weights): if new_weight.experimental_ref() not in ignore_weights_ref: new_weight.assign(weight) for slot_name in new_optimizer.get_slot_names(): if slot_name not in optimizer.get_slot_names(): continue new_slot = new_optimizer.get_slot(new_weight, slot_name) slot = optimizer.get_slot(weight, slot_name) new_slot.assign(slot)
def _get_optimizer_preconditioner(optimizer: tf.keras.optimizers.Optimizer, model_weights: tff.learning.ModelWeights): """Get the preconditioner states from the optimizer.""" config = optimizer.get_config() config_name = config['name'] if config_name == 'Adagrad': eps = optimizer.epsilon v = tf.nest.map_structure( lambda var: optimizer.get_slot(var, 'accumulator'), model_weights.trainable) return tf.nest.map_structure( lambda a: tf.math.divide_no_nan( # pylint: disable=g-long-lambda 1.0, tf.math.sqrt(a) + eps), v) elif config_name in {'Adam', 'Yogi'}: eps = optimizer.epsilon v = tf.nest.map_structure(lambda var: optimizer.get_slot(var, 'v'), model_weights.trainable) return tf.nest.map_structure( lambda a: tf.math.divide_no_nan( # pylint: disable=g-long-lambda 1.0, tf.math.sqrt(a) + eps), v) elif config_name == 'SGD': return tf.nest.map_structure(tf.ones_like, model_weights.trainable) else: raise TypeError( 'client optimizer should be one of these: SGD, Adagrad, Adam, Yogi.' )
def _get_optimizer_preconditioner(optimizer: tf.keras.optimizers.Optimizer, model_weights: tff.learning.ModelWeights): """Get the preconditioner states from the optimizer.""" config = optimizer.get_config() config_name = config['name'] if config_name == 'Adagrad': return tf.nest.map_structure( lambda var: optimizer.get_slot(var, 'accumulator'), model_weights.trainable) elif config_name in {'Adam', 'Yogi'}: return tf.nest.map_structure(lambda var: optimizer.get_slot(var, 'v'), model_weights.trainable) elif config_name == 'SGD': return tf.nest.map_structure(tf.ones_like, model_weights.trainable) else: raise TypeError( 'client optimizer should be one of these: SGD, Adagrad, Adam, Yogi.')
def update_variable_and_slots( update: VariableUpdate, mapping: List[int], ref_optimizer: tf.keras.optimizers.Optimizer, new_optimizer: tf.keras.optimizers.Optimizer, ) -> List[tf.Variable]: """Update a vocabulary variable and its associated optimizer slots (if any).""" variables = [update_variable(update, mapping)] ref_slot_names = ref_optimizer.get_slot_names() new_slot_names = new_optimizer.get_slot_names() for slot_name in ref_slot_names: if slot_name not in new_slot_names: continue ref_slot = ref_optimizer.get_slot(update.ref_variable, slot_name) new_slot = new_optimizer.get_slot(update.new_variable, slot_name) slot_update = VariableUpdate(ref_slot, new_slot, update.vocab_axis) variables.append(update_variable(slot_update, mapping)) return variables