Adds documentation on concrete functions for TFLite 2.0.

PiperOrigin-RevId: 238319646

tensorflow/lite/g3doc/r2/convert/concrete_function.md

@@ -0,0 +1,208 @@
+# Generating a concrete function
+
+In order to convert TensorFlow 2.0 models to TensorFlow Lite, the model needs to
+be exported as a concrete function. This document outlines what a concrete
+function is and how to generate one for an existing model.
+
+[TOC]
+
+## Background
+
+In TensorFlow 2.0, eager execution is on by default. TensorFlow's eager
+execution is an imperative programming environment that evaluates operations
+immediately, without building graphs. Operations return concrete values instead
+of constructing a computational graph to run later. A detailed guide on eager
+execution is available
+[here](https://github.com/tensorflow/docs/blob/master/site/en/r2/guide/eager.ipynb).
+
+While running imperatively with eager execution makes development and debugging
+more interactive, it doesn't allow for deploying on-device. The `tf.function`
+API makes it possible to save models as graphs, which is required to run
+TensorFlow Lite in 2.0. All operations wrapped in the `tf.function` decorator
+can be exported as a graph which can then be converted to the TensorFlow Lite
+FlatBuffer format.
+
+## Terminology
+
+The following terminology is used in this document:
+
+*   **Signature** - The inputs and outputs for a set of operations.
+*   **Concrete function** - Graph with a single signature.
+*   **Polymorphic function** - Python callable that encapsulates several
+    concrete function graphs behind one API.
+
+## Methodology
+
+This section describes how to export a concrete function.
+
+### Annotate functions with `tf.function`
+
+Annotating a function with `tf.function` generates a *polymorphic function*
+containing those operations. All operations that are not annotated with
+`tf.function` will be evaluated with eager execution. The examples below show
+how to use `tf.function`.
+
+```python
+@tf.function
+def pow(x):
+  return x ** 2
+```
+
+```python
+tf.function(lambda x : x ** 2)
+```
+
+### Create an object to save
+
+The `tf.function` can be optionally stored as part of a `tf.Module` object.
+Variables should only be defined once within the `tf.Module`. The examples below
+show two different approaches for creating a class that derives `Checkpoint`.
+
+```python
+class BasicModel(tf.Module):
+
+  def __init__(self):
+    self.const = None
+
+  @tf.function
+  def pow(self, x):
+    if self.const is None:
+      self.const = tf.Variable(2.)
+    return x ** self.const
+
+root = BasicModel()
+```
+
+```python
+root = tf.Module()
+root.const = tf.Variable(2.)
+root.pow = tf.function(lambda x : x ** root.const)
+```
+
+### Exporting the concrete function
+
+The concrete function defines a graph that can be converted to TensorFlow Lite
+model or be exported to a SavedModel. In order to export a concrete function
+from the polymorphic function, the signature needs to be defined. The signature
+can be defined the following ways:
+
+*   Define `input_signature` parameter in `tf.function`.
+*   Pass in `tf.TensorSpec` into `get_concrete_function`: e.g.
+    `tf.TensorSpec(shape=[1], dtype=tf.float32)`.
+*   Pass in a sample input tensor into `get_concrete_function`: e.g.
+    `tf.constant(1., shape=[1])`.
+
+The follow example shows how to define the `input_signature` parameter for
+`tf.function`.
+
+```python
+class BasicModel(tf.Module):
+
+  def __init__(self):
+    self.const = None
+
+  @tf.function(input_signature=[tf.TensorSpec(shape=[1], dtype=tf.float32)])
+  def pow(self, x):
+    if self.const is None:
+      self.const = tf.Variable(2.)
+    return x ** self.const
+
+# Create the tf.Module object.
+root = BasicModel()
+
+# Get the concrete function.
+concrete_func = root.pow.get_concrete_function()
+```
+
+The example below passes in a sample input tensor into `get_concrete_function`.
+
+```python
+# Create the tf.Module object.
+root = tf.Module()
+root.const = tf.Variable(2.)
+root.pow = tf.function(lambda x : x ** root.const)
+
+# Get the concrete function.
+input_data = tf.constant(1., shape=[1])
+concrete_func = root.pow.get_concrete_function(input_data)
+```
+
+## Example program
+
+```python
+import tensorflow as tf
+
+# Initialize the tf.Module object.
+root = tf.Module()
+
+# Instantiate the variable once.
+root.var = None
+
+# Define a function so that the operations aren't computed in advance.
+@tf.function
+def exported_function(x):
+  # Each variable can only be defined once. The variable can be defined within
+  # the function but needs to contain a reference outside of the function.
+  if root.var is None:
+    root.var = tf.Variable(tf.random.uniform([2, 2]))
+  root.const = tf.constant([[37.0, -23.0], [1.0, 4.0]])
+  root.mult = tf.matmul(root.const, root.var)
+  return root.mult * x
+
+# Save the function as part of the tf.Module object.
+root.func = exported_function
+
+# Get the concrete function.
+concrete_func = root.func.get_concrete_function(
+  tf.TensorSpec([1, 1], tf.float32))
+```
+
+## Common Questions
+
+### How do I save a concrete function as a SavedModel?
+
+Users who want to save their TensorFlow model before converting it to TensorFlow
+Lite should save it as a SavedModel. After getting the concrete function, call
+`tf.saved_model.save` to save the model. The example above can be saved using
+the following instruction.
+
+```python
+tf.saved_model.save(root, export_dir, concrete_func)
+```
+
+Reference the
+[SavedModel guide](https://github.com/tensorflow/docs/blob/master/site/en/r2/guide/saved_model.ipynb)
+for detailed instructions on using SavedModels.
+
+### How do I get a concrete function from the SavedModel?
+
+Each concrete function within a SavedModel can be identified by a signature key.
+The default signature key is `tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY`.
+The example below shows how to get the concrete function from a model.
+
+```python
+model = tf.saved_model.load(export_dir)
+concrete_func = model.signatures[
+  tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+```
+
+### How do I get a concrete function for a `tf.Keras` model?
+
+There are two approaches that you can use:
+
+1.  Save the model as a SavedModel. A concrete function will be generated during
+    the saving process, which can be accessed upon loading the model.
+2.  Annotate the model with `tf.function` as seen below.
+
+```python
+model = tf.keras.Sequential([tf.keras.layers.Dense(units=1, input_shape=[1])])
+model.compile(optimizer='sgd', loss='mean_squared_error')
+model.fit(x=[-1, 0, 1, 2, 3, 4], y=[-3, -1, 1, 3, 5, 7], epochs=50)
+
+# Get the concrete function from the Keras model.
+run_model = tf.function(lambda x : model(x))
+
+# Save the concrete function.
+concrete_func = run_model.get_concrete_function(
+    tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
+```

tensorflow/lite/g3doc/r2/convert/index.md

@@ -1,14 +1,14 @@
-# TensorFlow Lite Converter
+# TensorFlow Lite converter
 
-The TensorFlow Lite converter takes a TensorFlow model represented as a concrete
+The TensorFlow Lite converter takes a TensorFlow model represented as a
-function, and generates a TensorFlow Lite
+[concrete function](concrete_function.md), and generates a TensorFlow Lite
 [`FlatBuffer`](https://google.github.io/flatbuffers/) file (`.tflite`).
 
 Note: This page contains documentation on the converter API for TensorFlow 2.0.
 The API for TensorFlow 1.X is available
 [here](https://www.tensorflow.org/lite/convert/).
 
-## Device Deployment
+## Device deployment
 
 The TensorFlow Lite `FlatBuffer` file is then deployed to a client device (e.g.
 mobile, embedded) and run locally using the TensorFlow Lite interpreter. This
@@ -16,9 +16,9 @@ conversion process is shown in the diagram below:
 
 ![TFLite converter workflow](../images/convert/workflow.svg)
 
-## Converting Models
+## Converting models
 
-The TensorFlow Lite Converter can be used from the [Python API](python_api.md).
+The TensorFlow Lite converter can be used from the [Python API](python_api.md).
 Using the Python API makes it easier to convert models as part of a model
 development pipeline and helps mitigate
 [compatibility](../guide/ops_compatibility.md) issues early on.

tensorflow/lite/g3doc/r2/convert/python_api.md

@@ -1,7 +1,7 @@
 # Converter Python API guide
 
 This page provides examples on how to use the
-[TensorFlow Lite Converter](index.md) using the Python API in TensorFlow 2.0.
+[TensorFlow Lite converter](index.md) using the Python API in TensorFlow 2.0.
 
 [TOC]
 
@@ -10,6 +10,7 @@ This page provides examples on how to use the
 The Python API for converting TensorFlow models to TensorFlow Lite in TensorFlow
 2.0 is
 [`tf.lite.TFLiteConverter.from_concrete_function()`](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/lite/TFLiteConverter).
+Documentation on concrete functions is available [here](concrete_function.md).
 
 This document contains [example usages](#examples) of the API, a detailed list
 of [changes in the API between 1.X and 2.0](#differences), and
@@ -94,8 +95,8 @@ model.compile(optimizer='sgd', loss='mean_squared_error')
 model.fit(x, y, epochs=50)
 
 # Get the concrete function from the Keras model.
-to_save = tf.function(lambda x : model(x))
+run_model = tf.function(lambda x : model(x))
-concrete_func = to_save.get_concrete_function(
+concrete_func = run_model.get_concrete_function(
     tf.TensorSpec([None, 1], tf.float32))
 
 # Convert the model.