Repository: OpenNMT/OpenNMT-tf
Pull request #789 (2021-02-18 22:34:22)

TensorFlow Lite Support
**Introduction**
This pull request adds support to convert 4 models to TensorFlow Lite.  This addresses the feature request mentioned in issue #490.

**Supported Models**
The following models are supported for TensorFlow Lite conversion: 
_NMTSmallV1, NMTMediumV1, NMTBigV1, Luong Attention_


If the model is not yet supported for conversion to TensorFlow Lite, then an error is thrown.

Note: I have not looked at converting other models, I suspect models that inherit from SequenceToSequence may not be too difficult, though this remains to be investigated.

**Running**
I added a _save_tflite_ parameter in addition to the parameters {_train, eval, infer, score_}.
This parameter is used when you want to convert the model to TensorFlow Lite.

Example input:
`python opennmt/bin/main.py --model_type NMTSmallV1 --config opensubs/data.yml --checkpoint_path opensubs/smallnmt/checkpoint save_tflite --output_dir ./opensubs/smallnmt/tflite`

The conversion of a TensorFlow-Lite should take about 1 to 2 minutes.

**Android Model Input and Output**
To run on android, by default the model will require an integer array of any size to run an inference on Android. 
The output of the model will by default be an integer array of size 250.  The output size is modifiable when you convert the model to TFLite in the data configuration file.
```
params:
	tflite_output_size: 100
```

I've tested the models on Android, and they produce the same results as Python.

Here is some of the code I used for running on Android.

**Loading the model**
```
	AssetManager assetManager = this.context.getResources().getAssets();
        ByteBuffer buffer = loadModelFile(assetManager, modelPath);
        if (buffer == null){
            Log.e("nmt-tf", "Could not load model");
            return;
        }
        Interpreter.Options opt = new Interpreter.Options();
        opt.setNumThreads(NUM_LITE_THREADS);
        tflite = new Interpreter(buffer, opt);
```

**Running an inference**
```
	int[] input_ids = acsentenceToIds(“Hello World”);
        int[] output_ids = new int[250];
        tflite.run(input_ids, output_ids);
	String translatedSentence = IdsToSentence(output_ids);
	System.out.println(“Translated Sentence: ” + translatedSentence);
```

**Vocabulary Files**
On Android, you will have to write a function to read the vocabulary files and to convert sentences you want to feed onto the model as integer arrays.  To get a sentence back from the model, the integer array will have to be converted to a sentence with the vocabulary file. 

Android libraries require the nightly version of TensorFlow, I have not tested with 2.4 TensorFlow Android libraries.

```
implementation 'org.tensorflow:tensorflow-lite:0.0.0-nightly'
implementation 'org.tensorflow:tensorflow-lite-select-tf-ops:0.0.0-nightly'
implementation 'org.tensorflow:tensorflow-lite-support:0.0.0-nightly'
```

**Testing**
I have included a test file.  Below I describe 4 tests that are implemented in this file opennmt/tests/tflite_test.py.


<table>
    <thead>
        <tr>
            <th>Test name</th>
            <th>Model</th>
            <th>Test Description</th>
        </tr>
    </thead>
    <tbody>
        <tr>
            <td >testLuongAttentionTFLiteOutput </td>
<td >Luong Attention</td>
            <td rowspan=2>Checks that the model prediction is the same when outputting with the TensorFlow Lite conversion.</td>
        </tr>
        <tr>
<td >testNMTBigV1TFLiteOutput </td>
<td >NMTBigV1</td>
        </tr>
        <tr>
            <td >testLuongAttentionTFLiteFile </td>
<td >Luong Attention</td>
            <td rowspan=2>Checks that the code is   able to convert the model and output the model into a file.</td>
        </tr>
        <tr>
            <td >testNMTBigTFLiteFile </td>
<td >NMTBigV1</td>
        </tr>
    </tbody>
</table>

**Code Changes**
The TensorFlow Lite model was converted by running the _TFLiteConverter.from_concrete_functions_ function.

I created an inference function that takes as a parameter an integer array, and outputs an integer array (ids).

The custom inference function runs very similar to the regular inference function, with some differences:

- Decoding is stopped early after it outputs the ids, and just outputs the ids as an integer array of ids
- Argument no longer takes in a dict object with Tensors and now just requires an Tensor Integer array
- Word embedding lookup is modified when running with TFLite due to a bug with converting to TensorFlow Lite - This is mentioned in issue #42410 in the TensorFlow Repository, and the fix mentioned there is applied.
- Decoding function has an extra parameter to run differently with TFLite
- Decoding change when running with TFLite in the file opennmt/utils/decoding.py file, if running with TFLite, then a batch size of 1 is assumed and TensorArrays are initiated to be static.

**Source Files Modified**
File | Reason
-- | --
opennmt/bin/main.py | Added ‘save_tflite’   argument to pass in when converting a model to TFLite (mobile friendly).
opennmt/runner.py | Added save_tflite function.  This function is called when the save_tflite   argument is given.  The model throws an   error if the model is not compatible with conversion to TensorFlow Lite.  Elsewise, the model is loaded and converted.
opennmt/inputters/text_inputter.py | Adds the function tflite_call   which gets the word embeddings in a way that is safe for TensorFlow   Lite.  The fix and problem description   can be found in TensorFlow Repository issue  #42410.
opennmt/models/sequence_to_sequence.py | Creates the custom inference   function which is the concrete function that will be converted.  A parameter is added to _dynamic_decode   to tell the function that it is running in a TensorFlow Lite safe way, and if   running with TensorFlow Lite, then returns ids.
opennmt/decoders/decoder.py | Modifies dynamic_decode   function to use the modified word embedding function.
opennmt/utils/decoding.py | Modifies the decoding.  The original function uses TensorArrays   which are not compatible with TensorFlow Lite unless they are specified as   static and are pre-allocated; this is changed.

**Test Files Modified/Added**
File | Reason
-- | --
opennmt/models/__init__.py | Made NMTBigV1 public, so   we can test this model.
opennmt/tests/tflite_test.py | Contains four test cases created to   check if the models NMTBigV1 and Luong Attention produce the same output when   running as a TensorFlow Lite model and to check if they were able to convert   the model and save it.





**Acknowledgement**
This project for converting models to TensorFlow Lite was created under contract for the U.S. DEVCOM Army Research Laboratory

gcervantes8: I've done all the changes you recommended! Let me know if you believe I should make any more changes.

Thank you for reviewing.

gcervantes8: Thanks for reviewing the pull request again!  I have made the suggested changes.  If you see anywhere else that needs changes, let me know.

Patch:
diff --git a/docs/configuration.md b/docs/configuration.md
index 6cce3d13..f37b9b1f 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -231,6 +231,8 @@ params:
   contrastive_learning: false
   # (optional) The value of the parameter eta in the max-margin loss (default: 0.1).
   max_margin_eta: 0.1
+  # (optional) Size of input when saving a model as TFLite
+  tflite_output_size: 250
 
 
 # Training options.
diff --git a/opennmt/decoders/decoder.py b/opennmt/decoders/decoder.py
index cf566d0a..1506ebb5 100644
--- a/opennmt/decoders/decoder.py
+++ b/opennmt/decoders/decoder.py
@@ -401,6 +401,7 @@ class Decoder(tf.keras.layers.Layer):
         sampler=None,
         maximum_iterations=None,
         minimum_iterations=0,
+        tflite_output_size=None,
     ):
         """Decodes dynamically from :obj:`start_ids`.
 
@@ -416,6 +417,7 @@ class Decoder(tf.keras.layers.Layer):
             predictions from the model output. Defaults to an argmax sampling.
           maximum_iterations: The maximum number of iterations to decode for.
           minimum_iterations: The minimum number of iterations to decode for.
+          tflite_output_size: If not None will run TFLite safe, is the size of 1D output tensor.
 
         Returns:
           A :class:`opennmt.utils.DecodingResult` instance.
@@ -423,7 +425,9 @@ class Decoder(tf.keras.layers.Layer):
         See Also:
           :func:`opennmt.utils.dynamic_decode`
         """
-        if isinstance(embeddings, text_inputter.WordEmbedder):
+        if tflite_output_size is not None:
+            input_fn = lambda ids: embeddings.tflite_call(ids)
+        elif isinstance(embeddings, text_inputter.WordEmbedder):
             input_fn = lambda ids: embeddings({"ids": ids})
         else:
             input_fn = lambda ids: tf.nn.embedding_lookup(embeddings, ids)
@@ -446,6 +450,7 @@ class Decoder(tf.keras.layers.Layer):
             attention_size=tf.shape(self.memory)[1]
             if self.support_alignment_history
             else None,
+            tflite_output_size=tflite_output_size,
         )
 
     def map_v1_weights(self, weights):
diff --git a/opennmt/inputters/text_inputter.py b/opennmt/inputters/text_inputter.py
index a3392a48..1b724eba 100644
--- a/opennmt/inputters/text_inputter.py
+++ b/opennmt/inputters/text_inputter.py
@@ -442,6 +442,16 @@ class WordEmbedder(TextInputter):
         outputs = common.dropout(outputs, self.dropout, training=training)
         return outputs
 
+    def tflite_call(self, ids):
+        """Replicates call function and changes parameter.
+        Avoids a TFLite conversion issue, TensorFlow issue #42410
+
+        Args:
+          ids: A 1-dimensional tensor with the ids that you want to get word embeddings of.
+        """
+        ids = tf.expand_dims(ids, axis=-1)
+        return tf.gather_nd(self.embedding, ids)
+
     def visualize(self, model_root, log_dir):
         save_embeddings_metadata(
             log_dir,
diff --git a/opennmt/models/sequence_to_sequence.py b/opennmt/models/sequence_to_sequence.py
index 16d40c3b..eada0dd1 100644
--- a/opennmt/models/sequence_to_sequence.py
+++ b/opennmt/models/sequence_to_sequence.py
@@ -191,6 +191,32 @@ class SequenceToSequence(model.SequenceGenerator):
 
         return outputs, predictions
 
+    def infer_tflite(self, ids):
+        """Runs a prediction, returns a 1-dim Tensor with the target ids it predicted, is TFLite safe.
+        This is the function that gets converted was saving as a TFLite model
+
+        Args:
+          ids: A 1-dimensional tensor with the ids of the sentence you want to predict
+        """
+        ids = tf.expand_dims(ids, axis=0)
+        source_inputs = self.features_inputter.tflite_call(ids)
+        source_length = tf.convert_to_tensor([tf.math.count_nonzero(ids)])
+
+        encoder_outputs, encoder_state, encoder_sequence_length = self.encoder(
+            source_inputs, sequence_length=source_length, training=False
+        )
+
+        predictions = self._dynamic_decode(
+            ids,
+            encoder_outputs,
+            encoder_state,
+            encoder_sequence_length,
+            tflite_run=True,
+        )
+
+        predictions = tf.squeeze(predictions, axis=1)
+        return predictions
+
     def _decode_target(
         self,
         labels,
@@ -245,7 +271,12 @@ class SequenceToSequence(model.SequenceGenerator):
         return outputs
 
     def _dynamic_decode(
-        self, features, encoder_outputs, encoder_state, encoder_sequence_length
+        self,
+        features,
+        encoder_outputs,
+        encoder_state,
+        encoder_sequence_length,
+        tflite_run=False,
     ):
         params = self.params
         batch_size = tf.shape(tf.nest.flatten(encoder_outputs)[0])[0]
@@ -285,7 +316,14 @@ class SequenceToSequence(model.SequenceGenerator):
             sampler=decoding.Sampler.from_params(params),
             maximum_iterations=params.get("maximum_decoding_length", 250),
             minimum_iterations=params.get("minimum_decoding_length", 0),
+            tflite_output_size=params.get("tflite_output_size", 250)
+            if tflite_run
+            else None,
         )
+
+        if tflite_run:
+            return sampled_ids
+
         target_tokens = self.labels_inputter.ids_to_tokens.lookup(
             tf.cast(sampled_ids, tf.int64)
         )
diff --git a/opennmt/utils/decoding.py b/opennmt/utils/decoding.py
index 120fc1be..232b2811 100644
--- a/opennmt/utils/decoding.py
+++ b/opennmt/utils/decoding.py
@@ -427,6 +427,7 @@ def dynamic_decode(
     minimum_iterations=0,
     attention_history=False,
     attention_size=None,
+    tflite_output_size=None,
 ):
     """Dynamic decoding.
 
@@ -445,6 +446,7 @@ def dynamic_decode(
       attention_history: Gather attention history during the decoding.
       attention_size: If known, the size of the attention vectors (i.e. the
         maximum source length).
+      tflite_output_size: If not None will run TFLite safe, is the size of 1D output tensor.
 
     Returns:
       A :class:`opennmt.utils.DecodingResult` instance.
@@ -455,6 +457,7 @@ def dynamic_decode(
         decoding_strategy = GreedySearch()
     if sampler is None:
         sampler = BestSampler()
+    is_tflite_run = tflite_output_size is not None
 
     def _cond(
         step, finished, state, inputs, outputs, attention, cum_log_probs, extra_vars
@@ -507,14 +510,18 @@ def dynamic_decode(
             **extra_vars,
         )
 
-        # Update loop vars.
-        if attention_history:
-            if attn is None:
-                raise ValueError(
-                    "attention_history is set but the model did not return attention"
-                )
-            attention = attention.write(step, tf.cast(attn, tf.float32))
-        outputs = outputs.write(step, output)
+        # For TensorArray assuming batch size of 1 to be TFlite safe
+        if is_tflite_run:
+            outputs = outputs.write(step, output[0])
+        else:
+            # Update loop vars.
+            if attention_history:
+                if attn is None:
+                    raise ValueError(
+                        "attention_history is set but the model did not return attention"
+                    )
+                attention = attention.write(step, tf.cast(attn, tf.float32))
+            outputs = outputs.write(step, output)
         cum_log_probs = tf.where(finished, x=cum_log_probs, y=next_cum_log_probs)
         finished = tf.logical_or(finished, tf.equal(output, end_id))
         return (
@@ -535,8 +542,32 @@ def dynamic_decode(
         start_ids, attention_size=attention_size
     )
     step = tf.constant(0, dtype=tf.int32)
-    outputs = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
-    attention = tf.TensorArray(tf.float32, size=0, dynamic_size=True)
+
+    if is_tflite_run:
+        output_shape = tf.TensorShape(())
+        outputs = tf.TensorArray(
+            tf.int32,
+            size=tflite_output_size,
+            dynamic_size=False,
+            element_shape=output_shape,
+        )
+        attn_shape = tf.TensorShape(())
+        attention = tf.TensorArray(
+            tf.float32,
+            size=tflite_output_size,
+            dynamic_size=False,
+            element_shape=attn_shape,
+        )
+        maximum_iterations = (
+            tflite_output_size
+            if maximum_iterations > tflite_output_size
+            else maximum_iterations
+        )
+    else:
+        output_shape = tf.TensorShape(None)
+        outputs = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
+        attn_shape = tf.TensorShape(None)
+        attention = tf.TensorArray(tf.float32, size=0, dynamic_size=True)
 
     _, _, state, _, outputs, attention, log_probs, extra_vars = tf.while_loop(
         _cond,
@@ -556,8 +587,8 @@ def dynamic_decode(
             finished.shape,
             tf.nest.map_structure(_get_shape_invariants, initial_state),
             start_ids.shape,
-            tf.TensorShape(None),
-            tf.TensorShape(None),
+            output_shape,
+            attn_shape,
             initial_log_probs.shape,
             tf.nest.map_structure(_get_shape_invariants, extra_vars),
         ),
@@ -568,10 +599,10 @@ def dynamic_decode(
     ids, attention, lengths = decoding_strategy.finalize(
         outputs,
         end_id,
-        attention=attention if attention_history else None,
+        attention=attention if attention_history and not is_tflite_run else None,
         **extra_vars,
     )
-    if attention is not None:
+    if attention is not None and not is_tflite_run:
         attention = attention[:, :, :-1]  # Ignore attention for 
.
     log_probs = tf.reshape(log_probs, [-1, decoding_strategy.num_hypotheses])
     ids = tf.cast(ids, ids_dtype)
diff --git a/opennmt/utils/exporters.py b/opennmt/utils/exporters.py
index f12bc5d7..c96551e3 100644
--- a/opennmt/utils/exporters.py
+++ b/opennmt/utils/exporters.py
@@ -7,6 +7,7 @@ import tempfile
 import tensorflow as tf
 
 from opennmt.utils import misc
+from opennmt.models import catalog
 
 
 class Exporter(abc.ABC):
@@ -80,6 +81,76 @@ class CheckpointExporter(Exporter):
         checkpoint.write(os.path.join(export_dir, "ckpt"))
 
 
+@register_exporter(name="tflite")
+class TFLiteExporter(Exporter):
+    """TensorFlow Lite exporter."""
+
+    def __init__(self, quantization=None):
+        self._quantization = quantization
+
+    def _export_model(self, model, export_dir):
+
+        # Models currently supported for TFLite exporting
+        tflite_supported_models = [
+            catalog.NMTSmallV1,
+            catalog.NMTMediumV1,
+            catalog.NMTBigV1,
+            catalog.LuongAttention,
+        ]
+
+        # If it isn't any of the supported models, raise an exception
+        if not any(
+            [
+                isinstance(model, supported_model)
+                for supported_model in tflite_supported_models
+            ]
+        ):
+            raise TypeError(
+                "Unsupported model to export to TFLite, supported models are:"
+                "NMTSmallV1, NMTMediumV1, NMTBigV1, LuongAttention"
+            )
+
+        if not model.built:
+            model.create_variables()
+
+        # Tries to run prediction with TensorFlow Lite method it will convert
+        tflite_concrete_fn = tf.function(
+            model.infer_tflite,
+            input_signature=[tf.TensorSpec([None], dtype=tf.dtypes.int32, name="ids")],
+        ).get_concrete_function()
+
+        # Saving
+        converter = tf.lite.TFLiteConverter.from_concrete_functions(
+            [tflite_concrete_fn]
+        )
+        converter.target_spec.supported_ops = [
+            tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.
+            tf.lite.OpsSet.SELECT_TF_OPS,  # enable TensorFlow ops.
+        ]
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+
+        if self._quantization is not None:
+            if self._quantization in "float16":
+                converter.target_spec.supported_types = [tf.float16]
+
+        tflite_model_path = os.path.join(export_dir, "opennmt.tflite")
+        tflite_model = converter.convert()
+        with tf.io.gfile.GFile(tflite_model_path, "wb") as f:
+            f.write(tflite_model)
+        print(
+            "Finished Model Conversion, Converted model can be found at: "
+            + tflite_model_path.replace("\\", "/")
+        )
+
+
+@register_exporter(name="tflite_float16")
+class TFLiteFloat16Exporter(TFLiteExporter):
+    """TensorFlow Lite exporter with float16 quantization."""
+
+    def __init__(self):
+        super().__init__(quantization="float16")
+
+
 @register_exporter(name="ctranslate2")
 class CTranslate2Exporter(Exporter):
     """CTranslate2 exporter."""


Test patch:
diff --git a/opennmt/tests/tflite_test.py b/opennmt/tests/tflite_test.py
new file mode 100644
index 00000000..97321c62
--- /dev/null
+++ b/opennmt/tests/tflite_test.py
@@ -0,0 +1,102 @@
+import os
+
+import tensorflow as tf
+
+from opennmt import data
+from opennmt.tests import test_util
+from opennmt.models import catalog
+from opennmt.utils import exporters
+from parameterized import parameterized
+
+
+def _create_vocab(temp_dir):
+    vocab_path = os.path.join(temp_dir, "vocab.txt")
+    vocab = test_util.make_vocab(vocab_path, ["a", "b", "c"])
+    return vocab, vocab_path
+
+
+def _make_model(model_template, vocab):
+    model = model_template()
+    model.initialize(dict(source_vocabulary=vocab, target_vocabulary=vocab))
+    model.create_variables()
+    return model
+
+
+def _create_dataset(model, temp_dir):
+    data_path = os.path.join(temp_dir, "data.txt")
+    test_util.make_data_file(data_path, ["a a a b b d", "a b b b", "c c"])
+    dataset = model.examples_inputter.make_inference_dataset(data_path, 1)
+    return dataset
+
+
+def _get_predictions(model, dataset, vocab_path):
+    _, tokens_to_ids, _ = data.vocab.create_lookup_tables(vocab_path)
+
+    elem = next(iter(dataset))
+    elem_ids = tf.cast(tf.squeeze(elem["ids"]), dtype=tf.dtypes.int32)
+
+    _, pred = model(elem)
+    pred_ids = tf.squeeze(tokens_to_ids.lookup(pred["tokens"]))
+    tflite_concrete_fn = tf.function(
+        model.infer_tflite,
+        input_signature=[tf.TensorSpec([None], dtype=tf.dtypes.int32, name="ids")],
+    ).get_concrete_function()
+    tflite_pred_ids = tflite_concrete_fn(elem_ids)
+
+    # Modify tflite ids tensor to be the same size as normal model output
+    if tf.size(pred_ids) < tf.size(tflite_pred_ids):
+        tflite_pred_ids = tf.slice(tflite_pred_ids, [0], [tf.size(pred_ids)])
+    elif tf.size(pred_ids) > tf.size(tflite_pred_ids):
+        tflite_pred_ids = tf.pad(
+            tflite_pred_ids, [[tf.size(pred_ids) - tf.size(tflite_pred_ids)]]
+        )
+
+    return pred_ids, tflite_pred_ids
+
+
+def _convert_tflite(model_template, export_dir):
+    vocab, _ = _create_vocab(export_dir)
+    model = _make_model(model_template, vocab)
+    exporter = exporters.TFLiteExporter()
+    exporter.export(model, export_dir)
+
+
+def dir_has_tflite_file(check_dir):
+    extensions = [".lite", ".tflite"]
+    for file in os.listdir(check_dir):
+        for ext in extensions:
+            if file.endswith(ext):
+                return True
+    return False
+
+
+class TFLiteTest(tf.test.TestCase):
+    @parameterized.expand(
+        [
+            [catalog.LuongAttention],
+            [catalog.NMTBigV1],
+        ]
+    )
+    def testTFLiteOutput(self, model):
+        print(model)
+        vocab, vocab_path = _create_vocab(self.get_temp_dir())
+        created_model = _make_model(model, vocab)
+        dataset = _create_dataset(created_model, self.get_temp_dir())
+        pred, tflite_pred = _get_predictions(created_model, dataset, vocab_path)
+        self.assertAllEqual(pred, tflite_pred)
+
+    @parameterized.expand(
+        [
+            [catalog.LuongAttention],
+            [catalog.NMTBigV1],
+        ]
+    )
+    def testTFLiteOutputFile(self, model):
+        print(model)
+        export_dir = self.get_temp_dir()
+        _convert_tflite(model, export_dir)
+        self.assertTrue(dir_has_tflite_file(export_dir))
+
+
+if __name__ == "__main__":
+    tf.test.main()