def test_variable_capture(self):
with variable_utils.record_variable_creation_scope() as variable_list:
v1 = tf.Variable(1.0)
v2 = tf.Variable('abc', name='my_test_var')
v3 = tf.compat.v1.get_variable(
name='v1_var',
shape=(),
initializer=tf.compat.v1.initializers.constant)
# Explicitly add a variable that is not added to any collections.
v4 = tf.compat.v1.get_variable(
name='v1_var_no_collections',
shape=(),
initializer=tf.compat.v1.initializers.constant,
collections=[])
self.assertEqual([v1, v2, v3, v4], variable_list)
def test_construct_from_keras_converges(self):
functional_model = functional.functional_model_from_keras(
keras_model=create_test_keras_model(),
loss_fn=tf.keras.losses.MeanSquaredError(),
input_spec=(tf.TensorSpec([None, 1], dtype=tf.float32),
tf.TensorSpec([None, 1], dtype=tf.int32)))
with tf.Graph().as_default() as test_graph:
# Capture all the variables for later initialization in the session,
# otherwise it's hard to get our hands on the Keras-owned variables.
with variable_utils.record_variable_creation_scope(
) as captured_variables:
# Create data satisfying y = 2*x + 1
dataset = tf.data.Dataset.from_tensor_slices((
# Features
[[1.0], [2.0], [3.0]],
# Labels.
[[3.0], [5.0], [7.0]],
)).batch(1)
variables = tf.nest.map_structure(tf.Variable,
functional_model.initial_weights)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
@tf.function
def train():
weights = tf.nest.map_structure(lambda v: v.read_value(), variables)
initial_loss = loss = functional_model.forward_pass(
weights, next(iter(dataset)), training=True).loss
trainable = variables[0]
for batch in dataset.repeat(30):
with tf.GradientTape() as tape:
weights = tf.nest.map_structure(lambda v: v.read_value(),
variables)
tape.watch(weights[0])
batch_output = functional_model.forward_pass(
weights, batch, training=True)
gradients = tape.gradient(batch_output.loss, weights[0])
optimizer.apply_gradients(zip(gradients, trainable))
loss = batch_output.loss
return initial_loss, loss
initial_loss, final_loss = train()
with tf.compat.v1.Session(graph=test_graph) as sess:
sess.run(tf.compat.v1.initializers.variables(captured_variables))
initial_loss, final_loss = sess.run([initial_loss, final_loss])
# Expect some amount of convergence after a few epochs of the dataset.
self.assertGreater(initial_loss, 2.0)
self.assertLess(final_loss, 0.2)
def tf_computation_serializer(parameter_type: Optional[computation_types.Type],
context_stack):
"""Serializes a TF computation with a given parameter type.
Args:
parameter_type: The parameter type specification if the target accepts a
parameter, or `None` if the target doesn't declare any parameters. Either
an instance of `computation_types.Type`.
context_stack: The context stack to use.
Yields:
The first yielded value will be a Python object (such as a dataset,
a placeholder, or a `structure.Struct`) to be passed to the function to
serialize. The result of the function should then be passed to the
following `send` call.
The next yielded value will be
a tuple of (`pb.Computation`, `tff.Type`), where the computation contains
the instance with the `pb.TensorFlow` variant set, and the type is an
instance of `tff.Type`, potentially including Python container annotations,
for use by TensorFlow computation wrappers.
Raises:
TypeError: If the arguments are of the wrong types.
ValueError: If the signature of the target is not compatible with the given
parameter type.
"""
# TODO(b/113112108): Support a greater variety of target type signatures,
# with keyword args or multiple args corresponding to elements of a tuple.
# Document all accepted forms with examples in the API, and point to there
# from here.
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
if parameter_type is not None:
py_typecheck.check_type(parameter_type, computation_types.Type)
with tf.Graph().as_default() as graph:
if parameter_type is not None:
parameter_value, parameter_binding = tensorflow_utils.stamp_parameter_in_graph(
'arg', parameter_type, graph)
else:
parameter_value = None
parameter_binding = None
context = tensorflow_computation_context.TensorFlowComputationContext(
graph)
with context_stack.install(context):
with variable_utils.record_variable_creation_scope(
) as all_variables:
result = yield parameter_value
initializer_ops = []
if all_variables:
# Use a readable but not-too-long name for the init_op.
name = 'init_op_for_' + '_'.join(
[v.name.replace(':0', '') for v in all_variables])
if len(name) > 50:
name = 'init_op_for_{}_variables'.format(
len(all_variables))
initializer_ops.append(
tf.compat.v1.initializers.variables(all_variables,
name=name))
initializer_ops.extend(
tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.TABLE_INITIALIZERS))
if initializer_ops:
# Before running the main new init op, run any initializers for sub-
# computations from context.init_ops. Variables from import_graph_def
# will not make it into the global collections, and so will not be
# initialized without this code path.
with tf.compat.v1.control_dependencies(context.init_ops):
init_op_name = tf.group(*initializer_ops,
name='grouped_initializers').name
elif context.init_ops:
init_op_name = tf.group(*context.init_ops,
name='subcomputation_init_ops').name
else:
init_op_name = None
result_type, result_binding = tensorflow_utils.capture_result_from_graph(
result, graph)
type_signature = computation_types.FunctionType(parameter_type,
result_type)
tensorflow = pb.TensorFlow(graph_def=serialization_utils.pack_graph_def(
graph.as_graph_def()),
parameter=parameter_binding,
result=result_binding,
initialize_op=init_op_name)
yield pb.Computation(
type=type_serialization.serialize_type(type_signature),
tensorflow=tensorflow), type_signature
def functional_model_from_keras(
keras_model: Union[tf.keras.Model, Callable[[], tf.keras.Model]],
loss_fn: tf.keras.losses.Loss,
input_spec: Union[Sequence[Any], Mapping[str, Any]],
) -> FunctionalModel:
"""Converts a `tf.keras.Model` to a `tff.learning.models.FunctionalModel`.
NOTE: This method only supports models where calling that model with
`training=True` and `training=False` produce the same graph. Keras layers
such as batch normalization will fail because they require updating internal
state when `training=True` which is not suported.
IMPORTANT: The returned model must only be used in a graph context (for
example inside a `tff.tf_computation` decorated callable). It will raise an
error otherwise.
Args:
keras_model: A `tf.keras.Model` object, should be uncompiled. If compiled,
the metrics, optimizer, and loss function will be ignored. Note: models
that have multiple outputs will send all outputs to the `loss_fn`.
loss_fn: A `tf.keras.losses.Loss` object.
input_spec: A structure of `tf.TensorSpec` defining the input to the model.
Returns:
A `tff.learning.models.FunctionalModel`.
Raises:
KerasFunctionalModelError: the model has a batch normalization layer.
"""
# We're going to do something fancy here:
#
# 1. Get a copy of all the variables, in the order they are created during
# model construction, when in a graph context.
# 2. Use this ordering to construct a type signature of the model weights in
# such a way that we can inject TENSORS (those that are coming in as
# arguments) in place of variable creation during a call to
# `tf.keras.models.clone_model()`, which gives us a newly constructed Keras
# model in the context we want.
# 3. Profit by having variableless graphs!
#
# **WARNING** Caveats:
#
# 1. This model _must_ be used inside a graph context (e.g. a
# `tff.tf_computation` decorated callable, aka a `tff.Computation`). Keras
# appears to create extra variables in the eager context that are not part
# of the user specified model, and end up not being compatible.
#
# 2. We have found that this trick does NOT work with non-trainable variables
# that are updated during training. Namely layers such as
# BatchNormalization try to update means/variances during training and are
# not compatible with this approach. We generally recommend
# GroupNormalization in place of BatchNormalization at the current time.
#
# 3. This does not support multiple outputs with different loss functions, or
# laywerise regularization losses TODO(b/156629927).
if isinstance(keras_model, tf.keras.Model):
for layer in keras_model.layers:
# There may be other layers that are problematic, at this time updating
# the mean/variance in batchnorm layer is the only known such instance.
if isinstance(layer, tf.keras.layers.BatchNormalization):
raise KerasFunctionalModelError(
'Keras model contains a batch normalization layer, which is '
'incompatible with `tff.learning.models.FunctionalModel`. Consider '
'using group normalization instead.')
if keras_model.non_trainable_variables:
raise KerasFunctionalModelError(
'Received a Keras model with non-trainable variables. Keras models with '
'non-trainable variables are currently not supported by FunctionalModel'
'. Most training algorithms (e.g. Federated Averaging) will not '
'aggregate them, and they are not updated locally by the optimizer. '
'We can relax this in the future if we have APIs that support updating '
'non-trainable variables.')
elif not callable(keras_model):
raise ValueError(
'`keras_model` must be a `tf.keras.Model` or a no-arg '
'callable that returns a `tf.keras.Model`.')
# Clone the keras model inside a graph context so that we only get the
# variables for the layers (otherwise keras adds other non-user variables). We
# also setup ops to inject the current model weights, because the cloned model
# will be re-initialized from scratch.
with tf.Graph().as_default() as g:
with variable_utils.record_variable_creation_scope(
) as captured_variables:
if isinstance(keras_model, tf.keras.Model):
try:
cloned_model = tf.keras.models.clone_model(keras_model)
except RuntimeError as e:
raise KerasFunctionalModelError(
'Encountered a error converting the Keras model. Often this '
'occurs when the `tf.keras.Model` has a layer that receives '
'inputs from other layers directly (e.g. shared embeddings).'
'To avoid the problem, wrap the `tf.keras.Model` construction in '
'a no-arg callable (e.g. lambda) and pass that callable to '
'`functional_model_from_keras`') from e
if len(cloned_model.variables) != len(keras_model.variables):
raise KerasFunctionalModelError(
'The input Keras model is likely sharing variables across layers '
'which is unsupported. Cloning the model will duplicate these '
'variables and result in unexpected training gradients.'
)
else:
cloned_model = keras_model()
# Ensure our cloned model has the same weights as the current model.
# We'll feed in the current model waits into the palceholders for
# assignmnet in a session below.
def assign_placeholder(v):
p = tf.compat.v1.placeholder(dtype=v.dtype)
return v.assign(p), p
assign_ops, placeholders = zip(*(assign_placeholder(v)
for v in cloned_model.variables))
trainable_variables = tuple(v for v in captured_variables if v.trainable)
non_trainable_variables = tuple(v for v in captured_variables
if not v.trainable)
# Here we get the initial weights from the incoming keras model in the order
# they are constructed; and also ensure that the values are set to the
# incoming model weights rather than their fresh initialization.
if isinstance(keras_model, tf.keras.Model):
model_for_variables = keras_model
else:
model_for_variables = keras_model()
current_model_weights = tf.nest.map_structure(
lambda v: v.read_value().numpy(), model_for_variables.variables)
with tf.compat.v1.Session(graph=g) as sess:
sess.run(tf.compat.v1.initializers.variables(captured_variables))
sess.run(fetches=assign_ops,
feed_dict=dict(zip(placeholders, current_model_weights)))
initial_weights = sess.run(fetches=(trainable_variables,
non_trainable_variables))
@tf.function
def predict_on_batch(model_weights: ModelWeights,
x: Any,
training: bool = True) -> Any:
with tf.init_scope():
if tf.executing_eagerly():
raise KerasFunctionalModelError(
'tf.keras.Model used as a FunctionalModel is only usable inside a '
'tff.tf_computation decorated callable or a graph context.'
)
# Make a copy of the weights container; can't mutate Python containers
# inside a tf.function.
trainable, non_trainable = (list(w) for w in model_weights)
# Here were intercept variable creation requests during the
# `tf.keras.models.clone_model()` call.
#
# Instead of forwarding the variable request to TF core and getting a
# `tf.Variable` back, we skip that and return only the `tf.Tensor` that
# corresponds to the `tf.Variable` recreation request (avoiding any variable
# creation). This works because TF operations that accept `tf.Variable`
# inputs automatically call `variable.read_value()` and then operate on that
# resulting tensor. We're relying on shortcutting that and providing the
# tensor straight away.
#
# For example, `tf.matmul` doesn't notice its input is `tf.Variable` or
# `tf.Tensor`:
#
# v = tf.Variable([[1], [2], [3]])
# tf.matmul(v, [[4, 5, 6]])
#
# and
#
# v = tf.constant([[1], [2], [3]])
# tf.matmul(v, [[4, 5, 6]])
#
# both result in:
#
# <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
# array([[ 4, 5, 6],
# [ 8, 10, 12],
# [12, 15, 18]], dtype=int32)>
def swap_tensor_parameter_for_variable(_, **kwargs):
if kwargs.get('trainable', True):
return trainable.pop(0)
else:
return non_trainable.pop(0)
with tf.variable_creator_scope(swap_tensor_parameter_for_variable):
if isinstance(keras_model, tf.keras.Model):
variableless_model = tf.keras.models.clone_model(keras_model)
else:
variableless_model = keras_model()
return variableless_model(x, training)
@tf.function
def forward_pass(model_weights: ModelWeights,
batch_input: Any,
training: bool = True) -> model_lib.BatchOutput:
if isinstance(batch_input, collections.abc.Mapping):
x = batch_input['x']
y = batch_input['y']
elif isinstance(batch_input, collections.abc.Sequence):
x, y = batch_input
else:
raise ValueError(
'`batch_input` must be either a mapping with keys `x` '
f'and `y` or a sequence of `(x, y)`. Got: {batch_input!r}')
predictions = predict_on_batch(model_weights, x, training)
batch_loss = loss_fn(y_true=y, y_pred=predictions)
# TODO(b/207033265): more work needed to support models with multiple loss
# functions.
def nrows(t):
return t.nrows() if isinstance(t,
tf.RaggedTensor) else tf.shape(t)[0]
return model_lib.BatchOutput(loss=batch_loss,
predictions=predictions,
num_examples=nrows(
tf.nest.flatten(batch_input)[0]))
return FunctionalModel(initial_weights=initial_weights,
forward_pass_fn=forward_pass,
predict_on_batch_fn=predict_on_batch,
input_spec=input_spec)
def serialize_py_fn_as_tf_computation(target, parameter_type, context_stack):
"""Serializes the 'target' as a TF computation with a given parameter type.
See also `serialize_tf2_as_tf_computation` for TensorFlow 2
serialization.
Args:
target: The entity to convert into and serialize as a TF computation. This
can currently only be a Python function. In the future, we will add here
support for serializing the various kinds of non-eager and eager
functions, and eventually aim at full support for and compliance with TF
2.0. This function is currently required to declare either zero parameters
if `parameter_type` is `None`, or exactly one parameter if it's not
`None`. The nested structure of this parameter must correspond to the
structure of the 'parameter_type'. In the future, we may support targets
with multiple args/keyword args (to be documented in the API and
referenced from here).
parameter_type: The parameter type specification if the target accepts a
parameter, or `None` if the target doesn't declare any parameters. Either
an instance of `types.Type`, or something that's convertible to it by
`types.to_type()`.
context_stack: The context stack to use.
Returns:
A tuple of (`pb.Computation`, `tff.Type`), where the computation contains
the instance with the `pb.TensorFlow` variant set, and the type is an
instance of `tff.Type`, potentially including Python container annotations,
for use by TensorFlow computation wrappers.
Raises:
TypeError: If the arguments are of the wrong types.
ValueError: If the signature of the target is not compatible with the given
parameter type.
"""
# TODO(b/113112108): Support a greater variety of target type signatures,
# with keyword args or multiple args corresponding to elements of a tuple.
# Document all accepted forms with examples in the API, and point to there
# from here.
py_typecheck.check_type(target, types.FunctionType)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
parameter_type = computation_types.to_type(parameter_type)
signature = function_utils.get_signature(target)
with tf.Graph().as_default() as graph:
if parameter_type is not None:
if len(signature.parameters) != 1:
raise ValueError(
'Expected the target to declare exactly one parameter, found {!r}.'
.format(signature.parameters))
parameter_name = next(iter(signature.parameters))
parameter_value, parameter_binding = tensorflow_utils.stamp_parameter_in_graph(
parameter_name, parameter_type, graph)
else:
if signature.parameters:
raise ValueError(
'Expected the target to declare no parameters, found {!r}.'
.format(signature.parameters))
parameter_value = None
parameter_binding = None
context = tensorflow_computation_context.TensorFlowComputationContext(
graph)
with context_stack.install(context):
with variable_utils.record_variable_creation_scope(
) as all_variables:
if parameter_value is not None:
result = target(parameter_value)
else:
result = target()
initializer_ops = []
if all_variables:
# Use a readable but not-too-long name for the init_op.
name = 'init_op_for_' + '_'.join(
[v.name.replace(':0', '') for v in all_variables])
if len(name) > 50:
name = 'init_op_for_{}_variables'.format(
len(all_variables))
initializer_ops.append(
tf.compat.v1.initializers.variables(all_variables,
name=name))
initializer_ops.extend(
tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.TABLE_INITIALIZERS))
if initializer_ops:
# Before running the main new init op, run any initializers for sub-
# computations from context.init_ops. Variables from import_graph_def
# will not make it into the global collections, and so will not be
# initialized without this code path.
with tf.compat.v1.control_dependencies(context.init_ops):
init_op_name = tf.group(*initializer_ops,
name='grouped_initializers').name
elif context.init_ops:
init_op_name = tf.group(*context.init_ops,
name='subcomputation_init_ops').name
else:
init_op_name = None
result_type, result_binding = tensorflow_utils.capture_result_from_graph(
result, graph)
type_signature = computation_types.FunctionType(parameter_type,
result_type)
# WARNING: we do not really want to be modifying the graph here if we can
# avoid it. This is purely to work around performance issues uncovered with
# the non-standard usage of Tensorflow and have been discussed with the
# Tensorflow core team before being added.
clean_graph_def = _clean_graph_def(graph.as_graph_def())
tensorflow = pb.TensorFlow(
graph_def=serialization_utils.pack_graph_def(clean_graph_def),
parameter=parameter_binding,
result=result_binding,
initialize_op=init_op_name)
return pb.Computation(
type=type_serialization.serialize_type(type_signature),
tensorflow=tensorflow), type_signature