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