Check in seq_flow_lite (#10750)

research/seq_flow_lite/WORKSPACE

@@ -16,14 +16,11 @@ http_archive(
 
 http_archive(
     name = "org_tensorflow",
-    sha256 = "40d3203ab5f246d83bae328288a24209a2b85794f1b3e2cd0329458d8e7c1985",
-    strip_prefix = "tensorflow-2.6.0",
-    urls = [
-        "https://github.com/tensorflow/tensorflow/archive/v2.6.0.zip",
-    ],
+    strip_prefix = "tensorflow-2.9.1",
+    sha256 = "9f2dac244e5af6c6a13a7dad6481e390174ac989931942098e7a4373f1bccfc2",
+    urls = ["https://github.com/tensorflow/tensorflow/archive/v2.9.1.zip"],
 )
 
-
 http_archive(
     name = "org_tflite_support",
     strip_prefix = "tflite-support-0861599711ef31de58f62ed3ff6bbcc1e4817ef6",

research/seq_flow_lite/export_to_tflite.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """A tool to export TFLite model."""
 
 import importlib
@@ -22,7 +21,7 @@ import os
 from absl import app
 from absl import flags
 import tensorflow.compat.v1 as tf
-
+import tensorflow_text as tftext
 from layers import base_layers # import seq_flow_lite module
 from layers import projection_layers # import seq_flow_lite module
 from utils import tflite_utils # import seq_flow_lite module
@@ -48,25 +47,33 @@ def main(_):
   with tf.Graph().as_default() as graph:
     with tf.Session(graph=graph) as session:
       text = tf.placeholder(tf.string, shape=[1], name="Input")
-      prxlayer = projection_layers.ProjectionLayer(model_config,
-                                                   base_layers.TFLITE)
-      encoder = model.Encoder(model_config, base_layers.TFLITE)
-      projection, seq_lengh = prxlayer(text)
-      logits = encoder(projection, seq_lengh)
+      inputs = [text]
+      if "pqrnn" in runner_config["name"]:
+        prxlayer = projection_layers.ProjectionLayer(model_config,
+                                                     base_layers.TFLITE)
+        encoder = model.Encoder(model_config, base_layers.TFLITE)
+        projection, seq_length = prxlayer(text)
+        logits = encoder(projection, seq_length)
+      else:
+        byte_int = tftext.ByteSplitter().split(text)
+        token_ids = tf.cast(byte_int, tf.int32).to_tensor()
+        token_ids = tf.reshape(token_ids, [1, -1])
+        token_ids += 3
+        encoder = model.Encoder(model_config, base_layers.TFLITE)
+        logits = encoder(token_ids, None)
       if FLAGS.output == "logits":
-        outputs = logits
+        outputs = [logits]
       elif FLAGS.output == "sigmoid":
-        outputs = tf.math.sigmoid(logits)
+        outputs = [tf.math.sigmoid(logits)]
       else:
         assert FLAGS.output == "softmax", "Unexpected output"
-        outputs = tf.nn.softmax(logits)
+        outputs = [tf.nn.softmax(logits)]
 
       session.run(tf.global_variables_initializer())
       session.run(tf.local_variables_initializer())
       saver = tf.train.Saver()
       saver.restore(session, tf.train.latest_checkpoint(FLAGS.output_dir))
-      tflite_fb = tflite_utils.generate_tflite(session, graph, [text],
-                                               [outputs])
+      tflite_fb = tflite_utils.generate_tflite(session, graph, inputs, outputs)
       output_file_name = os.path.join(FLAGS.output_dir, "tflite.fb")
       with tf.gfile.Open(output_file_name, "wb") as f:
         f.write(tflite_fb)

research/seq_flow_lite/input_fn_reader.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Methods related to input datasets and readers."""
 
 import functools
@@ -21,6 +20,7 @@ import sys
 from absl import logging
 
 import tensorflow as tf
+from tensorflow import estimator as tf_estimator
 import tensorflow_datasets as tfds
 import tensorflow_text as tftext
 
@@ -83,13 +83,13 @@ def create_input_fn(runner_config, mode, drop_remainder):
 
   def _input_fn(params):
     """Method to be used for reading the data."""
-    assert mode != tf.estimator.ModeKeys.PREDICT
-    split = "train" if mode == tf.estimator.ModeKeys.TRAIN else "test"
+    assert mode != tf_estimator.ModeKeys.PREDICT
+    split = "train" if mode == tf_estimator.ModeKeys.TRAIN else "test"
     ds = tfds.load(runner_config["dataset"], split=split)
     ds = ds.batch(params["batch_size"], drop_remainder=drop_remainder)
     ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
     ds = ds.shuffle(buffer_size=100)
-    ds = ds.repeat(count=1 if mode == tf.estimator.ModeKeys.EVAL else None)
+    ds = ds.repeat(count=1 if mode == tf_estimator.ModeKeys.EVAL else None)
     ds = ds.map(
         functools.partial(_post_processor, batch_size=params["batch_size"]),
         num_parallel_calls=tf.data.experimental.AUTOTUNE,

research/seq_flow_lite/layers/BUILD

@@ -82,12 +82,12 @@ py_strict_library(
     srcs = ["misc_layers.py"],
     srcs_version = "PY3",
     deps = [
-        # package tensorflow
         ":embedding_layers",
+        # package tensorflow
         "//layers:base_layers",  # sequence projection
-        "//layers:conv_layers",
+        "//layers:conv_layers",  # sequence projection
         "//layers:dense_layers",  # sequence projection
-        "//layers:normalization_layers",
+        "//layers:normalization_layers",  # sequence projection
         "//layers:quantization_layers",  # sequence projection
     ],
 )
@@ -112,8 +112,8 @@ py_strict_library(
     srcs_version = "PY3",
     deps = [
         # package tensorflow
-        "//layers:base_layers",
-        "//layers:quantization_layers",
+        "//layers:base_layers",  # sequence projection
+        "//layers:quantization_layers",  # sequence projection
     ],
 )
 
@@ -124,11 +124,11 @@ py_strict_library(
     deps = [
         ":embedding_layers",
         # package tensorflow
-        "//layers:base_layers",
-        "//layers:dense_layers",
-        "//layers:normalization_layers",
-        "//layers:quantization_layers",
-        "//tf_ops:tf_custom_ops",
-        "//tf_ops:tf_custom_ops_py",
+        "//layers:base_layers",  # sequence projection
+        "//layers:dense_layers",  # sequence projection
+        "//layers:normalization_layers",  # sequence projection
+        "//layers:quantization_layers",  # sequence projection
+        # "//tf_ops:tf_custom_ops"  # sequence projection
+        "//tf_ops:tf_custom_ops_py",  # sequence projection
     ],
 )

research/seq_flow_lite/layers/base_layers.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Base layer for building models trained with quantization."""
 
 import tensorflow as tf
@@ -57,7 +56,7 @@ class BaseLayer(tf.keras.layers.Layer):
   def add_weight_wrapper(self, shape):
     """Return a weight variable for the given shape."""
     if self.parameters.initializer is not None:
-      initializer = self.parameters.initializer
+      initializer = clone_initializer(self.parameters.initializer)
     else:
       initializer = tf.keras.initializers.GlorotUniform()
     weight = self.add_weight(
@@ -136,3 +135,9 @@ class BaseLayer(tf.keras.layers.Layer):
         maxval=(1.0 - zero_probability),
         dtype=tensor.dtype)
     return tf.math.ceil(rnd)
+
+
+def clone_initializer(initializer):
+  if isinstance(initializer, tf.keras.initializers.Initializer):
+    return initializer.__class__.from_config(initializer.get_config())
+  return initializer

research/seq_flow_lite/layers/conv_layers.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Base layer for convolution."""
 import tensorflow as tf
 

research/seq_flow_lite/layers/dense_layers.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Basic dense layers."""
 import tensorflow as tf
 
@@ -30,6 +29,7 @@ class BaseQDense(base_layers.BaseLayer):
                bias=True,
                rank=2,
                normalize=True,
+               quantize_output=True,
                **kwargs):
     self.units = units
     self.rank = rank
@@ -37,7 +37,9 @@ class BaseQDense(base_layers.BaseLayer):
     self.activation = activation
     self.bias = bias
     self.normalize = normalize
-    self.qoutput = quantization_layers.ActivationQuantization(**kwargs)
+    self.quantize_output = quantize_output
+    if quantize_output:
+      self.qoutput = quantization_layers.ActivationQuantization(**kwargs)
     self._create_normalizer(**kwargs)
     super(BaseQDense, self).__init__(**kwargs)
 
@@ -62,7 +64,10 @@ class BaseQDense(base_layers.BaseLayer):
       outputs = normalize_method(outputs)
     if self.activation:
       outputs = self.activation(outputs)
-    return self.qoutput(outputs)
+    if self.quantize_output:
+      return self.qoutput(outputs)
+    else:
+      return outputs
 
   def _dense_r34(self, inputs, normalize_method):
     bsz = self.get_batch_dimension(inputs)

research/seq_flow_lite/layers/embedding_layers.py

@@ -15,8 +15,8 @@
 """Layers for embedding."""
 import tensorflow as tf
 
-from layers import base_layers
-from layers import quantization_layers
+from layers import base_layers # import seq_flow_lite module
+from layers import quantization_layers # import seq_flow_lite module
 
 
 class EmbeddingLayer(base_layers.BaseLayer):

research/seq_flow_lite/layers/misc_layers.py

@@ -12,15 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Layers for embedding."""
 import math
 import tensorflow as tf
 
 from layers import base_layers # import seq_flow_lite module
-from layers import conv_layers
+from layers import conv_layers # import seq_flow_lite module
 from layers import dense_layers # import seq_flow_lite module
-from layers import embedding_layers
+from layers import embedding_layers # import seq_flow_lite module
 from layers import quantization_layers # import seq_flow_lite module
 
 

research/seq_flow_lite/layers/normalization_layers.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Layers for normalization."""
 import tensorflow as tf
 

research/seq_flow_lite/layers/qrnn_layers.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Layers for QRNN."""
 import tensorflow as tf
 

research/seq_flow_lite/layers/quantization_layers.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Layers for quantization."""
 
 import tensorflow as tf

research/seq_flow_lite/layers/transformer_layers.py

@@ -16,12 +16,12 @@
 # pylint: disable=arguments-renamed
 import tensorflow as tf
 
-from layers import base_layers
-from layers import dense_layers
-from layers import embedding_layers
-from layers import normalization_layers
-from layers import quantization_layers
-from tf_ops import tf_custom_ops_py
+from layers import base_layers # import seq_flow_lite module
+from layers import dense_layers # import seq_flow_lite module
+from layers import embedding_layers # import seq_flow_lite module
+from layers import normalization_layers # import seq_flow_lite module
+from layers import quantization_layers # import seq_flow_lite module
+from tf_ops import tf_custom_ops_py # import seq_flow_lite module
 
 
 class SelfAttention(base_layers.BaseLayer):

research/seq_flow_lite/metric_functions.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Metric functions."""
 import tensorflow.compat.v1 as tf
 

research/seq_flow_lite/models/BUILD

@@ -45,13 +45,13 @@ py_library(
     srcs_version = "PY3",
     deps = [
         # package tensorflow
-        "//layers:base_layers",
-        "//layers:dense_layers",
-        "//layers:embedding_layers",
-        "//layers:misc_layers",
-        "//layers:qrnn_layers",
-        # //tf_ops:tf_custom_ops",
-        "//tf_ops:tf_custom_ops_py",
+        "//layers:base_layers",  # sequence projection
+        "//layers:dense_layers",  # sequence projection
+        "//layers:embedding_layers",  # sequence projection
+        "//layers:misc_layers",  # sequence projection
+        "//layers:qrnn_layers",  # sequence projection
+        # "//tf_ops:tf_custom_ops"  # sequence projection
+        "//tf_ops:tf_custom_ops_py",  # sequence projection
     ],
 )
 
@@ -62,13 +62,13 @@ py_library(
     deps = [
         ":transformer_encoder",
         # package tensorflow
-        "//layers:base_layers",
-        "//layers:embedding_layers",
-        "//layers:misc_layers",
-        "//layers:normalization_layers",
-        "//layers:quantization_layers",
-        # "//tf_ops:tf_custom_ops",
-        "//tf_ops:tf_custom_ops_py",
+        "//layers:base_layers",  # sequence projection
+        "//layers:embedding_layers",  # sequence projection
+        "//layers:misc_layers",  # sequence projection
+        "//layers:normalization_layers",  # sequence projection
+        "//layers:quantization_layers",  # sequence projection
+        # "//tf_ops:tf_custom_ops"  # sequence projection
+        "//tf_ops:tf_custom_ops_py",  # sequence projection
     ],
 )
 
@@ -79,11 +79,11 @@ py_library(
     deps = [
         # package absl/logging
         # package tensorflow
-        "//layers:base_layers",
-        "//layers:embedding_layers",
-        "//layers:transformer_layers",
-        # "//tf_ops:tf_custom_ops",
-        "//tf_ops:tf_custom_ops_py",
+        "//layers:base_layers",  # sequence projection
+        "//layers:embedding_layers",  # sequence projection
+        "//layers:transformer_layers",  # sequence projection
+        # "//tf_ops:tf_custom_ops"  # sequence projection
+        "//tf_ops:tf_custom_ops_py",  # sequence projection
     ],
 )
 
@@ -93,13 +93,13 @@ py_library(
     srcs_version = "PY3",
     deps = [
         # package absl/logging
+        # package tensor2tensor/utils:beam_search
         # package tensorflow
-        # tensor2tensor/utils:beam_search",
-        "//layers:base_layers",
-        "//layers:embedding_layers",
-        "//layers:misc_layers",
-        "//layers:transformer_layers",
-        "//tf_ops:tf_custom_ops",
-        "//tf_ops:tf_custom_ops_py",
+        "//layers:base_layers",  # sequence projection
+        "//layers:embedding_layers",  # sequence projection
+        "//layers:misc_layers",  # sequence projection
+        "//layers:transformer_layers",  # sequence projection
+        # "//tf_ops:tf_custom_ops"  # sequence projection
+        "//tf_ops:tf_custom_ops_py",  # sequence projection
     ],
 )

research/seq_flow_lite/models/byteqrnn.py

@@ -33,11 +33,11 @@ Sample model params:
 from absl import logging
 import tensorflow as tf
 
-from layers import base_layers
-from layers import dense_layers
-from layers import embedding_layers
-from layers import misc_layers
-from layers import qrnn_layers
+from layers import base_layers # import seq_flow_lite module
+from layers import dense_layers # import seq_flow_lite module
+from layers import embedding_layers # import seq_flow_lite module
+from layers import misc_layers # import seq_flow_lite module
+from layers import qrnn_layers # import seq_flow_lite module
 
 
 class Encoder(tf.keras.layers.Layer):

research/seq_flow_lite/models/charformer.py

@@ -16,13 +16,13 @@
 from absl import logging
 import tensorflow as tf
 
-from layers import base_layers
-from layers import dense_layers
-from layers import embedding_layers
-from layers import misc_layers
-from layers import normalization_layers
-from layers import quantization_layers
-from models import transformer_encoder
+from layers import base_layers # import seq_flow_lite module
+from layers import dense_layers # import seq_flow_lite module
+from layers import embedding_layers # import seq_flow_lite module
+from layers import misc_layers # import seq_flow_lite module
+from layers import normalization_layers # import seq_flow_lite module
+from layers import quantization_layers # import seq_flow_lite module
+from models import transformer_encoder # import seq_flow_lite module
 
 
 class Encoder(tf.keras.layers.Layer):

research/seq_flow_lite/models/pqrnn.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Implementation of pQRNN model."""
 
 from absl import logging
@@ -43,7 +42,7 @@ class Encoder(tf.keras.layers.Layer):
     _get_params("qrnn_kernel_width", 3)
     _get_params("qrnn_zoneout_probability")
     _get_params("number_qrnn_layers")
-    _get_params("labels")
+    _get_params("labels", [])
     _get_params("regularizer_scale")
     _get_params("quantize")
 
@@ -66,11 +65,12 @@ class Encoder(tf.keras.layers.Layer):
     self.attention_pool = misc_layers.AttentionPooling(
         parameters=self.parameters)
 
-    self.final_fc = dense_layers.BaseQDense(
-        units=self.num_classes,
-        rank=2,
-        parameters=self.parameters,
-        activation=None)
+    if self.num_classes:
+      self.final_fc = dense_layers.BaseQDense(
+          units=self.num_classes,
+          rank=2,
+          parameters=self.parameters,
+          activation=None)
 
   def call(self, projection, seq_length):
     mask = tf.sequence_mask(
@@ -82,7 +82,11 @@ class Encoder(tf.keras.layers.Layer):
     bottleneck = self.bottleneck_layer(projection, maskr3, inverse_normalizer)
     outputs = self.qrnn_stack(bottleneck, maskr3, inverse_normalizer)
     pre_logits = self.attention_pool(outputs, maskr3, inverse_normalizer)
-    return self.final_fc(pre_logits)
+    if self.num_classes:
+      return self.final_fc(pre_logits)
+    else:
+      return pre_logits
+
 
 class Model(Encoder):
 

research/seq_flow_lite/models/prado.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Lint as: python3
 """Implementation of PRADO model."""
 
 import copy

research/seq_flow_lite/models/sgnn/sgnn_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-# Lint as: python3
 """Tests for seq_flow_lite.sgnn."""
 
 import tensorflow as tf