tff.learning.models.functional_model_from_keras

Converts a tf.keras.Model to a tff.learning.models.FunctionalModel.

tff.learning.models.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]],
    metrics_constructor: Optional[Union[keras_utils.MetricConstructor, keras_utils.
        MetricsConstructor, keras_utils.MetricConstructors]] = None
) -> 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 supported.

This method doesn't support loss functions scaled by sample weights at the current state. Keras models with non-None sample weights will fail because sample weights aren't supported in model serialization and deserialization.

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.
metrics_constructor	An optional callable that must be compatible with tff.learning.metrics.create_functional_metric_fns.

Returns
A tff.learning.models.FunctionalModel.

Raises

KerasFunctionalModelError

If the following conditions: 1) the Keras model contains a batch normalization layer, 2) the Keras model is with non-trainable variable, 3) error occurs when converting the Keras model, 4) the Keras model shares variable across layers, 5) the FunctionalModel is used outside of a tff.tf_computation decorated callable or a graph context, 6) the Keras model contains a loss function with non-None sample weights.