Rename sequence_projection to seq_flow_lite (#9448)

* Rename sequence_projection to seq_flow_lite

* Rename sequence_projection to seq_flow_lite

research/seq_flow_lite/trainer.py

+# Copyright 2020 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Lint as: python3
+"""A utility for PRADO model to do train, eval, inference and model export."""
+
+import importlib
+import json
+
+from absl import app
+from absl import flags
+from absl import logging
+import tensorflow.compat.v1 as tf
+
+import input_fn_reader # import root module
+import metric_functions # import root module
+
+tf.disable_v2_behavior()
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("config_path", None, "Path to a RunnerConfig.")
+flags.DEFINE_enum("runner_mode", None, ["train", "train_and_eval", "eval"],
+                  "Runner mode.")
+flags.DEFINE_string("master", None, "TensorFlow master URL.")
+flags.DEFINE_string(
+    "output_dir", None,
+    "The output directory where the model checkpoints will be written.")
+flags.DEFINE_bool("use_tpu", False, "Whether to use TPU or GPU/CPU.")
+flags.DEFINE_integer(
+    "num_tpu_cores", 8,
+    "Only used if `use_tpu` is True. Total number of TPU cores to use.")
+
+
+def load_runner_config():
+  with tf.gfile.GFile(FLAGS.config_path, "r") as f:
+    return json.loads(f.read())
+
+
+def create_model(model, model_config, features, mode):
+  """Creates a sequence labeling model."""
+  keras_model = model.Encoder(model_config, mode)
+  logits = keras_model(features["projection"], features["seq_length"])
+  if not model_config["multilabel"]:
+    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        labels=features["label"], logits=logits)
+  else:
+    loss = tf.nn.sigmoid_cross_entropy_with_logits(
+        labels=features["label"], logits=logits)
+  loss = tf.reduce_mean(loss)
+  loss += tf.add_n(keras_model.losses)
+
+  return (loss, logits)
+
+
+def create_optimizer(loss, runner_config):
+  """Returns a train_op using Adam optimizer."""
+  learning_rate = tf.train.exponential_decay(
+      learning_rate=runner_config["learning_rate"],
+      global_step=tf.train.get_global_step(),
+      decay_steps=runner_config["learning_rate_decay_steps"],
+      decay_rate=runner_config["learning_rate_decay_rate"],
+      staircase=True)
+  optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
+  if FLAGS.use_tpu:
+    optimizer = tf.tpu.CrossShardOptimizer(optimizer)
+
+  return optimizer.minimize(loss, global_step=tf.train.get_global_step())
+
+
+def model_fn_builder(runner_config):
+  """Returns `model_fn` closure for TPUEstimator."""
+
+  rel_module_path = "" # empty base dir
+  model = importlib.import_module(rel_module_path + runner_config["name"])
+
+  def model_fn(features, mode, params):
+    """The `model_fn` for TPUEstimator."""
+    del params
+    label_ids = None
+    if mode != tf.estimator.ModeKeys.PREDICT:
+      label_ids = features["label"]
+
+    model_config = runner_config["model_config"]
+    loss, logits = create_model(model, model_config, features, mode)
+
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      train_op = create_optimizer(loss, runner_config)
+      return tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
+          mode=mode, loss=loss, train_op=train_op)
+    elif mode == tf.estimator.ModeKeys.EVAL:
+      if not runner_config["model_config"]["multilabel"]:
+        metric_fn = metric_functions.classification_metric
+      else:
+        metric_fn = metric_functions.labeling_metric
+
+      eval_metrics = (metric_fn, [loss, label_ids, logits])
+      return tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
+          mode=mode, loss=loss, eval_metrics=eval_metrics)
+    else:
+      assert False, "Expected to be called in TRAIN or EVAL mode."
+
+  return model_fn
+
+
+def main(_):
+  runner_config = load_runner_config()
+
+  if FLAGS.output_dir:
+    tf.gfile.MakeDirs(FLAGS.output_dir)
+
+  is_per_host = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V2
+  run_config = tf.estimator.tpu.RunConfig(
+      master=FLAGS.master,
+      model_dir=FLAGS.output_dir,
+      save_checkpoints_steps=runner_config["save_checkpoints_steps"],
+      keep_checkpoint_max=20,
+      tpu_config=tf.estimator.tpu.TPUConfig(
+          iterations_per_loop=runner_config["iterations_per_loop"],
+          num_shards=FLAGS.num_tpu_cores,
+          per_host_input_for_training=is_per_host))
+
+  model_fn = model_fn_builder(runner_config)
+
+  # If TPU is not available, this will fall back to normal Estimator on CPU
+  # or GPU.
+  batch_size = runner_config["batch_size"]
+  estimator = tf.estimator.tpu.TPUEstimator(
+      use_tpu=FLAGS.use_tpu,
+      model_fn=model_fn,
+      config=run_config,
+      train_batch_size=batch_size,
+      eval_batch_size=batch_size,
+      predict_batch_size=batch_size)
+
+  if FLAGS.runner_mode == "train":
+    train_input_fn = input_fn_reader.create_input_fn(
+        runner_config=runner_config,
+        mode=tf.estimator.ModeKeys.TRAIN,
+        drop_remainder=True)
+    estimator.train(
+        input_fn=train_input_fn, max_steps=runner_config["train_steps"])
+  elif FLAGS.runner_mode == "eval":
+    # TPU needs fixed shapes, so if the last batch is smaller, we drop it.
+    eval_input_fn = input_fn_reader.create_input_fn(
+        runner_config=runner_config,
+        mode=tf.estimator.ModeKeys.EVAL,
+        drop_remainder=True)
+
+    for _ in tf.train.checkpoints_iterator(FLAGS.output_dir, timeout=600):
+      result = estimator.evaluate(input_fn=eval_input_fn)
+      for key in sorted(result):
+        logging.info("  %s = %s", key, str(result[key]))
+
+
+if __name__ == "__main__":
+  app.run(main)

research/seq_flow_lite/utils/BUILD

+py_strict_library = py_library
+
+licenses(["notice"])
+
+package(
+    default_visibility = ["//:friends"],  # sequence projection
+)
+
+py_strict_library(
+    name = "tflite_utils",
+    srcs = ["tflite_utils.py"],
+    srcs_version = "PY3",
+    deps = [
+        # package tensorflow
+    ],
+)
+
+py_strict_library(
+    name = "misc_utils",
+    srcs = ["misc_utils.py"],
+    srcs_version = "PY3",
+    deps = [
+        # package tensorflow
+    ],
+)

research/sequence_projection/prado/input_fn_reader.py → research/seq_flow_lite/utils/misc_utils.py

@@ -13,38 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 # Lint as: python3
-"""Methods related to input datasets and readers."""
-
-import functools
-import sys
-from typing import Any, Callable, Mapping, Optional, Tuple, Dict
-
-from absl import logging
-
-import tensorflow.compat.v1 as tf
-import tensorflow_datasets as tfds
-
-
-def imdb_reviews(features, _):
-  return features["text"], features["label"]
-
-
-def civil_comments(features, runner_config):
-  labels = runner_config["model_config"]["labels"]
-  label_tensor = tf.stack([features[label] for label in labels], axis=1)
-  label_tensor = tf.floor(label_tensor + 0.5)
-  return features["text"], label_tensor
-
-
-def goemotions(features, runner_config):
-  labels = runner_config["model_config"]["labels"]
-  label_tensor = tf.stack([features[label] for label in labels], axis=1)
-  return features["comment_text"], tf.cast(label_tensor, tf.float32)
+"""A module for miscelaneous utils."""
+import tensorflow as tf
 
 
 def random_substr(str_tensor, max_words):
   """Select random substring if the input has more than max_words."""
-  word_batch_r = tf.strings.split(str_tensor, result_type="RaggedTensor")
+  word_batch_r = tf.strings.split(str_tensor)
   row_splits = word_batch_r.row_splits
   words = word_batch_r.values
   start_idx = row_splits[:-1]
@@ -72,49 +47,3 @@ def random_substr(str_tensor, max_words):
   new_tensor = tf.RaggedTensor.from_row_splits(
       values=selected_words, row_splits=row_splits)
   return tf.strings.reduce_join(new_tensor, axis=1, separator=" ")
-
-
-def _post_processor(features, runner_config, mode, create_projection,
-                    batch_size):
-  """Post process the data to a form expected by model_fn."""
-  data_processor = getattr(sys.modules[__name__], runner_config["dataset"])
-  text, label = data_processor(features, runner_config)
-  if "max_seq_len" in runner_config["model_config"]:
-    max_seq_len = runner_config["model_config"]["max_seq_len"]
-    logging.info("Truncating text to have at most %d tokens", max_seq_len)
-    text = random_substr(text, max_seq_len)
-  text = tf.reshape(text, [batch_size])
-  num_classes = len(runner_config["model_config"]["labels"])
-  label = tf.reshape(label, [batch_size, num_classes])
-  projection, seq_length = create_projection(runner_config["model_config"],
-                                             mode, text)
-  return {"projection": projection, "seq_length": seq_length, "label": label}
-
-
-def create_input_fn(runner_config: Dict[str, Any], create_projection: Callable,
-                    mode: tf.estimator.ModeKeys, drop_remainder: bool):
-  """Returns an input function to use in the instantiation of tf.estimator.*."""
-
-  def _input_fn(
-      params: Mapping[str, Any]
-  ) -> Tuple[Mapping[str, tf.Tensor], Optional[Mapping[str, tf.Tensor]]]:
-    """Method to be used for reading the data."""
-    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.map(
-        functools.partial(
-            _post_processor,
-            runner_config=runner_config,
-            mode=mode,
-            create_projection=create_projection,
-            batch_size=params["batch_size"]),
-        num_parallel_calls=tf.data.experimental.AUTOTUNE,
-        deterministic=False)
-    return ds
-
-  return _input_fn

research/seq_flow_lite/utils/tflite_utils.py

+# Copyright 2020 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Lint as: python3
+"""Utils to convert to a TFLite model."""
+import tensorflow.compat.v1 as tf
+
+
+def _dump_graph_in_text_format(filename, graph_def):
+  """Dump a tensorflow graph in readable text format."""
+  f = open(filename, 'w')
+
+  for node in graph_def.node:
+    f.write('Node: %s (%s)\n' % (node.name, node.op))
+    for input_name in node.input:
+      f.write('\tInput: %s\n' % input_name)
+  f.close()
+
+
+class InterpreterWithCustomOps(tf.lite.Interpreter):
+
+  def __init__(self, model_content, custom_op_registerers):
+    self._custom_op_registerers = custom_op_registerers
+    super(InterpreterWithCustomOps, self).__init__(model_content=model_content)
+
+
+def set_output_quantized_for_custom_ops(graph_def):
+  """Set output types/quantized flag for custom/unsupported ops."""
+  quantized_custom_ops = {
+      'SequenceStringProjection': [tf.float32.as_datatype_enum],
+      'SequenceStringProjectionV2': [tf.float32.as_datatype_enum],
+      'PoolingOp': [tf.float32.as_datatype_enum],
+      'ExpectedValueOp': [tf.float32.as_datatype_enum],
+      'LayerNormV2': [tf.float32.as_datatype_enum],
+  }
+
+  for node in graph_def.node:
+    if node.op in quantized_custom_ops:
+      node.attr['_output_quantized'].b = True
+      node.attr['_output_types'].list.type[:] = quantized_custom_ops[node.op]
+
+
+def generate_tflite(session, graph, input_tensors, output_tensors):
+  """Generate TFLite model from a session, graph and input/output tensors."""
+  output_nodes = [tensor.name.split(':')[0] for tensor in output_tensors]
+  graph_def = tf.graph_util.convert_variables_to_constants(
+      session, graph.as_graph_def(), output_nodes)
+
+  set_output_quantized_for_custom_ops(graph_def)
+
+# TODO(b/171063452): Bug needs to be fixed to handle this correctly.
+#   def _node_name(tensor):
+#     return tensor.name.split(':')[0]
+
+#   input_arrays_with_shape = [
+#       (_node_name(tensor), None) for tensor in input_tensors
+#   ]
+#   output_arrays = [_node_name(tensor) for tensor in output_tensors]
+#   converter = tf.lite.TFLiteConverter(graph_def, None, None,
+#                                       input_arrays_with_shape, output_arrays)
+  converter = tf.lite.TFLiteConverter(graph_def, input_tensors, output_tensors)
+  converter.inference_type = tf.uint8
+  converter.default_ranges_stats = (127.5, 127.5)
+  converter.quantized_input_stats = {
+      tensor.op.name: (127.5, 127.5) for tensor in input_tensors
+  }
+  converter.allow_custom_ops = True
+  converter.experimental_new_converter = False
+  return converter.convert()

research/sequence_projection/BUILD

-licenses(["notice"])
-
-package(
-    default_visibility = [
-        "//:__subpackages__",
-    ],
-)

research/sequence_projection/demo/BUILD

-# Demo app for sequence projection model

research/sequence_projection/prado/common_layer.py

-# Copyright 2020 The TensorFlow Authors All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-# python3
-"""Common layer creator."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow.compat.v1 as tf
-
-from tensorflow.python.training import moving_averages  # pylint: disable=g-direct-tensorflow-import
-
-
-class CommonLayers(object):
-  """A base class that defines TfLite compatible NN layers."""
-
-  def __init__(self,
-               mode,
-               regularizer_scale=0.0,
-               weights_initializer=tf.keras.initializers.glorot_uniform(),
-               quantization_enabled=True):
-    """PoDLayers constructor.
-
-    Args:
-      mode: Graph creation mode.
-      regularizer_scale: Optional regularizer for the weights.
-      weights_initializer: Optional initializer for the weights.
-      quantization_enabled: Enables quantization of weights and activation in
-        the DNN.
-    """
-    self._mode = mode
-    self._regularizer_scale = regularizer_scale
-    self._weights_initializer = weights_initializer
-    self._quantization_enabled = quantization_enabled
-    # Batch normalization is the default normalization scheme.
-    self._normalizer = self.batch_normalization
-    self._moment_fn = None
-
-  def qrange_sigmoid(self, tensor):
-    """Quantize the tensor in sigmoid range (0.0, 1.0)."""
-    return tf.fake_quant_with_min_max_args(
-        tensor, 0.0, 1.0) if self._quantization_enabled else tensor
-
-  def qrange_tanh(self, tensor):
-    """Quantize the tensor in tanh range (-1.0, 1.0)."""
-    return tf.fake_quant_with_min_max_args(
-        tensor, -1.0, 1.0) if self._quantization_enabled else tensor
-
-  def _quantized_tanh(self, tensor):
-    """Apply tanh op and quantize in the range (-1.0, 1.0)."""
-    return self.qrange_tanh(tf.tanh(tensor))
-
-  def _quantized_sigmoid(self, tensor):
-    """Apply sigmoid op and quantize in the range (0.0, 1.0)."""
-    return self.qrange_sigmoid(tf.sigmoid(tensor))
-
-  def set_moment_fn(self, moment_fn):
-    """Set a moment function that will be used by batch norm."""
-    self._moment_fn = moment_fn
-
-  def set_regularizer_scale(self, regularizer_scale):
-    """Override / set a new weights regularizer scale."""
-    self._regularizer_scale = regularizer_scale
-
-  def set_variable_length_moment_fn(self, sequence_length, max_sequence_length):
-    """Set variable length moment function for use in batch norm.
-
-    Args:
-      sequence_length: An vector of sequence lengths.
-      max_sequence_length: Padding length for the batch.
-
-    Returns:
-      Returns sequence mask.
-    """
-    mask = tf.sequence_mask(
-        sequence_length, maxlen=max_sequence_length, dtype=tf.float32)
-    mask = tf.expand_dims(mask, 2)
-
-    mask_r4 = tf.expand_dims(mask, 3)
-    mask_r2 = tf.reshape(mask, [-1, 1])
-    inverse_numsteps = tf.math.reciprocal(tf.reduce_sum(mask))
-
-    def _varlen_moment_fn(input_tensor, axes):
-      """Moment function to use with batch normalization."""
-      input_tensor_shape = input_tensor.get_shape().as_list()
-      input_tensor_rank = len(input_tensor_shape)
-      if input_tensor_rank == 2:
-        input_tensor = mask_r2 * input_tensor
-      elif input_tensor_rank == 4:
-        assert input_tensor_shape[2] == 1
-        input_tensor = mask_r4 * input_tensor
-      else:
-        assert False, "Supports rank2 and rank4 tensors."
-      ex = tf.reduce_sum(input_tensor, axis=axes) * inverse_numsteps
-      exx = tf.reduce_sum(
-          input_tensor * input_tensor, axis=axes) * inverse_numsteps
-      return ex, (exx - ex * ex)
-
-    self._moment_fn = _varlen_moment_fn
-    return mask
-
-  def batch_normalization(self, input_tensor, decay=0.999):
-    """Add batch normalization network structure after input_tensor.
-
-    It performs batch normalization of the input tensor. This routine is
-    verified to works for rank 4 or 2 tensors.
-
-    Args:
-      input_tensor: Input tensor that needs to be normalized.
-      decay: Moving average decay
-
-    Returns:
-      A tensor that is normalized.
-    """
-    input_tensor_shape = input_tensor.get_shape().as_list()
-    nstat = input_tensor_shape[-1]
-    reduce_dims = list(range(len(input_tensor_shape) - 1))
-
-    with tf.variable_scope(name_or_scope=None, default_name="batch_norm"):
-      offset = tf.get_variable(
-          "offset",
-          shape=[nstat],
-          initializer=tf.zeros_initializer,
-          trainable=True)
-      scale = tf.get_variable(
-          "scale",
-          shape=[nstat],
-          initializer=tf.ones_initializer,
-          trainable=True)
-      moving_mean = tf.get_variable(
-          "moving_mean",
-          shape=[nstat],
-          initializer=tf.zeros_initializer,
-          trainable=False)
-      moving_var = tf.get_variable(
-          "moving_variance",
-          shape=[nstat],
-          initializer=tf.ones_initializer,
-          trainable=False)
-
-      if self._mode == tf.estimator.ModeKeys.TRAIN:
-        # During training compute summay stats, update them to moving average
-        # variables and use the summary stas for batch normalization.
-        moment_fn = self._moment_fn or tf.nn.moments
-        mean_mom, var_mom = moment_fn(input_tensor, reduce_dims)
-        with tf.control_dependencies([
-            moving_averages.assign_moving_average(
-                moving_mean, mean_mom, decay, name="mean_op"),
-            moving_averages.assign_moving_average(
-                moving_var, var_mom, decay, name="variance_op")
-        ]):
-          tensor = tf.nn.batch_normalization(
-              input_tensor,
-              mean_mom,
-              var_mom,
-              offset,
-              scale,
-              1e-9,
-              name="batch_norm_core")
-      else:
-        # During eval/inference use the moving average variable for batch
-        # normalization. The variables would be frozen to constants before
-        # saving graph.
-        tensor = tf.nn.batch_normalization(
-            input_tensor,
-            moving_mean,
-            moving_var,
-            offset,
-            scale,
-            1e-9,
-            name="batch_norm_core")
-    return tensor
-
-  def get_quantization_ranges(self, tensor, ema_decay=0.99):
-    """Perform fake quantization of the tensor.
-
-    The method computes ranges for quantization by first computing the
-    batch min/max and then computing a moving average of the min/max across
-    batches. The moving average of min/max is used for quantization during
-    inference. During training the batch min/maxs are used directly.
-
-    Args:
-      tensor: Input tensor that needs to be quantized.
-      ema_decay: Moving average decay
-
-    Returns:
-      Min/Max for fake quantization.
-    """
-    # If neither quantization is enabled, nor are we calculating ranges for
-    # floating point models, this method is a no-op.
-    if not self._quantization_enabled:
-      return None, None
-
-    # Calculate min/max for the tensor.
-    min_var = tf.get_variable("min", initializer=0.0, trainable=False)
-    max_var = tf.get_variable("max", initializer=1.0, trainable=False)
-
-    if self._mode == tf.estimator.ModeKeys.TRAIN:
-      # During training estimate moving average for min/max. Use the min/max
-      # values directly for quantization.
-      ops = []
-      batch_min = tf.reduce_min(tensor, name="BatchMin")
-      # Toco expects 0.0 to be part of the quantization range.
-      batch_min = tf.minimum(batch_min, 0.0)
-      ops.append(
-          moving_averages.assign_moving_average(min_var, batch_min, ema_decay))
-
-      batch_max = tf.reduce_max(tensor, name="BatchMax")
-      # Toco expects 0.0 to be part of the quantization range.
-      batch_max = tf.maximum(batch_max, 0.0)
-      ops.append(
-          moving_averages.assign_moving_average(max_var, batch_max, ema_decay))
-
-      with tf.control_dependencies(ops):
-        return tf.identity(batch_min), tf.identity(batch_max)
-    else:
-      # During inference/eval use the moving average min/maxs for
-      # quantization.
-      return min_var, max_var
-
-  def quantization(self, tensor, ema_decay=0.99, num_bits=8):
-    """Perform fake quantization of the tensor.
-
-    The method performs fake quantization of the tensor by first computing the
-    batch min/max and then computing a moving average of the min/max across
-    batches. The moving average of min/max is used for quantization during
-    inference. During training the batch min/maxs are used directly.
-
-    Args:
-      tensor: Input tensor that needs to be quantized.
-      ema_decay: Moving average decay
-      num_bits: Number of bits used for quantization
-
-    Returns:
-      Quantized tensor.
-    """
-    with tf.variable_scope(
-        name_or_scope=None, default_name="MovingAvgQuantize"):
-      min_tensor, max_tensor = self.get_quantization_ranges(tensor, ema_decay)
-      if min_tensor is None or max_tensor is None:
-        return tensor
-      else:
-        return tf.fake_quant_with_min_max_vars(
-            tensor, min_tensor, max_tensor, num_bits=num_bits)
-
-  def _weight_quantization(self, tensor, num_bits=8):
-    """Quantize weights when enabled."""
-    if not self._quantization_enabled:
-      return tensor
-
-    # For infer mode, toco computes the min/max from the weights offline to
-    # quantize it. During train/eval this is computed from the current value
-    # in the session by the graph itself.
-    modes = set([tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL])
-    if self._mode in modes:
-      batch_min = tf.reduce_min(tensor, name="BatchMin")
-      # Toco expects 0.0 to be part of the quantization range.
-      batch_min = tf.minimum(batch_min, 0.0)
-
-      batch_max = tf.reduce_max(tensor, name="BatchMax")
-      # Toco expects 0.0 to be part of the quantization range.
-      batch_max = tf.maximum(batch_max, 0.0)
-
-      return tf.fake_quant_with_min_max_vars(
-          tensor, batch_min, batch_max, num_bits=num_bits)
-    else:
-      return tensor
-
-  def _get_weight(self, shape, num_bits=8):
-    """Return a weight variable for the given shape.
-
-    The disable_pruning flag overrides the global pruning_obj object. When set
-    to True, the returned weight tensor is not pruned.
-    Args:
-      shape: Shape of the weight tensor
-      num_bits: Number of bits to use for the variable.
-
-    Returns:
-      Quantized tensor with the mask and threshold variables needed for pruning.
-
-    """
-    weight = tf.get_variable(
-        "weight", shape, initializer=self._weights_initializer)
-    if self._regularizer_scale > 0.0:
-      reg_loss = tf.nn.l2_loss(weight) * tf.convert_to_tensor(
-          self._regularizer_scale)
-      tf.losses.add_loss(
-          reg_loss, loss_collection=tf.GraphKeys.REGULARIZATION_LOSSES)
-    return self._weight_quantization(weight, num_bits=num_bits)
-
-  def _get_bias(self, shape):
-    weight = tf.get_variable("bias", shape, initializer=tf.zeros_initializer())
-    if self._regularizer_scale > 0.0:
-      reg_loss = tf.nn.l2_loss(weight) * tf.convert_to_tensor(
-          self._regularizer_scale)
-      tf.losses.add_loss(
-          reg_loss, loss_collection=tf.GraphKeys.REGULARIZATION_LOSSES)
-    return weight
-
-  def zero_beyond_sequence_length(self, sequence_length, gate):
-    """Generate a binary mask for the sequence based on the timestep's validity.
-
-    Args:
-      sequence_length: The sequence length tensor of [batch size] elements.
-      gate: A gate tensor used by the QuasiRNN cell to infer shape from it.
-
-    Returns:
-      Mask tensor with one for valid time and zero for invalid timestep.
-    """
-    mask = tf.sequence_mask(
-        sequence_length, maxlen=tf.shape(gate)[1], dtype=tf.float32)
-    return tf.expand_dims(mask, 2)
-
-  def _convolution2d(self,
-                     inputs,
-                     kernel_size,
-                     filters,
-                     stride,
-                     padding,
-                     dilations=None,
-                     weight_mask=None,
-                     scope="convolution2d"):
-    """Linear part of the convolution layer."""
-    if isinstance(stride, int):
-      strides = [1, stride, stride, 1]
-    else:
-      if not isinstance(stride, list) or len(stride) != 2:
-        raise ValueError("`Stride` should be an integer or a list of length 2")
-      strides = [1, stride[0], stride[1], 1]
-    if dilations is not None:
-      if not isinstance(dilations, list) or len(dilations) != 2:
-        raise ValueError("`Dilations` should be an integer list of length 2")
-      dilations = [1, dilations[0], dilations[1], 1]
-    else:
-      dilations = [1, 1, 1, 1]
-
-    with tf.variable_scope(name_or_scope=None, default_name=scope):
-      input_channels = inputs.get_shape().as_list()[-1]
-      kernel_shape = kernel_size + [input_channels, filters]
-      weight = self._get_weight(kernel_shape)
-      if weight_mask is not None:
-        # Tensor multiply for disabling backprop
-        weight = weight * weight_mask
-      bias = self._get_bias([filters])
-
-      features = tf.nn.conv2d(
-          inputs, weight, strides, padding, dilations=dilations)
-      return tf.nn.bias_add(features, bias)
-
-  def convolution2d(self,
-                    inputs,
-                    kernel_size,
-                    filters,
-                    scope="convolution2d",
-                    stride=1,
-                    padding="SAME",
-                    dilations=None,
-                    weight_mask=None,
-                    activation=tf.nn.relu,
-                    normalization=True):
-    """Creates a 2d convolution layer.
-
-    Performs batch normalization to the tensor pre activation and fake
-    quantization post activation.
-
-    Args:
-      inputs: Input tensor, that is expected to be a rank 4 tensor.
-      kernel_size: 2D convolution kernel size (2 tuple).
-      filters: Number of output channels (integer).
-      scope: A string that would be used as variable scope for the layer.
-      stride: Convolution stride, can be a constant or a 2 tuple.
-      padding: Padding to use for the convolution.
-      dilations: tuple of size 2 specifying the dilation rates for input height
-        and width respectively. Refer to tf.nn.conv2d API for more details.
-      weight_mask: A floating point numpy array or constant tensor mask to turn
-        off weights in the convolution kernel.
-      activation: Activation function to be used, Relu is used by default.
-      normalization: A boolean flag indicating if batchnorm should be performed.
-
-    Returns:
-      Tensor result of the convolution layer.
-
-    Raises:
-      ValueError: If inputs is not a rank 4 tensor
-      ValueError: If kernel_size is not a list or tuple of length 2
-    """
-    if len(inputs.get_shape().as_list()) != 4:
-      raise ValueError("`inputs` should be a rank 4 tensor. "
-                       "Was: {}.".format(len(inputs.get_shape().as_list())))
-
-    kernel_size = list(kernel_size)
-    if len(kernel_size) != 2:
-      raise ValueError("`kernel_size` should be a tuple or list of length 2. "
-                       "Was: {}.".format(kernel_size))
-
-    features_rank4 = self._convolution2d(
-        inputs,
-        kernel_size,
-        filters,
-        stride,
-        padding,
-        dilations,
-        weight_mask=weight_mask,
-        scope=scope)
-
-    if normalization and self._normalizer:
-      features_rank4 = self._normalizer(features_rank4)
-    if activation is not None:
-      features_rank4 = activation(features_rank4)
-
-    return self.quantization(features_rank4)
-
-  def _fully_connected(self,
-                       features,
-                       output_size,
-                       scope="fully_connected",
-                       use_bias=True):
-    """Performs fully connected operation."""
-    with tf.variable_scope(name_or_scope=None, default_name=scope):
-      weight = self._get_weight(
-          [features.get_shape().as_list()[-1], output_size])
-      bias = self._get_bias([output_size])
-      features = tf.matmul(features, weight)
-      return tf.nn.bias_add(features, bias) if use_bias else features
-
-  def fully_connected(self,
-                      features,
-                      output_size,
-                      scope="fully_connected",
-                      activation=tf.nn.relu,
-                      normalization=True,
-                      use_bias=True):
-    """Creates a fully connected layer.
-
-    Performs batch normalization to the tensor pre activation and fake
-    quantization post activation.
-
-    Args:
-      features: Input features to the fully connected layer.
-      output_size: Number of output features.
-      scope: A variable scope for the connected layer.
-      activation: activation function to be used, Relu is used by default.
-      normalization: A flag indicating if batchnorm should be performed.
-      use_bias: If True, bias is added to the result
-
-    Returns:
-      Tensor result of the fully connected layer.
-
-    Raises:
-      ValueError: If last dimension of features is dynamic (shape = None).
-    """
-    input_shape = features.get_shape().as_list()
-    if not input_shape[-1]:
-      raise ValueError("Last dimension of features should be static")
-
-    need_reshape = len(input_shape) > 2
-    input_tensor = features
-    if need_reshape:
-      features = tf.reshape(features, [-1, input_shape[-1]])
-
-    features = self._fully_connected(
-        features, output_size, scope=scope, use_bias=use_bias)
-
-    if normalization and self._normalizer:
-      features = self._normalizer(features)
-
-    if activation:
-      # Batch normalization is done pre activation as suggested in the original
-      # paper. Quantization is done post activation because the range will
-      # change after applying the squashing function.
-      features = activation(features)
-    features = self.quantization(features)
-    if not need_reshape:
-      return features
-    else:
-      # The fully connected layer changes the last dimension to output_size.
-      # If a reshape was done before applying the fully connected layer, change
-      # it back to the right rank. If the input dimensions are known use the
-      # static shape otherwise use the shape tensor.
-      if sum([val is None for val in input_shape]) <= 1:
-        # Just one dynamic shape, we can reshape with -1
-        output_shape = [-1 if val is None else val for val in input_shape]
-      else:
-        input_shape_tensor = tf.shape(input_tensor)
-        output_shape = [
-            shape or input_shape_tensor[index]
-            for index, shape in enumerate(input_shape)
-        ]
-      output_shape[-1] = output_size
-      return tf.reshape(features, output_shape)

research/sequence_projection/prado/prado_model.py

-# Copyright 2020 The TensorFlow Authors All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Tensorflow graph creator for PRADO model."""
-
-import collections
-import functools
-
-from typing import Mapping, Dict, Any
-from absl import logging
-import tensorflow.compat.v1 as tf
-
-from prado import common_layer # import sequence_projection module
-from tf_ops import sequence_string_projection_op as ssp # import sequence_projection module
-
-_NGRAM_INFO = [
-    {
-        "name": "unigram",
-        "padding": 0,
-        "kernel_size": [1, 1],
-        "mask": None
-    },
-    {
-        "name": "bigram",
-        "padding": 1,
-        "kernel_size": [2, 1],
-        "mask": None
-    },
-    {
-        "name": "trigram",
-        "padding": 2,
-        "kernel_size": [3, 1],
-        "mask": None
-    },
-    {
-        "name": "bigramskip1",
-        "padding": 2,
-        "kernel_size": [3, 1],
-        "mask": [[[[1]]], [[[0]]], [[[1]]]]
-    },
-    {
-        "name": "bigramskip2",
-        "padding": 3,
-        "kernel_size": [4, 1],
-        "mask": [[[[1]]], [[[0]]], [[[0]]], [[[1]]]]
-    },
-    {
-        "name": "fourgram",
-        "padding": 3,
-        "kernel_size": [4, 1],
-        "mask": None
-    },
-    {
-        "name": "fivegram",
-        "padding": 4,
-        "kernel_size": [5, 1],
-        "mask": None
-    },
-]
-
-
-def _get_params(model_config, varname, default_value=None):
-  value = model_config[varname] if varname in model_config else default_value
-  logging.info("%s = %s", varname, value)
-  return value
-
-
-def create_projection(model_config, mode, inputs):
-  """Create projection."""
-  feature_size = _get_params(model_config, "feature_size")
-  text_distortion_probability = _get_params(model_config,
-                                            "text_distortion_probability", 0.0)
-  max_seq_len = _get_params(model_config, "max_seq_len", 0)
-  add_eos_tag = _get_params(model_config, "add_eos_tag")
-  is_training = mode == tf.estimator.ModeKeys.TRAIN
-  distortion_probability = text_distortion_probability if is_training else 0.0
-  raw_string = tf.identity(inputs, "Input")
-  features, _, seq_length = ssp.sequence_string_projection(
-      input=raw_string,
-      feature_size=feature_size,
-      max_splits=max_seq_len - 1,
-      distortion_probability=distortion_probability,
-      split_on_space=True,
-      add_eos_tag=add_eos_tag,
-      vocabulary="")
-
-  if mode != tf.estimator.ModeKeys.PREDICT and max_seq_len > 0:
-    pad_value = [[0, 0], [0, max_seq_len - tf.shape(features)[1]], [0, 0]]
-    features = tf.pad(features, pad_value)
-    batch_size = inputs.get_shape().as_list()[0]
-    features = tf.reshape(features,
-                          [batch_size, max_seq_len, feature_size])
-  return features, seq_length
-
-
-def _fully_connected(pod_layers, tensor, num_features, mode, bsz, keep_prob):
-  """Fully connected layer."""
-  tensor_out = pod_layers.fully_connected(tensor, num_features)
-  if mode == tf.estimator.ModeKeys.TRAIN:
-    tensor_out = tf.nn.dropout(tensor_out, rate=(1 - keep_prob))
-  return tf.reshape(tensor_out, [bsz, -1, 1, num_features])
-
-
-def _get_convolutional_layer(pod_layers, head_type, channels, valid_step_mask,
-                             tensor, invalid_value):
-  """Get convolutional layer."""
-  info = _NGRAM_INFO[head_type]
-  pad = info["padding"]
-  weight_mask = info["mask"]
-  kernel_size = info["kernel_size"]
-  paddings = [[0, 0], [0, pad], [0, 0], [0, 0]]
-  # Padding before convolution and using 'valid' instead of 'same' padding
-  # structure ensures that the convolution output is identical between
-  # train/eval and inference models. It also ensures that they lineup
-  # correctly with the valid_step_mask.
-  tensor = tf.pad(tensor, paddings) if pad != 0 else tensor
-  # Not using activation allows a bigram feature to de-emphasize a feature
-  # that triggers positive for unigram for example. The output weights
-  # should be allowed to be positve or negative for this to happen.
-  tensor = pod_layers.convolution2d(
-      tensor,
-      kernel_size,
-      channels,
-      padding="VALID",
-      weight_mask=weight_mask,
-      activation=None)
-  if valid_step_mask is not None:
-    tensor = tensor * valid_step_mask + (1 - valid_step_mask) * invalid_value
-  return tensor
-
-
-def _get_predictions(pod_layers, head_type, keys, values, channels,
-                     valid_step_mask):
-  """Get predictions using one ngram head."""
-  conv_layer = functools.partial(_get_convolutional_layer, pod_layers,
-                                 head_type, channels, valid_step_mask)
-  return conv_layer(keys, -100), conv_layer(values, 0)
-
-
-def reduce_tensors(pod_layers, bsz, attention_logits, values):
-  """Reduce information using attention."""
-  channels = attention_logits.get_shape().as_list()[-1]
-  attention_logits = tf.reshape(attention_logits, [bsz, -1, channels])
-  values = tf.reshape(values, [bsz, -1, channels])
-
-  with tf.variable_scope("attention_expected_value"):
-    attention_logits = tf.identity(attention_logits, "attention_logits_in")
-    values = tf.identity(values, "values_in")
-    attention_logits = tf.transpose(attention_logits, [0, 2, 1])
-    values = tf.transpose(values, [0, 2, 1])
-    attention = tf.nn.softmax(attention_logits, axis=2)
-    evalue = tf.reduce_sum(attention * values, axis=[2])
-    evalue = tf.identity(evalue, "expected_value_out")
-  return pod_layers.quantization(evalue)
-
-
-def ngram_attention_args_v2(projection, seq_length, mode, num_classes,
-                            model_args):
-  """Implements an ngram attention network.
-
-  Args:
-    projection: Projection features from text.
-    seq_length: Sequence length.
-    mode: Model creation mode (train, eval or predict).
-    num_classes: Number of classes to be predicted.
-    model_args: A namedtuple containing all model arguments.
-
-  Returns:
-    A tensor corresponding to the logits of the graph.
-  """
-
-  pod_layers = common_layer.CommonLayers(
-      mode, quantization_enabled=model_args.quantize)
-
-  features = pod_layers.qrange_tanh(projection)
-  bsz = features.get_shape().as_list()[0] or tf.shape(features)[0]
-
-  # Regularizer just for the word embedding.
-  pod_layers.set_regularizer_scale(model_args.embedding_regularizer_scale)
-  values = _fully_connected(pod_layers, features, model_args.embedding_size,
-                            mode, bsz, model_args.keep_prob)
-  keys = _fully_connected(pod_layers, features, model_args.embedding_size, mode,
-                          bsz, model_args.keep_prob)
-
-  # Regularizer for the rest of the network.
-  pod_layers.set_regularizer_scale(model_args.network_regularizer_scale)
-
-  valid_step_mask = None
-  if mode != tf.estimator.ModeKeys.PREDICT:
-    valid_step_mask = pod_layers.zero_beyond_sequence_length(
-        seq_length, features)
-    valid_step_mask = tf.expand_dims(valid_step_mask, 3)
-    # Mask out the sentence beyond valid sequence length for training graph.
-    # This ensures that these values are all zeroed out. Without masking, the
-    # fully connected layer before will make them take an arbitrary constant
-    # value during training/eval in the minibatches. But these values won't
-    # be present during inference as the inference is not batched.
-    keys = valid_step_mask * keys
-    values = valid_step_mask * values
-    pod_layers.set_variable_length_moment_fn(seq_length, tf.shape(features)[1])
-
-  multi_head_predictions = []
-  for head_type, head in zip(model_args.head_types, model_args.heads):
-    if not head:
-      continue
-    att_logits, att_values = _get_predictions(pod_layers, head_type, keys,
-                                              values, head, valid_step_mask)
-    multi_head_predictions.append(
-        reduce_tensors(pod_layers, bsz, att_logits, att_values))
-  multi_head_predictions = tf.concat(multi_head_predictions, axis=1)
-  multi_head_predictions = pod_layers.quantization(multi_head_predictions)
-  # Sequence dimension has been summed out, so we don't need special moment
-  # function.
-  pod_layers.set_moment_fn(None)
-
-  output = multi_head_predictions
-
-  # Add FC layers before the logits.
-  for fc_layer_size in model_args.pre_logits_fc_layers:
-    output = pod_layers.fully_connected(
-        output, fc_layer_size, activation=tf.nn.relu)
-
-  return pod_layers.fully_connected(output, num_classes, activation=None)
-
-
-def create_encoder(model_config: Dict[str, Any], projection: tf.Tensor,
-                   seq_length: tf.Tensor,
-                   mode: tf.estimator.ModeKeys) -> Mapping[str, tf.Tensor]:
-  """Implements a simple attention network for brand safety."""
-
-  args = {}
-
-  def _get_params(varname, default_value=None):
-    value = model_config[varname] if varname in model_config else default_value
-    logging.info("%s = %s", varname, value)
-    args[varname] = value
-
-  _get_params("labels")
-  _get_params("quantize", True)
-  _get_params("max_seq_len", 0)
-  _get_params("max_seq_len_inference", 0)
-  _get_params("split_on_space", True)
-  _get_params("exclude_nonalphaspace_unicodes", False)
-  _get_params("embedding_regularizer_scale", 35e-3)
-  _get_params("embedding_size", 64)
-  _get_params("heads", [0, 64, 64, 0, 0])
-  _get_params("feature_size", 512)
-  _get_params("network_regularizer_scale", 1e-4)
-  _get_params("keep_prob", 0.5)
-  _get_params("word_novelty_bits", 0)
-  _get_params("doc_size_levels", 0)
-  _get_params("pre_logits_fc_layers", [])
-  args["head_types"] = list(range(len(args["heads"])))
-  args["text_distortion_probability"] = 0.0
-  if mode == tf.estimator.ModeKeys.TRAIN:
-    _get_params("text_distortion_probability", 0.25)
-  model_args = collections.namedtuple("ModelArgs", sorted(args))(**args)
-  num_classes = len(model_args.labels)
-  logits = ngram_attention_args_v2(
-      projection=projection,
-      seq_length=seq_length,
-      mode=mode,
-      num_classes=num_classes,
-      model_args=model_args)
-  outputs = {
-      "logits":
-          tf.identity(logits, "Logits"),
-      "label_map":
-          tf.constant(list(model_args.labels), tf.string, name="LabelMap")
-  }
-  return outputs

research/sequence_projection/prado/runner.py

-# Copyright 2020 The TensorFlow Authors All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-# Lint as: python3
-"""A utility for PRADO model to do train, eval, inference and model export."""
-
-import json
-import os
-from typing import Any, Mapping, Optional, Sequence, Tuple, Dict
-
-from absl import logging
-import tensorflow.compat.v1 as tf
-
-from tensorflow.core.framework import types_pb2 as tf_types
-from tensorflow.python.tools import optimize_for_inference_lib  # pylint: disable=g-direct-tensorflow-import
-from prado import input_fn_reader # import sequence_projection module
-from prado import metric_functions # import sequence_projection module
-from prado import prado_model as model # import sequence_projection module
-
-tf.disable_v2_behavior()
-
-
-FLAGS = tf.flags.FLAGS
-
-tf.flags.DEFINE_string("config_path", None, "Path to a RunnerConfig.")
-tf.flags.DEFINE_enum("runner_mode", None,
-                     ["train", "train_and_eval", "eval", "export"],
-                     "Runner mode.")
-tf.flags.DEFINE_string("master", None, "TensorFlow master URL.")
-tf.flags.DEFINE_string(
-    "output_dir", None,
-    "The output directory where the model checkpoints will be written.")
-tf.flags.DEFINE_bool("use_tpu", False, "Whether to use TPU or GPU/CPU.")
-tf.flags.DEFINE_integer(
-    "num_tpu_cores", 8,
-    "Only used if `use_tpu` is True. Total number of TPU cores to use.")
-
-
-def load_runner_config() -> Dict[str, Any]:
-  with tf.gfile.GFile(FLAGS.config_path, "r") as f:
-    return json.loads(f.read())
-
-
-def create_model(
-    model_config: Dict[str, Any], projection: tf.Tensor, seq_length: tf.Tensor,
-    mode: tf.estimator.ModeKeys, label_ids: tf.Tensor
-) -> Tuple[tf.Tensor, tf.Tensor, Mapping[str, tf.Tensor]]:
-  """Creates a sequence labeling model."""
-  outputs = model.create_encoder(model_config, projection, seq_length, mode)
-  with tf.variable_scope("loss"):
-    loss = None
-    per_example_loss = None
-    if mode != tf.estimator.ModeKeys.PREDICT:
-      if not model_config["multilabel"]:
-        per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
-            labels=label_ids, logits=outputs["logits"])
-      else:
-        per_label_loss = tf.nn.sigmoid_cross_entropy_with_logits(
-            labels=label_ids, logits=outputs["logits"])
-        per_example_loss = tf.reduce_mean(per_label_loss, axis=1)
-      loss = tf.reduce_mean(per_example_loss)
-      loss += tf.add_n(tf.compat.v1.losses.get_regularization_losses())
-
-  return (loss, per_example_loss, outputs)
-
-
-def create_optimizer(loss: tf.Tensor, runner_config: Dict[str,
-                                                          Any]) -> tf.Operation:
-  """Returns a train_op using Adam optimizer."""
-  learning_rate = tf.train.exponential_decay(
-      learning_rate=runner_config["learning_rate"],
-      global_step=tf.train.get_global_step(),
-      decay_steps=runner_config["learning_rate_decay_steps"],
-      decay_rate=runner_config["learning_rate_decay_rate"],
-      staircase=True)
-  optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
-  if FLAGS.use_tpu:
-    optimizer = tf.tpu.CrossShardOptimizer(optimizer)
-  else:
-    tf.compat.v1.summary.scalar("learning_rate", learning_rate)
-
-  tvars = tf.trainable_variables()
-  grads = tf.gradients(loss, tvars)
-  train_op = optimizer.apply_gradients(
-      zip(grads, tvars), global_step=tf.train.get_global_step())
-  return train_op
-
-
-def model_fn_builder(runner_config: Dict[str, Any]):
-  """Returns `model_fn` closure for TPUEstimator."""
-
-  def model_fn(
-      features: Mapping[str, tf.Tensor],
-      mode: tf.estimator.ModeKeys,
-      params: Optional[Mapping[str, Any]]  # pylint: disable=unused-argument
-  ) -> tf.compat.v1.estimator.tpu.TPUEstimatorSpec:
-    """The `model_fn` for TPUEstimator."""
-
-    projection = features["projection"]
-    seq_length = features["seq_length"]
-    label_ids = None
-    if mode != tf.estimator.ModeKeys.PREDICT:
-      label_ids = features["label"]
-
-    (total_loss, per_example_loss,
-     model_outputs) = create_model(runner_config["model_config"], projection,
-                                   seq_length, mode, label_ids)
-
-    if mode == tf.estimator.ModeKeys.TRAIN:
-      train_op = create_optimizer(total_loss, runner_config)
-      return tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
-          mode=mode, loss=total_loss, train_op=train_op)
-
-    if mode == tf.estimator.ModeKeys.EVAL:
-
-      if not runner_config["model_config"]["multilabel"]:
-        metric_fn = metric_functions.classification_metric
-      else:
-        metric_fn = metric_functions.labeling_metric
-
-      eval_metrics = (metric_fn,
-                      [per_example_loss, label_ids, model_outputs["logits"]])
-      return tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
-          mode=mode, loss=total_loss, eval_metrics=eval_metrics)
-
-    # Prediction mode
-    return tf.compat.v1.estimator.tpu.TPUEstimatorSpec(
-        mode=mode, predictions=model_outputs)
-
-  return model_fn
-
-
-def set_output_types_and_quantized(graph_def, quantize):
-  """Set _output_types and _output_quantized for custom ops."""
-  for node in graph_def.node:
-    if node.op == "SequenceStringProjection":
-      node.attr["_output_quantized"].b = quantize
-      node.attr["_output_types"].list.type[:] = [tf_types.DT_FLOAT]
-      node.op = "SEQUENCE_STRING_PROJECTION"
-    elif node.op == "SequenceStringProjectionV2":
-      node.attr["_output_quantized"].b = quantize
-      node.attr["_output_types"].list.type[:] = [tf_types.DT_FLOAT]
-      node.op = "SEQUENCE_STRING_PROJECTION_V2"
-
-
-def export_frozen_graph_def(
-    session: tf.compat.v1.Session, model_config: Dict[str, Any],
-    input_tensors: Sequence[tf.Tensor],
-    output_tensors: Sequence[tf.Tensor]) -> tf.compat.v1.GraphDef:
-  """Returns a GraphDef object holding a processed network ready for exporting.
-
-  Args:
-    session: Active TensorFlow session containing the variables.
-    model_config: `ModelConfig` of the exported model.
-    input_tensors: A list of input tensors.
-    output_tensors: A list of output tensors.
-
-  Returns:
-    A frozen GraphDef object holding a processed network ready for exporting.
-  """
-  graph_def = session.graph_def
-
-  input_node_names = [tensor.op.name for tensor in input_tensors]
-  output_node_names = [tensor.op.name for tensor in output_tensors]
-  input_node_types = [tensor.dtype.as_datatype_enum for tensor in input_tensors]
-
-  graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
-      session, graph_def, output_node_names)
-  set_output_types_and_quantized(
-      graph_def, quantize=model_config["quantize"])
-
-  # Optimize the graph for inference by removing unused nodes. Also removes
-  # nodes related to training, which are not going to be used for inference.
-  graph_def = optimize_for_inference_lib.optimize_for_inference(
-      graph_def, input_node_names, output_node_names, input_node_types)
-
-  return graph_def
-
-
-def convert_frozen_graph_def_to_tflite(
-    graph_def: tf.compat.v1.GraphDef, model_config: Dict[str, Any],
-    input_tensors: Sequence[tf.Tensor],
-    output_tensors: Sequence[tf.Tensor]) -> bytes:
-  """Converts a TensorFlow GraphDef into a serialized TFLite Flatbuffer."""
-  converter = tf.lite.TFLiteConverter(graph_def, input_tensors, output_tensors)
-  if model_config["quantize"]:
-    converter.inference_type = tf.uint8
-    converter.inference_input_type = tf.uint8
-    converter.default_ranges_stats = (0., 1.)
-    converter.quantized_input_stats = {
-        tensor.op.name: (0., 1.) for tensor in input_tensors
-    }
-  # Custom ops 'PoolingOp' and 'SequenceStringProjection' are used.
-  converter.allow_custom_ops = True
-  converter.experimental_new_converter = False
-  return converter.convert()
-
-
-def export_tflite_model(model_config: Dict[str, Any], saved_model_dir: str,
-                        export_dir: str) -> None:
-  """Exports a saved_model into a tflite format."""
-  graph = tf.Graph()
-  with graph.as_default():
-    with tf.Session(graph=graph) as session:
-      metagraph_def = tf.compat.v1.saved_model.loader.load(
-          session, [tf.saved_model.tag_constants.SERVING], saved_model_dir)
-
-      serving_signature_key = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
-      signature_def = metagraph_def.signature_def[serving_signature_key]
-
-      def _get_tensors(tensor_infos):
-        tensor_names = [tensor_info.name for tensor_info in tensor_infos]
-        # Always use reverse lexicographic order for consistency and
-        # compatibility with PoD inference libraries.
-        tensor_names.sort(reverse=True)
-        return [graph.get_tensor_by_name(name) for name in tensor_names]
-
-      input_tensors = _get_tensors(signature_def.inputs.values())
-      output_tensors = _get_tensors(signature_def.outputs.values())
-
-      graph_def = export_frozen_graph_def(session, model_config, input_tensors,
-                                          output_tensors)
-      tflite_model = convert_frozen_graph_def_to_tflite(graph_def, model_config,
-                                                        input_tensors,
-                                                        output_tensors)
-
-      export_path = os.path.join(export_dir, "model.tflite")
-      with tf.gfile.GFile(export_path, "wb") as handle:
-        handle.write(tflite_model)
-      logging.info("TFLite model written to: %s", export_path)
-
-
-def main(_):
-  runner_config = load_runner_config()
-
-  if FLAGS.output_dir:
-    tf.gfile.MakeDirs(FLAGS.output_dir)
-
-  is_per_host = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V2
-  run_config = tf.estimator.tpu.RunConfig(
-      master=FLAGS.master,
-      model_dir=FLAGS.output_dir,
-      save_checkpoints_steps=runner_config["save_checkpoints_steps"],
-      keep_checkpoint_max=20,
-      tpu_config=tf.estimator.tpu.TPUConfig(
-          iterations_per_loop=runner_config["iterations_per_loop"],
-          num_shards=FLAGS.num_tpu_cores,
-          per_host_input_for_training=is_per_host))
-
-  model_fn = model_fn_builder(runner_config=runner_config)
-
-  # If TPU is not available, this will fall back to normal Estimator on CPU
-  # or GPU.
-  estimator = tf.estimator.tpu.TPUEstimator(
-      use_tpu=FLAGS.use_tpu,
-      model_fn=model_fn,
-      config=run_config,
-      train_batch_size=runner_config["batch_size"],
-      eval_batch_size=runner_config["batch_size"],
-      predict_batch_size=runner_config["batch_size"])
-
-  if FLAGS.runner_mode == "train":
-    train_input_fn = input_fn_reader.create_input_fn(
-        runner_config=runner_config,
-        create_projection=model.create_projection,
-        mode=tf.estimator.ModeKeys.TRAIN,
-        drop_remainder=True)
-    estimator.train(
-        input_fn=train_input_fn, max_steps=runner_config["train_steps"])
-
-  if FLAGS.runner_mode == "eval":
-
-    # TPU needs fixed shapes, so if the last batch is smaller, we drop it.
-    eval_input_fn = input_fn_reader.create_input_fn(
-        runner_config=runner_config,
-        create_projection=model.create_projection,
-        mode=tf.estimator.ModeKeys.EVAL,
-        drop_remainder=True)
-
-    for _ in tf.train.checkpoints_iterator(FLAGS.output_dir):
-      result = estimator.evaluate(input_fn=eval_input_fn)
-      for key in sorted(result):
-        logging.info("  %s = %s", key, str(result[key]))
-
-  if FLAGS.runner_mode == "export":
-    logging.info("Exporting the model...")
-
-    def serving_input_fn():
-      """Input function of the exported model."""
-
-      def _input_fn():
-        text = tf.placeholder(tf.string, shape=[1], name="Input")
-        projection, seq_length = model.create_projection(
-            model_config=runner_config["model_config"],
-            mode=tf.estimator.ModeKeys.PREDICT,
-            inputs=text)
-        features = {"projection": projection, "seq_length": seq_length}
-        return tf.estimator.export.ServingInputReceiver(
-            features=features, receiver_tensors=features)
-
-      return _input_fn
-
-    saved_model_dir = estimator.export_saved_model(FLAGS.output_dir,
-                                                   serving_input_fn())
-
-    export_tflite_model(runner_config["model_config"], saved_model_dir,
-                        FLAGS.output_dir)
-
-
-if __name__ == "__main__":
-  tf.app.run()

research/sequence_projection/requirements.txt

-numpy
-tensorflow
-tensorflow-text

research/sequence_projection/third_party/eigen.BUILD

-# Description:
-#   Eigen is a C++ template library for linear algebra: vectors,
-#   matrices, and related algorithms.
-
-licenses([
-    # Note: Eigen is an MPL2 library that includes GPL v3 and LGPL v2.1+ code.
-    #       We've taken special care to not reference any restricted code.
-    "reciprocal",  # MPL2
-    "notice",  # Portions BSD
-])
-
-exports_files(["COPYING.MPL2"])
-
-EIGEN_FILES = [
-    "Eigen/**",
-    "unsupported/Eigen/CXX11/**",
-    "unsupported/Eigen/FFT",
-    "unsupported/Eigen/KroneckerProduct",
-    "unsupported/Eigen/src/FFT/**",
-    "unsupported/Eigen/src/KroneckerProduct/**",
-    "unsupported/Eigen/MatrixFunctions",
-    "unsupported/Eigen/SpecialFunctions",
-    "unsupported/Eigen/src/MatrixFunctions/**",
-    "unsupported/Eigen/src/SpecialFunctions/**",
-]
-
-# Files known to be under MPL2 license.
-EIGEN_MPL2_HEADER_FILES = glob(
-    EIGEN_FILES,
-    exclude = [
-        # Guarantees that any non-MPL2 file added to the list above will fail to
-        # compile.
-        "Eigen/src/Core/util/NonMPL2.h",
-        "Eigen/**/CMakeLists.txt",
-    ],
-)
-
-cc_library(
-    name = "eigen",
-    hdrs = EIGEN_MPL2_HEADER_FILES,
-    defines = [
-        # This define (mostly) guarantees we don't link any problematic
-        # code. We use it, but we do not rely on it, as evidenced above.
-        "EIGEN_MPL2_ONLY",
-        "EIGEN_MAX_ALIGN_BYTES=64",
-        "EIGEN_HAS_TYPE_TRAITS=0",
-    ],
-    includes = ["."],
-    visibility = ["//visibility:public"],
-)
-
-filegroup(
-    name = "eigen_header_files",
-    srcs = EIGEN_MPL2_HEADER_FILES,
-    visibility = ["//visibility:public"],
-)