Merged commit includes the following changes: (#6726)

246873701  by menglong:

    Missing __init__.py under meta_architectures/

--
246857392  by menglong:

    Standardize proto namespace: lstm_object_detection.protos

--
246625127  by menglong:

    Internal changes.

--
246596481  by menglong:

    Add License

--
246580605  by menglong:

    Internal changes

--
246344626  by menglong:

    Open source interleaved mobilenet v2 model.

--
244893883  by menglong:

    Introduce multi_input_decoder for interleaved model.

--
244461016  by menglong:

    Add pre-bottleneck operation to lstm cells to support interleaved model.

--
244052176  by menglong:

    Update README

--
244020495  by menglong:

    Add test to rnn_decoder.

--
243704250  by menglong:

    Duplicate assignment.

--
243091836  by menglong:

    Move LSTMSSD meta arch into separate folder

--
242900337  by menglong:

    Modified mobilenet definition for LSTM-SSD

--
242...

research/lstm_object_detection/README

-Tensorflow mobile video object detection implementation proposed in the following paper:
-Mobile Video Object Detection with Temporally-Aware Feature Maps (CVPR 2018).
-
-http://openaccess.thecvf.com/content_cvpr_2018/papers/Liu_Mobile_Video_Object_CVPR_2018_paper.pdf
-
-@article{liu2017mobile,
-  title={Mobile Video Object Detection with Temporally-Aware Feature Maps},
-  author={Liu, Mason and Zhu, Menglong},
-  journal={CVPR},
-  year={2018}
-}
-
-If you have any questions regarding this codebase, please contact us:
-masonliuw@gmail.com
-yinxiao@google.com
-menglong@google.com
\ No newline at end of file

research/lstm_object_detection/README.md

+# Tensorflow Mobile Video Object Detection
+
+Tensorflow mobile video object detection implementation proposed in the
+following papers:
+
+<p align="center">
+  <img src="g3doc/lstm_ssd_intro.png" width=640 height=360>
+</p>
+
+```
+"Mobile Video Object Detection with Temporally-Aware Feature Maps",
+Liu, Mason and Zhu, Menglong, CVPR 2018.
+```
+\[[link](http://openaccess.thecvf.com/content_cvpr_2018/papers/Liu_Mobile_Video_Object_CVPR_2018_paper.pdf)\]\[[bibtex](
+https://scholar.googleusercontent.com/scholar.bib?q=info:hq5rcMUUXysJ:scholar.google.com/&output=citation&scisig=AAGBfm0AAAAAXLdwXcU5g_wiMQ40EvbHQ9kTyvfUxffh&scisf=4&ct=citation&cd=-1&hl=en)\]
+
+
+<p align="center">
+  <img src="g3doc/Interleaved_Intro.png" width=480 height=360>
+</p>
+
+```
+"Looking Fast and Slow: Memory-Guided Mobile Video Object Detection",
+Liu, Mason and Zhu, Menglong and White, Marie and Li, Yinxiao and Kalenichenko, Dmitry
+```
+\[[link](https://arxiv.org/abs/1903.10172)\]\[[bibtex](
+https://scholar.googleusercontent.com/scholar.bib?q=info:rLqvkztmWYgJ:scholar.google.com/&output=citation&scisig=AAGBfm0AAAAAXLdwNf-LJlm2M1ymQHbq2wYA995MHpJu&scisf=4&ct=citation&cd=-1&hl=en)\]
+
+
+## Maintainers
+* masonliuw@gmail.com
+* yinxiao@google.com
+* menglong@google.com

research/lstm_object_detection/builders/graph_rewriter_builder.py

+# Copyright 2019 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.
+# ==============================================================================
+
+"""Custom version for quantized training and evaluation functions.
+
+The main difference between this and the third_party graph_rewriter_builder.py
+is that this version uses experimental_create_training_graph which allows the
+customization of freeze_bn_delay.
+"""
+
+import re
+import tensorflow as tf
+from tensorflow.contrib.quantize.python import common
+from tensorflow.contrib.quantize.python import input_to_ops
+from tensorflow.contrib.quantize.python import quant_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+
+
+def build(graph_rewriter_config,
+          quant_overrides_config=None,
+          is_training=True,
+          is_export=False):
+  """Returns a function that modifies default graph based on options.
+
+  Args:
+    graph_rewriter_config: graph_rewriter_pb2.GraphRewriter proto.
+    quant_overrides_config: quant_overrides_pb2.QuantOverrides proto.
+    is_training: whether in training or eval mode.
+    is_export: whether exporting the graph.
+  """
+  def graph_rewrite_fn():
+    """Function to quantize weights and activation of the default graph."""
+    if (graph_rewriter_config.quantization.weight_bits != 8 or
+        graph_rewriter_config.quantization.activation_bits != 8):
+      raise ValueError('Only 8bit quantization is supported')
+
+    graph = tf.get_default_graph()
+
+    # Insert custom quant ops.
+    if quant_overrides_config is not None:
+      input_to_ops_map = input_to_ops.InputToOps(graph)
+      for q in quant_overrides_config.quant_configs:
+        producer = graph.get_operation_by_name(q.op_name)
+        if producer is None:
+          raise ValueError('Op name does not exist in graph.')
+        context = _get_context_from_op(producer)
+        consumers = input_to_ops_map.ConsumerOperations(producer)
+        if q.fixed_range:
+          _insert_fixed_quant_op(
+              context,
+              q.quant_op_name,
+              producer,
+              consumers,
+              init_min=q.min,
+              init_max=q.max,
+              quant_delay=q.delay if is_training else 0)
+        else:
+          raise ValueError('Learned ranges are not yet supported.')
+
+    # Quantize the graph by inserting quantize ops for weights and activations
+    if is_training:
+      tf.contrib.quantize.experimental_create_training_graph(
+          input_graph=graph,
+          quant_delay=graph_rewriter_config.quantization.delay,
+          freeze_bn_delay=graph_rewriter_config.quantization.delay)
+    else:
+      tf.contrib.quantize.experimental_create_eval_graph(
+          input_graph=graph,
+          quant_delay=graph_rewriter_config.quantization.delay
+          if not is_export else 0)
+
+    tf.contrib.layers.summarize_collection('quant_vars')
+  return graph_rewrite_fn
+
+
+def _get_context_from_op(op):
+  """Gets the root context name from the op name."""
+  context_re = re.search(r'^(.*)/([^/]+)', op.name)
+  if context_re:
+    return context_re.group(1)
+  return ''
+
+
+def _insert_fixed_quant_op(context,
+                           name,
+                           producer,
+                           consumers,
+                           init_min=-6.0,
+                           init_max=6.0,
+                           quant_delay=None):
+  """Adds a fake quant op with fixed ranges.
+
+  Args:
+    context: The parent scope of the op to be quantized.
+    name: The name of the fake quant op.
+    producer: The producer op to be quantized.
+    consumers: The consumer ops to the producer op.
+    init_min: The minimum range for the fake quant op.
+    init_max: The maximum range for the fake quant op.
+    quant_delay: Number of steps to wait before activating the fake quant op.
+
+  Raises:
+    ValueError: When producer operation is not directly connected to the
+      consumer operation.
+  """
+  name_prefix = name if not context else context + '/' + name
+  inputs = producer.outputs[0]
+  quant = quant_ops.FixedQuantize(
+      inputs, init_min=init_min, init_max=init_max, scope=name_prefix)
+
+  if quant_delay and quant_delay > 0:
+    activate_quant = math_ops.greater_equal(
+        common.CreateOrGetQuantizationStep(),
+        quant_delay,
+        name=name_prefix + '/activate_quant')
+    quant = control_flow_ops.cond(
+        activate_quant,
+        lambda: quant,
+        lambda: inputs,
+        name=name_prefix + '/delayed_quant')
+
+  if consumers:
+    tensors_modified_count = common.RerouteTensor(
+        quant, inputs, can_modify=consumers)
+    # Some operations can have multiple output tensors going to the same
+    # consumer. Since consumers is a set, we need to ensure that
+    # tensors_modified_count is greater than or equal to the length of the set
+    # of consumers.
+    if tensors_modified_count < len(consumers):
+      raise ValueError('No inputs quantized for ops: [%s]' % ', '.join(
+          [consumer.name for consumer in consumers]))

research/lstm_object_detection/builders/graph_rewriter_builder_test.py

+# Copyright 2019 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.
+# ==============================================================================
+
+"""Tests for graph_rewriter_builder."""
+import mock
+import tensorflow as tf
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from lstm_object_detection.builders import graph_rewriter_builder
+from lstm_object_detection.protos import quant_overrides_pb2
+from object_detection.protos import graph_rewriter_pb2
+
+
+class QuantizationBuilderTest(tf.test.TestCase):
+
+  def testQuantizationBuilderSetsUpCorrectTrainArguments(self):
+    with mock.patch.object(
+        tf.contrib.quantize,
+        'experimental_create_training_graph') as mock_quant_fn:
+      with mock.patch.object(tf.contrib.layers,
+                             'summarize_collection') as mock_summarize_col:
+        graph_rewriter_proto = graph_rewriter_pb2.GraphRewriter()
+        graph_rewriter_proto.quantization.delay = 10
+        graph_rewriter_proto.quantization.weight_bits = 8
+        graph_rewriter_proto.quantization.activation_bits = 8
+        graph_rewrite_fn = graph_rewriter_builder.build(
+            graph_rewriter_proto, is_training=True)
+        graph_rewrite_fn()
+        _, kwargs = mock_quant_fn.call_args
+        self.assertEqual(kwargs['input_graph'], tf.get_default_graph())
+        self.assertEqual(kwargs['quant_delay'], 10)
+        mock_summarize_col.assert_called_with('quant_vars')
+
+  def testQuantizationBuilderSetsUpCorrectEvalArguments(self):
+    with mock.patch.object(tf.contrib.quantize,
+                           'experimental_create_eval_graph') as mock_quant_fn:
+      with mock.patch.object(tf.contrib.layers,
+                             'summarize_collection') as mock_summarize_col:
+        graph_rewriter_proto = graph_rewriter_pb2.GraphRewriter()
+        graph_rewriter_proto.quantization.delay = 10
+        graph_rewrite_fn = graph_rewriter_builder.build(
+            graph_rewriter_proto, is_training=False)
+        graph_rewrite_fn()
+        _, kwargs = mock_quant_fn.call_args
+        self.assertEqual(kwargs['input_graph'], tf.get_default_graph())
+        mock_summarize_col.assert_called_with('quant_vars')
+
+  def testQuantizationBuilderAddsQuantOverride(self):
+    graph = ops.Graph()
+    with graph.as_default():
+      self._buildGraph()
+
+      quant_overrides_proto = quant_overrides_pb2.QuantOverrides()
+      quant_config = quant_overrides_proto.quant_configs.add()
+      quant_config.op_name = 'test_graph/add_ab'
+      quant_config.quant_op_name = 'act_quant'
+      quant_config.fixed_range = True
+      quant_config.min = 0
+      quant_config.max = 6
+      quant_config.delay = 100
+
+      graph_rewriter_proto = graph_rewriter_pb2.GraphRewriter()
+      graph_rewriter_proto.quantization.delay = 10
+      graph_rewriter_proto.quantization.weight_bits = 8
+      graph_rewriter_proto.quantization.activation_bits = 8
+
+      graph_rewrite_fn = graph_rewriter_builder.build(
+          graph_rewriter_proto,
+          quant_overrides_config=quant_overrides_proto,
+          is_training=True)
+      graph_rewrite_fn()
+
+      act_quant_found = False
+      quant_delay_found = False
+      for op in graph.get_operations():
+        if (quant_config.quant_op_name in op.name and
+            op.type == 'FakeQuantWithMinMaxArgs'):
+          act_quant_found = True
+          min_val = op.get_attr('min')
+          max_val = op.get_attr('max')
+          self.assertEqual(min_val, quant_config.min)
+          self.assertEqual(max_val, quant_config.max)
+        if ('activate_quant' in op.name and
+            quant_config.quant_op_name in op.name and op.type == 'Const'):
+          tensor = op.get_attr('value')
+          if tensor.int64_val[0] == quant_config.delay:
+            quant_delay_found = True
+
+      self.assertTrue(act_quant_found)
+      self.assertTrue(quant_delay_found)
+
+  def _buildGraph(self, scope='test_graph'):
+    with ops.name_scope(scope):
+      a = tf.constant(10, dtype=dtypes.float32, name='input_a')
+      b = tf.constant(20, dtype=dtypes.float32, name='input_b')
+      ab = tf.add(a, b, name='add_ab')
+      c = tf.constant(30, dtype=dtypes.float32, name='input_c')
+      abc = tf.multiply(ab, c, name='mul_ab_c')
+      return abc
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/lstm_object_detection/configs/lstm_ssd_mobilenet_v1_imagenet.config

@@ -22,7 +22,7 @@
 
 model {
   ssd {
-    num_classes: 30
+    num_classes: 30  # Num of class for imagenet vid dataset.
     box_coder {
       faster_rcnn_box_coder {
         y_scale: 10.0
@@ -197,9 +197,9 @@ train_input_reader: {
   min_after_dequeue: 4
   label_map_path: "path/to/label_map"
   external_input_reader {
-    [lstm_object_detection.input_readers.GoogleInputReader.google_input_reader] {
+    [lstm_object_detection.protos.GoogleInputReader.google_input_reader] {
       tf_record_video_input_reader: {
-        input_path: "your/cns/path"
+        input_path: "path/to/sequence_example/data"
         data_type: TF_SEQUENCE_EXAMPLE
         video_length: 4
       }
@@ -208,7 +208,7 @@ train_input_reader: {
 }
 
 eval_config: {
-  metrics_set: "coco_evaluation_last_frame"
+  metrics_set: "coco_evaluation_all_frames"
   use_moving_averages: true
   min_score_threshold: 0.5
   max_num_boxes_to_visualize: 300
@@ -219,9 +219,9 @@ eval_config: {
 eval_input_reader: {
   label_map_path: "path/to/label_map"
   external_input_reader {
-    [lstm_object_detection.input_readers.GoogleInputReader.google_input_reader] {
+    [lstm_object_detection.protos.GoogleInputReader.google_input_reader] {
       tf_record_video_input_reader: {
-        input_path: "your/cns/path"
+        input_path: "path/to/sequence_example/data"
         data_type: TF_SEQUENCE_EXAMPLE
         video_length: 4
       }

research/lstm_object_detection/evaluator.py

@@ -20,7 +20,6 @@ DetectionModel.
 
 """
 
-import logging
 import tensorflow as tf
 from lstm_object_detection.metrics import coco_evaluation_all_frames
 from object_detection import eval_util
@@ -215,7 +214,7 @@ def evaluate(create_input_dict_fn,
   model = create_model_fn()
 
   if eval_config.ignore_groundtruth and not eval_config.export_path:
-    logging.fatal('If ignore_groundtruth=True then an export_path is '
+    tf.logging.fatal('If ignore_groundtruth=True then an export_path is '
                   'required. Aborting!!!')
 
   tensor_dicts = _extract_prediction_tensors(
@@ -252,14 +251,14 @@ def evaluate(create_input_dict_fn,
         third_party eval_util.py.
     """
     if batch_index % 10 == 0:
-      logging.info('Running eval ops batch %d', batch_index)
+      tf.logging.info('Running eval ops batch %d', batch_index)
     if not losses_dict:
       losses_dict = {}
     try:
       result_dicts, result_losses_dict = sess.run([tensor_dicts, losses_dict])
       counters['success'] += 1
     except tf.errors.InvalidArgumentError:
-      logging.info('Skipping image')
+      tf.logging.info('Skipping image')
       counters['skipped'] += 1
       return {}
     num_images = len(tensor_dicts)

research/lstm_object_detection/inputs/seq_dataset_builder.py

@@ -23,7 +23,6 @@ Detection configuration framework, they should define their own builder function
 that wraps the build function.
 """
 import tensorflow as tf
-import tensorflow.google as google_tf
 from tensorflow.contrib.training.python.training import sequence_queueing_state_saver as sqss
 from lstm_object_detection.inputs import tf_sequence_example_decoder
 from lstm_object_detection.protos import input_reader_google_pb2
@@ -116,12 +115,12 @@ def build(input_reader_config,
                      'input_reader_pb2.InputReader.')
 
   external_reader_config = input_reader_config.external_input_reader
-  google_input_reader_config = external_reader_config.Extensions[
+  external_input_reader_config = external_reader_config.Extensions[
       input_reader_google_pb2.GoogleInputReader.google_input_reader]
-  input_reader_type = google_input_reader_config.WhichOneof('input_reader')
+  input_reader_type = external_input_reader_config.WhichOneof('input_reader')
 
   if input_reader_type == 'tf_record_video_input_reader':
-    config = google_input_reader_config.tf_record_video_input_reader
+    config = external_input_reader_config.tf_record_video_input_reader
     reader_type_class = tf.TFRecordReader
   else:
     raise ValueError(

research/lstm_object_detection/inputs/seq_dataset_builder_test.py

@@ -20,7 +20,6 @@ import numpy as np
 import tensorflow as tf
 
 from google.protobuf import text_format
-from google3.testing.pybase import parameterized
 from tensorflow.core.example import example_pb2
 from tensorflow.core.example import feature_pb2
 from lstm_object_detection.inputs import seq_dataset_builder
@@ -32,7 +31,7 @@ from object_detection.protos import pipeline_pb2
 from object_detection.protos import preprocessor_pb2
 
 
-class DatasetBuilderTest(parameterized.TestCase):
+class DatasetBuilderTest(tf.test.TestCase):
 
   def _create_tf_record(self):
     path = os.path.join(self.get_temp_dir(), 'tfrecord')
@@ -104,7 +103,7 @@ class DatasetBuilderTest(parameterized.TestCase):
     """
 
     model_text_proto = """
-    [object_detection.protos.lstm_model] {
+    [lstm_object_detection.protos.lstm_model] {
       train_unroll_length: 4
       eval_unroll_length: 4
     }
@@ -211,7 +210,7 @@ class DatasetBuilderTest(parameterized.TestCase):
   def _get_input_proto(self, input_reader):
     return """
         external_input_reader {
-          [lstm_object_detection.input_readers.GoogleInputReader.google_input_reader] {
+          [lstm_object_detection.protos.GoogleInputReader.google_input_reader] {
             %s: {
               input_path: '{0}'
               data_type: TF_SEQUENCE_EXAMPLE
@@ -221,11 +220,11 @@ class DatasetBuilderTest(parameterized.TestCase):
         }
       """ % input_reader
 
-  @parameterized.named_parameters(('tf_record', 'tf_record_video_input_reader'))
-  def test_video_input_reader(self, video_input_type):
+  def test_video_input_reader(self):
     input_reader_proto = input_reader_pb2.InputReader()
     text_format.Merge(
-        self._get_input_proto(video_input_type), input_reader_proto)
+        self._get_input_proto('tf_record_video_input_reader'),
+        input_reader_proto)
 
     configs = self._get_model_configs_from_proto()
     tensor_dict = seq_dataset_builder.build(

research/lstm_object_detection/inputs/tf_sequence_example_decoder.py

@@ -17,8 +17,6 @@
 
 A decoder to decode string tensors containing serialized
 tensorflow.SequenceExample protos.
-TODO(yinxiao): When TensorFlow object detection API officially supports
-tensorflow.SequenceExample, merge this decoder.
 """
 import tensorflow as tf
 from object_detection.core import data_decoder

research/lstm_object_detection/lstm/lstm_cells.py

@@ -17,6 +17,7 @@
 
 import tensorflow as tf
 from tensorflow.contrib.framework.python.ops import variables
+import lstm_object_detection.lstm.utils as lstm_utils
 
 slim = tf.contrib.slim
 
@@ -44,10 +45,11 @@ class BottleneckConvLSTMCell(tf.contrib.rnn.RNNCell):
                num_units,
                forget_bias=1.0,
                activation=tf.tanh,
-               flattened_state=False,
+               flatten_state=False,
                clip_state=False,
                output_bottleneck=False,
-               visualize_gates=True):
+               pre_bottleneck=False,
+               visualize_gates=False):
     """Initializes the basic LSTM cell.
 
     Args:
@@ -56,11 +58,13 @@ class BottleneckConvLSTMCell(tf.contrib.rnn.RNNCell):
       num_units: int, The number of channels in the LSTM cell.
       forget_bias: float, The bias added to forget gates (see above).
       activation: Activation function of the inner states.
-      flattened_state: if True, state tensor will be flattened and stored as
+      flatten_state: if True, state tensor will be flattened and stored as
         a 2-d tensor. Use for exporting the model to tfmini.
       clip_state: if True, clip state between [-6, 6].
       output_bottleneck: if True, the cell bottleneck will be concatenated
         to the cell output.
+      pre_bottleneck: if True, cell assumes that bottlenecking was performing
+        before the function was called.
       visualize_gates: if True, add histogram summaries of all gates
         and outputs to tensorboard.
     """
@@ -70,9 +74,10 @@ class BottleneckConvLSTMCell(tf.contrib.rnn.RNNCell):
     self._forget_bias = forget_bias
     self._activation = activation
     self._viz_gates = visualize_gates
-    self._flattened_state = flattened_state
+    self._flatten_state = flatten_state
     self._clip_state = clip_state
     self._output_bottleneck = output_bottleneck
+    self._pre_bottleneck = pre_bottleneck
     self._param_count = self._num_units
     for dim in self._output_size:
       self._param_count *= dim
@@ -103,18 +108,20 @@ class BottleneckConvLSTMCell(tf.contrib.rnn.RNNCell):
       a LSTMStateTuple of the state at the current timestep.
     """
     scope = scope or 'conv_lstm_cell'
-    with tf.variable_scope(scope):
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
       c, h = state
 
       # unflatten state if necessary
-      if self._flattened_state:
+      if self._flatten_state:
         c = tf.reshape(c, [-1] + self.output_size)
         h = tf.reshape(h, [-1] + self.output_size)
 
       # summary of input passed into cell
       if self._viz_gates:
         slim.summaries.add_histogram_summary(inputs, 'cell_input')
-
+      if self._pre_bottleneck:
+        bottleneck = inputs
+      else:
         bottleneck = tf.contrib.layers.separable_conv2d(
             tf.concat([inputs, h], 3),
             self._num_units,
@@ -154,7 +161,7 @@ class BottleneckConvLSTMCell(tf.contrib.rnn.RNNCell):
         output = tf.concat([new_h, bottleneck], axis=3)
 
       # reflatten state to store it
-      if self._flattened_state:
+      if self._flatten_state:
         new_c = tf.reshape(new_c, [-1, self._param_count])
         new_h = tf.reshape(new_h, [-1, self._param_count])
 
@@ -174,7 +181,7 @@ class BottleneckConvLSTMCell(tf.contrib.rnn.RNNCell):
       The created initial state.
     """
     state_size = (
-        self.state_size_flat if self._flattened_state else self.state_size)
+        self.state_size_flat if self._flatten_state else self.state_size)
     # list of 2 zero tensors or variables tensors, depending on if
     # learned_state is true
     ret_flat = [(variables.model_variable(
@@ -197,3 +204,541 @@ class BottleneckConvLSTMCell(tf.contrib.rnn.RNNCell):
       r.set_shape([None] + s)
     return tf.contrib.framework.nest.pack_sequence_as(
         structure=[1, 1], flat_sequence=ret_flat)
+
+  def pre_bottleneck(self, inputs, state, input_index):
+    """Apply pre-bottleneck projection to inputs.
+
+    Pre-bottleneck operation maps features of different channels into the same
+    dimension. The purpose of this op is to share the features from both large
+    and small models in the same LSTM cell.
+
+    Args:
+      inputs: 4D Tensor with shape [batch_size x width x height x input_size].
+      state: 4D Tensor with shape [batch_size x width x height x state_size].
+      input_index: integer index indicating which base features the inputs
+        correspoding to.
+    Returns:
+      inputs: pre-bottlenecked inputs.
+    Raises:
+      ValueError: If pre_bottleneck is not set or inputs is not rank 4.
+    """
+    # Sometimes state is a tuple, in which case it cannot be modified, e.g.
+    # during training, tf.contrib.training.SequenceQueueingStateSaver
+    # returns the state as a tuple. This should not be an issue since we
+    # only need to modify state[1] during export, when state should be a
+    # list.
+    if not len(inputs.shape) == 4:
+      raise ValueError('Expect rank 4 feature tensor.')
+    if not self._flatten_state and not len(state.shape) == 4:
+      raise ValueError('Expect rank 4 state tensor.')
+    if self._flatten_state and not len(state.shape) == 2:
+      raise ValueError('Expect rank 2 state tensor when flatten_state is set.')
+    with tf.name_scope(None):
+      state = tf.identity(state, name='raw_inputs/init_lstm_h')
+    if self._flatten_state:
+      batch_size = inputs.shape[0]
+      height = inputs.shape[1]
+      width = inputs.shape[2]
+      state = tf.reshape(state, [batch_size, height, width, -1])
+    with tf.variable_scope('conv_lstm_cell', reuse=tf.AUTO_REUSE):
+      scope_name = 'bottleneck_%d' % input_index
+      inputs = tf.contrib.layers.separable_conv2d(
+          tf.concat([inputs, state], 3),
+          self.output_size[-1],
+          self._filter_size,
+          depth_multiplier=1,
+          activation_fn=tf.nn.relu6,
+          normalizer_fn=None,
+          scope=scope_name)
+      # For exporting inference graph, we only mark the first timestep.
+      with tf.name_scope(None):
+        inputs = tf.identity(
+            inputs, name='raw_outputs/base_endpoint_%d' % (input_index + 1))
+    return inputs
+
+
+class GroupedConvLSTMCell(tf.contrib.rnn.RNNCell):
+  """Basic LSTM recurrent network cell using separable convolutions.
+
+  The implementation is based on: https://arxiv.org/abs/1903.10172.
+
+  We add forget_bias (default: 1) to the biases of the forget gate in order to
+  reduce the scale of forgetting in the beginning of the training.
+
+  This LSTM first projects inputs to the size of the output before doing gate
+  computations. This saves params unless the input is less than a third of the
+  state size channel-wise. Computation of bottlenecks and gates is divided
+  into independent groups for further savings.
+  """
+
+  def __init__(self,
+               filter_size,
+               output_size,
+               num_units,
+               is_training,
+               forget_bias=1.0,
+               activation=tf.tanh,
+               use_batch_norm=False,
+               flatten_state=False,
+               groups=4,
+               clip_state=False,
+               scale_state=False,
+               output_bottleneck=False,
+               pre_bottleneck=False,
+               is_quantized=False,
+               visualize_gates=False):
+    """Initialize the basic LSTM cell.
+
+    Args:
+      filter_size: collection, conv filter size
+      output_size: collection, the width/height dimensions of the cell/output
+      num_units: int, The number of channels in the LSTM cell.
+      is_training: Whether the LSTM is in training mode.
+      forget_bias: float, The bias added to forget gates (see above).
+      activation: Activation function of the inner states.
+      use_batch_norm: if True, use batch norm after convolution
+      flatten_state: if True, state tensor will be flattened and stored as
+        a 2-d tensor. Use for exporting the model to tfmini
+      groups: Number of groups to split the state into. Must evenly divide
+        num_units.
+      clip_state: if True, clips state between [-6, 6].
+      scale_state: if True, scales state so that all values are under 6 at all
+        times.
+      output_bottleneck: if True, the cell bottleneck will be concatenated
+        to the cell output.
+      pre_bottleneck: if True, cell assumes that bottlenecking was performing
+        before the function was called.
+      is_quantized: if True, the model is in quantize mode, which requires
+        quantization friendly concat and separable_conv2d ops.
+      visualize_gates: if True, add histogram summaries of all gates
+        and outputs to tensorboard
+
+    Raises:
+      ValueError: when both clip_state and scale_state are enabled.
+    """
+    if clip_state and scale_state:
+      raise ValueError('clip_state and scale_state cannot both be enabled.')
+
+    self._filter_size = list(filter_size)
+    self._output_size = list(output_size)
+    self._num_units = num_units
+    self._is_training = is_training
+    self._forget_bias = forget_bias
+    self._activation = activation
+    self._use_batch_norm = use_batch_norm
+    self._viz_gates = visualize_gates
+    self._flatten_state = flatten_state
+    self._param_count = self._num_units
+    self._groups = groups
+    self._scale_state = scale_state
+    self._clip_state = clip_state
+    self._output_bottleneck = output_bottleneck
+    self._pre_bottleneck = pre_bottleneck
+    self._is_quantized = is_quantized
+    for dim in self._output_size:
+      self._param_count *= dim
+
+  @property
+  def state_size(self):
+    return tf.contrib.rnn.LSTMStateTuple(self._output_size + [self._num_units],
+                                         self._output_size + [self._num_units])
+
+  @property
+  def state_size_flat(self):
+    return tf.contrib.rnn.LSTMStateTuple([self._param_count],
+                                         [self._param_count])
+
+  @property
+  def output_size(self):
+    return self._output_size + [self._num_units]
+
+  @property
+  def filter_size(self):
+    return self._filter_size
+
+  @property
+  def num_groups(self):
+    return self._groups
+
+  def __call__(self, inputs, state, scope=None):
+    """Long short-term memory cell (LSTM) with bottlenecking.
+
+    Includes logic for quantization-aware training. Note that all concats and
+    activations use fixed ranges unless stated otherwise.
+
+    Args:
+      inputs: Input tensor at the current timestep.
+      state: Tuple of tensors, the state at the previous timestep.
+      scope: Optional scope.
+    Returns:
+      A tuple where the first element is the LSTM output and the second is
+      a LSTMStateTuple of the state at the current timestep.
+    """
+    scope = scope or 'conv_lstm_cell'
+    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
+      c, h = state
+
+      # Set nodes to be under raw_inputs/ name scope for tfmini export.
+      with tf.name_scope(None):
+        c = tf.identity(c, name='raw_inputs/init_lstm_c')
+        # When pre_bottleneck is enabled, input h handle is in rnn_decoder.py
+        if not self._pre_bottleneck:
+          h = tf.identity(h, name='raw_inputs/init_lstm_h')
+
+      # unflatten state if necessary
+      if self._flatten_state:
+        c = tf.reshape(c, [-1] + self.output_size)
+        h = tf.reshape(h, [-1] + self.output_size)
+
+      c_list = tf.split(c, self._groups, axis=3)
+      if self._pre_bottleneck:
+        inputs_list = tf.split(inputs, self._groups, axis=3)
+      else:
+        h_list = tf.split(h, self._groups, axis=3)
+      out_bottleneck = []
+      out_c = []
+      out_h = []
+      # summary of input passed into cell
+      if self._viz_gates:
+        slim.summaries.add_histogram_summary(inputs, 'cell_input')
+
+      for k in range(self._groups):
+        if self._pre_bottleneck:
+          bottleneck = inputs_list[k]
+        else:
+          if self._use_batch_norm:
+            b_x = lstm_utils.quantizable_separable_conv2d(
+                inputs,
+                self._num_units / self._groups,
+                self._filter_size,
+                is_quantized=self._is_quantized,
+                depth_multiplier=1,
+                activation_fn=None,
+                normalizer_fn=None,
+                scope='bottleneck_%d_x' % k)
+            b_h = lstm_utils.quantizable_separable_conv2d(
+                h_list[k],
+                self._num_units / self._groups,
+                self._filter_size,
+                is_quantized=self._is_quantized,
+                depth_multiplier=1,
+                activation_fn=None,
+                normalizer_fn=None,
+                scope='bottleneck_%d_h' % k)
+            b_x = slim.batch_norm(
+                b_x, scale=True, is_training=self._is_training,
+                scope='BatchNorm_%d_X' % k)
+            b_h = slim.batch_norm(
+                b_h, scale=True, is_training=self._is_training,
+                scope='BatchNorm_%d_H' % k)
+            bottleneck = b_x + b_h
+          else:
+            # All concats use fixed quantization ranges to prevent rescaling
+            # at inference. Both |inputs| and |h_list| are tensors resulting
+            # from Relu6 operations so we fix the ranges to [0, 6].
+            bottleneck_concat = lstm_utils.quantizable_concat(
+                [inputs, h_list[k]],
+                axis=3,
+                is_training=False,
+                is_quantized=self._is_quantized,
+                scope='bottleneck_%d/quantized_concat' % k)
+
+            bottleneck = lstm_utils.quantizable_separable_conv2d(
+                bottleneck_concat,
+                self._num_units / self._groups,
+                self._filter_size,
+                is_quantized=self._is_quantized,
+                depth_multiplier=1,
+                activation_fn=self._activation,
+                normalizer_fn=None,
+                scope='bottleneck_%d' % k)
+
+        concat = lstm_utils.quantizable_separable_conv2d(
+            bottleneck,
+            4 * self._num_units / self._groups,
+            self._filter_size,
+            is_quantized=self._is_quantized,
+            depth_multiplier=1,
+            activation_fn=None,
+            normalizer_fn=None,
+            scope='concat_conv_%d' % k)
+
+        # Since there is no activation in the previous separable conv, we
+        # quantize here. A starting range of [-6, 6] is used because the
+        # tensors are input to a Sigmoid function that saturates at these
+        # ranges.
+        concat = lstm_utils.quantize_op(
+            concat,
+            is_training=self._is_training,
+            default_min=-6,
+            default_max=6,
+            is_quantized=self._is_quantized,
+            scope='gates_%d/act_quant' % k)
+
+        # i = input_gate, j = new_input, f = forget_gate, o = output_gate
+        i, j, f, o = tf.split(concat, 4, 3)
+
+        f_add = f + self._forget_bias
+        f_add = lstm_utils.quantize_op(
+            f_add,
+            is_training=self._is_training,
+            default_min=-6,
+            default_max=6,
+            is_quantized=self._is_quantized,
+            scope='forget_gate_%d/add_quant' % k)
+        f_act = tf.sigmoid(f_add)
+        # The quantization range is fixed for the sigmoid to ensure that zero
+        # is exactly representable.
+        f_act = lstm_utils.quantize_op(
+            f_act,
+            is_training=False,
+            default_min=0,
+            default_max=1,
+            is_quantized=self._is_quantized,
+            scope='forget_gate_%d/act_quant' % k)
+
+        a = c_list[k] * f_act
+        a = lstm_utils.quantize_op(
+            a,
+            is_training=self._is_training,
+            is_quantized=self._is_quantized,
+            scope='forget_gate_%d/mul_quant' % k)
+
+        i_act = tf.sigmoid(i)
+        # The quantization range is fixed for the sigmoid to ensure that zero
+        # is exactly representable.
+        i_act = lstm_utils.quantize_op(
+            i_act,
+            is_training=False,
+            default_min=0,
+            default_max=1,
+            is_quantized=self._is_quantized,
+            scope='input_gate_%d/act_quant' % k)
+
+        j_act = self._activation(j)
+        # The quantization range is fixed for the relu6 to ensure that zero
+        # is exactly representable.
+        j_act = lstm_utils.quantize_op(
+            j_act,
+            is_training=False,
+            default_min=0,
+            default_max=6,
+            is_quantized=self._is_quantized,
+            scope='new_input_%d/act_quant' % k)
+
+        b = i_act * j_act
+        b = lstm_utils.quantize_op(
+            b,
+            is_training=self._is_training,
+            is_quantized=self._is_quantized,
+            scope='input_gate_%d/mul_quant' % k)
+
+        new_c = a + b
+        # The quantization range is fixed to [0, 6] due to an optimization in
+        # TFLite. The order of operations is as fllows:
+        #     Add -> FakeQuant -> Relu6 -> FakeQuant -> Concat.
+        # The fakequant ranges to the concat must be fixed to ensure all inputs
+        # to the concat have the same range, removing the need for rescaling.
+        # The quantization ranges input to the relu6 are propagated to its
+        # output. Any mismatch between these two ranges will cause an error.
+        new_c = lstm_utils.quantize_op(
+            new_c,
+            is_training=False,
+            default_min=0,
+            default_max=6,
+            is_quantized=self._is_quantized,
+            scope='new_c_%d/add_quant' % k)
+
+        if not self._is_quantized:
+          if self._scale_state:
+            normalizer = tf.maximum(1.0,
+                                    tf.reduce_max(new_c, axis=(1, 2, 3)) / 6)
+            new_c /= tf.reshape(normalizer, [tf.shape(new_c)[0], 1, 1, 1])
+          elif self._clip_state:
+            new_c = tf.clip_by_value(new_c, -6, 6)
+
+        new_c_act = self._activation(new_c)
+        # The quantization range is fixed for the relu6 to ensure that zero
+        # is exactly representable.
+        new_c_act = lstm_utils.quantize_op(
+            new_c_act,
+            is_training=False,
+            default_min=0,
+            default_max=6,
+            is_quantized=self._is_quantized,
+            scope='new_c_%d/act_quant' % k)
+
+        o_act = tf.sigmoid(o)
+        # The quantization range is fixed for the sigmoid to ensure that zero
+        # is exactly representable.
+        o_act = lstm_utils.quantize_op(
+            o_act,
+            is_training=False,
+            default_min=0,
+            default_max=1,
+            is_quantized=self._is_quantized,
+            scope='output_%d/act_quant' % k)
+
+        new_h = new_c_act * o_act
+        # The quantization range is fixed since it is input to a concat.
+        # A range of [0, 6] is used since |new_h| is a product of ranges [0, 6]
+        # and [0, 1].
+        new_h_act = lstm_utils.quantize_op(
+            new_h,
+            is_training=False,
+            default_min=0,
+            default_max=6,
+            is_quantized=self._is_quantized,
+            scope='new_h_%d/act_quant' % k)
+
+        out_bottleneck.append(bottleneck)
+        out_c.append(new_c_act)
+        out_h.append(new_h_act)
+
+      # Since all inputs to the below concats are already quantized, we can use
+      # a regular concat operation.
+      new_c = tf.concat(out_c, axis=3)
+      new_h = tf.concat(out_h, axis=3)
+
+      # |bottleneck| is input to a concat with |new_h|. We must use
+      # quantizable_concat() with a fixed range that matches |new_h|.
+      bottleneck = lstm_utils.quantizable_concat(
+          out_bottleneck,
+          axis=3,
+          is_training=False,
+          is_quantized=self._is_quantized,
+          scope='out_bottleneck/quantized_concat')
+
+      # summary of cell output and new state
+      if self._viz_gates:
+        slim.summaries.add_histogram_summary(new_h, 'cell_output')
+        slim.summaries.add_histogram_summary(new_c, 'cell_state')
+
+      output = new_h
+      if self._output_bottleneck:
+        output = lstm_utils.quantizable_concat(
+            [new_h, bottleneck],
+            axis=3,
+            is_training=False,
+            is_quantized=self._is_quantized,
+            scope='new_output/quantized_concat')
+
+      # reflatten state to store it
+      if self._flatten_state:
+        new_c = tf.reshape(new_c, [-1, self._param_count], name='lstm_c')
+        new_h = tf.reshape(new_h, [-1, self._param_count], name='lstm_h')
+
+      # Set nodes to be under raw_outputs/ name scope for tfmini export.
+      with tf.name_scope(None):
+        new_c = tf.identity(new_c, name='raw_outputs/lstm_c')
+        new_h = tf.identity(new_h, name='raw_outputs/lstm_h')
+      states_and_output = tf.contrib.rnn.LSTMStateTuple(new_c, new_h)
+
+      return output, states_and_output
+
+  def init_state(self, state_name, batch_size, dtype, learned_state=False):
+    """Creates an initial state compatible with this cell.
+
+    Args:
+      state_name: name of the state tensor
+      batch_size: model batch size
+      dtype: dtype for the tensor values i.e. tf.float32
+      learned_state: whether the initial state should be learnable. If false,
+        the initial state is set to all 0's
+
+    Returns:
+      ret: the created initial state
+    """
+    state_size = (self.state_size_flat if self._flatten_state
+                  else self.state_size)
+    # list of 2 zero tensors or variables tensors,
+    # depending on if learned_state is true
+    ret_flat = [(variables.model_variable(
+        state_name + str(i),
+        shape=s,
+        dtype=dtype,
+        initializer=tf.truncated_normal_initializer(stddev=0.03))
+                 if learned_state else tf.zeros(
+                     [batch_size] + s, dtype=dtype, name=state_name))
+                for i, s in enumerate(state_size)]
+
+    # duplicates initial state across the batch axis if it's learned
+    if learned_state:
+      ret_flat = [tf.stack([tensor for i in range(int(batch_size))])
+                  for tensor in ret_flat]
+    for s, r in zip(state_size, ret_flat):
+      r = tf.reshape(r, [-1] + s)
+    ret = tf.contrib.framework.nest.pack_sequence_as(
+        structure=[1, 1], flat_sequence=ret_flat)
+    return ret
+
+  def pre_bottleneck(self, inputs, state, input_index):
+    """Apply pre-bottleneck projection to inputs.
+
+    Pre-bottleneck operation maps features of different channels into the same
+    dimension. The purpose of this op is to share the features from both large
+    and small models in the same LSTM cell.
+
+    Args:
+      inputs: 4D Tensor with shape [batch_size x width x height x input_size].
+      state: 4D Tensor with shape [batch_size x width x height x state_size].
+      input_index: integer index indicating which base features the inputs
+        correspoding to.
+    Returns:
+      inputs: pre-bottlenecked inputs.
+    Raises:
+      ValueError: If pre_bottleneck is not set or inputs is not rank 4.
+    """
+    # Sometimes state is a tuple, in which case it cannot be modified, e.g.
+    # during training, tf.contrib.training.SequenceQueueingStateSaver
+    # returns the state as a tuple. This should not be an issue since we
+    # only need to modify state[1] during export, when state should be a
+    # list.
+    if not self._pre_bottleneck:
+      raise ValueError('Only applied when pre_bottleneck is set to true.')
+    if not len(inputs.shape) == 4:
+      raise ValueError('Expect a rank 4 feature tensor.')
+    if not self._flatten_state and not len(state.shape) == 4:
+      raise ValueError('Expect rank 4 state tensor.')
+    if self._flatten_state and not len(state.shape) == 2:
+      raise ValueError('Expect rank 2 state tensor when flatten_state is set.')
+    with tf.name_scope(None):
+      state = tf.identity(state, name='raw_inputs/init_lstm_h')
+    if self._flatten_state:
+      batch_size = inputs.shape[0]
+      height = inputs.shape[1]
+      width = inputs.shape[2]
+      state = tf.reshape(state, [batch_size, height, width, -1])
+    with tf.variable_scope('conv_lstm_cell', reuse=tf.AUTO_REUSE):
+      state_split = tf.split(state, self._groups, axis=3)
+      with tf.variable_scope('bottleneck_%d' % input_index):
+        bottleneck_out = []
+        for k in range(self._groups):
+          with tf.variable_scope('group_%d' % k):
+            bottleneck_out.append(
+                lstm_utils.quantizable_separable_conv2d(
+                    lstm_utils.quantizable_concat(
+                        [inputs, state_split[k]],
+                        axis=3,
+                        is_training=self._is_training,
+                        is_quantized=self._is_quantized,
+                        scope='quantized_concat'),
+                    self.output_size[-1] / self._groups,
+                    self._filter_size,
+                    is_quantized=self._is_quantized,
+                    depth_multiplier=1,
+                    activation_fn=tf.nn.relu6,
+                    normalizer_fn=None,
+                    scope='project'))
+        inputs = lstm_utils.quantizable_concat(
+            bottleneck_out,
+            axis=3,
+            is_training=self._is_training,
+            is_quantized=self._is_quantized,
+            scope='bottleneck_out/quantized_concat')
+      # For exporting inference graph, we only mark the first timestep.
+      with tf.name_scope(None):
+        inputs = tf.identity(
+            inputs,
+            name='raw_outputs/base_endpoint_%d' % (input_index + 1))
+    return inputs

research/lstm_object_detection/lstm/lstm_cells_test.py

@@ -62,7 +62,7 @@ class BottleneckConvLstmCellsTest(tf.test.TestCase):
         filter_size=filter_size,
         output_size=output_size,
         num_units=num_units,
-        flattened_state=True)
+        flatten_state=True)
     init_state = cell.init_state(
         state_name, batch_size, dtype, learned_state)
     output, state_tuple = cell(inputs, init_state)
@@ -138,6 +138,275 @@ class BottleneckConvLstmCellsTest(tf.test.TestCase):
     self.assertAllEqual([4, 10, 10, 15], init_c.shape.as_list())
     self.assertAllEqual([4, 10, 10, 15], init_h.shape.as_list())
 
+  def test_unroll(self):
+    filter_size = [3, 3]
+    output_size = [10, 10]
+    num_units = 15
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    unroll = 10
+    learned_state = False
+
+    inputs = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
+    cell = lstm_cells.BottleneckConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units)
+    state = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+    for step in range(unroll):
+      output, state = cell(inputs, state)
+    self.assertAllEqual([4, 10, 10, 15], output.shape.as_list())
+    self.assertAllEqual([4, 10, 10, 15], state[0].shape.as_list())
+    self.assertAllEqual([4, 10, 10, 15], state[1].shape.as_list())
+
+  def test_prebottleneck(self):
+    filter_size = [3, 3]
+    output_size = [10, 10]
+    num_units = 15
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    unroll = 10
+    learned_state = False
+
+    inputs_large = tf.zeros([4, 10, 10, 5], dtype=tf.float32)
+    inputs_small = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
+    cell = lstm_cells.BottleneckConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        pre_bottleneck=True)
+    state = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+    for step in range(unroll):
+      if step % 2 == 0:
+        inputs = cell.pre_bottleneck(inputs_large, state[1], 0)
+      else:
+        inputs = cell.pre_bottleneck(inputs_small, state[1], 1)
+      output, state = cell(inputs, state)
+    self.assertAllEqual([4, 10, 10, 15], output.shape.as_list())
+    self.assertAllEqual([4, 10, 10, 15], state[0].shape.as_list())
+    self.assertAllEqual([4, 10, 10, 15], state[1].shape.as_list())
+
+  def test_flatten_state(self):
+    filter_size = [3, 3]
+    output_size = [10, 10]
+    num_units = 15
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    unroll = 10
+    learned_state = False
+
+    inputs_large = tf.zeros([4, 10, 10, 5], dtype=tf.float32)
+    inputs_small = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
+    cell = lstm_cells.BottleneckConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        pre_bottleneck=True,
+        flatten_state=True)
+    state = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+    for step in range(unroll):
+      if step % 2 == 0:
+        inputs = cell.pre_bottleneck(inputs_large, state[1], 0)
+      else:
+        inputs = cell.pre_bottleneck(inputs_small, state[1], 1)
+      output, state = cell(inputs, state)
+    with self.test_session() as sess:
+      sess.run(tf.global_variables_initializer())
+      output_result, state_result = sess.run([output, state])
+      self.assertAllEqual((4, 10, 10, 15), output_result.shape)
+      self.assertAllEqual((4, 10*10*15), state_result[0].shape)
+      self.assertAllEqual((4, 10*10*15), state_result[1].shape)
+
+
+class GroupedConvLstmCellsTest(tf.test.TestCase):
+
+  def test_run_lstm_cell(self):
+    filter_size = [3, 3]
+    output_size = [10, 10]
+    num_units = 16
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    learned_state = False
+
+    inputs = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
+    cell = lstm_cells.GroupedConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        is_training=True)
+    init_state = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+    output, state_tuple = cell(inputs, init_state)
+    self.assertAllEqual([4, 10, 10, 16], output.shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], state_tuple[0].shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], state_tuple[1].shape.as_list())
+
+  def test_run_lstm_cell_with_output_bottleneck(self):
+    filter_size = [3, 3]
+    output_dim = 10
+    output_size = [output_dim] * 2
+    num_units = 16
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    learned_state = False
+
+    inputs = tf.zeros([batch_size, output_dim, output_dim, 3], dtype=tf.float32)
+    cell = lstm_cells.GroupedConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        is_training=True,
+        output_bottleneck=True)
+    init_state = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+    output, state_tuple = cell(inputs, init_state)
+    self.assertAllEqual([4, 10, 10, 32], output.shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], state_tuple[0].shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], state_tuple[1].shape.as_list())
+
+  def test_get_init_state(self):
+    filter_size = [3, 3]
+    output_dim = 10
+    output_size = [output_dim] * 2
+    num_units = 16
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    learned_state = False
+
+    cell = lstm_cells.GroupedConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        is_training=True)
+    init_c, init_h = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+
+    self.assertEqual(tf.float32, init_c.dtype)
+    self.assertEqual(tf.float32, init_h.dtype)
+    with self.test_session() as sess:
+      init_c_res, init_h_res = sess.run([init_c, init_h])
+      self.assertAllClose(np.zeros((4, 10, 10, 16)), init_c_res)
+      self.assertAllClose(np.zeros((4, 10, 10, 16)), init_h_res)
+
+  def test_get_init_learned_state(self):
+    filter_size = [3, 3]
+    output_size = [10, 10]
+    num_units = 16
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    learned_state = True
+
+    cell = lstm_cells.GroupedConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        is_training=True)
+    init_c, init_h = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+
+    self.assertEqual(tf.float32, init_c.dtype)
+    self.assertEqual(tf.float32, init_h.dtype)
+    self.assertAllEqual([4, 10, 10, 16], init_c.shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], init_h.shape.as_list())
+
+  def test_unroll(self):
+    filter_size = [3, 3]
+    output_size = [10, 10]
+    num_units = 16
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    unroll = 10
+    learned_state = False
+
+    inputs = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
+    cell = lstm_cells.GroupedConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        is_training=True)
+    state = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+    for step in range(unroll):
+      output, state = cell(inputs, state)
+    self.assertAllEqual([4, 10, 10, 16], output.shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], state[0].shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], state[1].shape.as_list())
+
+  def test_prebottleneck(self):
+    filter_size = [3, 3]
+    output_size = [10, 10]
+    num_units = 16
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    unroll = 10
+    learned_state = False
+
+    inputs_large = tf.zeros([4, 10, 10, 5], dtype=tf.float32)
+    inputs_small = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
+    cell = lstm_cells.GroupedConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        is_training=True,
+        pre_bottleneck=True)
+    state = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+    for step in range(unroll):
+      if step % 2 == 0:
+        inputs = cell.pre_bottleneck(inputs_large, state[1], 0)
+      else:
+        inputs = cell.pre_bottleneck(inputs_small, state[1], 1)
+      output, state = cell(inputs, state)
+    self.assertAllEqual([4, 10, 10, 16], output.shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], state[0].shape.as_list())
+    self.assertAllEqual([4, 10, 10, 16], state[1].shape.as_list())
+
+  def test_flatten_state(self):
+    filter_size = [3, 3]
+    output_size = [10, 10]
+    num_units = 16
+    state_name = 'lstm_state'
+    batch_size = 4
+    dtype = tf.float32
+    unroll = 10
+    learned_state = False
+
+    inputs_large = tf.zeros([4, 10, 10, 5], dtype=tf.float32)
+    inputs_small = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
+    cell = lstm_cells.GroupedConvLSTMCell(
+        filter_size=filter_size,
+        output_size=output_size,
+        num_units=num_units,
+        is_training=True,
+        pre_bottleneck=True,
+        flatten_state=True)
+    state = cell.init_state(
+        state_name, batch_size, dtype, learned_state)
+    for step in range(unroll):
+      if step % 2 == 0:
+        inputs = cell.pre_bottleneck(inputs_large, state[1], 0)
+      else:
+        inputs = cell.pre_bottleneck(inputs_small, state[1], 1)
+      output, state = cell(inputs, state)
+    with self.test_session() as sess:
+      sess.run(tf.global_variables_initializer())
+      output_result, state_result = sess.run([output, state])
+      self.assertAllEqual((4, 10, 10, 16), output_result.shape)
+      self.assertAllEqual((4, 10*10*16), state_result[0].shape)
+      self.assertAllEqual((4, 10*10*16), state_result[1].shape)
+
 
 if __name__ == '__main__':
   tf.test.main()

research/lstm_object_detection/lstm/rnn_decoder.py

@@ -15,7 +15,7 @@
 
 """Custom RNN decoder."""
 
-from tensorflow.python.ops import variable_scope
+import tensorflow as tf
 
 
 def rnn_decoder(decoder_inputs,
@@ -23,7 +23,7 @@ def rnn_decoder(decoder_inputs,
                 cell,
                 loop_function=None,
                 scope=None):
-  """RNN decoder for the sequence-to-sequence model.
+  """RNN decoder for the LSTM-SSD model.
 
   This decoder returns a list of all states, rather than only the final state.
   Args:
@@ -43,24 +43,205 @@ def rnn_decoder(decoder_inputs,
     A tuple of the form (outputs, state), where:
       outputs: A list of the same length as decoder_inputs of 4D Tensors with
         shape [batch_size x output_size] containing generated outputs.
-      state: A list of the same length as decoder_inputs of the state of each
+      states: A list of the same length as decoder_inputs of the state of each
         cell at each time-step. It is a 2D Tensor of shape
         [batch_size x cell.state_size].
   """
-  with variable_scope.variable_scope(scope or 'rnn_decoder'):
-    state = initial_state
+  with tf.variable_scope(scope or 'rnn_decoder'):
+    state_tuple = initial_state
     outputs = []
     states = []
     prev = None
-    for i, decoder_input in enumerate(decoder_inputs):
+    for local_step, decoder_input in enumerate(decoder_inputs):
       if loop_function is not None and prev is not None:
-        with variable_scope.variable_scope('loop_function', reuse=True):
-          decoder_input = loop_function(prev, i)
-      if i > 0:
-        variable_scope.get_variable_scope().reuse_variables()
-      output, state = cell(decoder_input, state)
+        with tf.variable_scope('loop_function', reuse=True):
+          decoder_input = loop_function(prev, local_step)
+      output, state_tuple = cell(decoder_input, state_tuple)
       outputs.append(output)
-      states.append(state)
+      states.append(state_tuple)
       if loop_function is not None:
         prev = output
   return outputs, states
+
+def multi_input_rnn_decoder(decoder_inputs,
+                            initial_state,
+                            cell,
+                            sequence_step,
+                            selection_strategy='RANDOM',
+                            is_training=None,
+                            is_quantized=False,
+                            preprocess_fn_list=None,
+                            pre_bottleneck=False,
+                            flatten_state=False,
+                            scope=None):
+  """RNN decoder for the Interleaved LSTM-SSD model.
+
+  This decoder takes multiple sequences of inputs and selects the input to feed
+  to the rnn at each timestep using its selection_strategy, which can be random,
+  learned, or deterministic.
+  This decoder returns a list of all states, rather than only the final state.
+  Args:
+    decoder_inputs: A list of lists of 2D Tensors [batch_size x input_size].
+    initial_state: 2D Tensor with shape [batch_size x cell.state_size].
+    cell: rnn_cell.RNNCell defining the cell function and size.
+    sequence_step: Tensor [batch_size] of the step number of the first elements
+      in the sequence.
+    selection_strategy: Method for picking the decoder_input to use at each
+      timestep. Must be 'RANDOM', 'SKIPX' for integer X,  where X is the number
+      of times to use the second input before using the first.
+    is_training: boolean, whether the network is training. When using learned
+      selection, attempts exploration if training.
+    is_quantized: flag to enable/disable quantization mode.
+    preprocess_fn_list: List of functions accepting two tensor arguments: one
+      timestep of decoder_inputs and the lstm state. If not None,
+      decoder_inputs[i] will be updated with preprocess_fn[i] at the start of
+      each timestep.
+    pre_bottleneck: if True, use separate bottleneck weights for each sequence.
+      Useful when input sequences have differing numbers of channels. Final
+      bottlenecks will have the same dimension.
+    flatten_state: Whether the LSTM state is flattened.
+    scope: VariableScope for the created subgraph; defaults to "rnn_decoder".
+  Returns:
+    A tuple of the form (outputs, state), where:
+      outputs: A list of the same length as decoder_inputs of 2D Tensors with
+        shape [batch_size x output_size] containing generated outputs.
+      states: A list of the same length as decoder_inputs of the state of each
+        cell at each time-step. It is a 2D Tensor of shape
+        [batch_size x cell.state_size].
+  Raises:
+    ValueError: If selection_strategy is not recognized or unexpected unroll
+      length.
+  """
+  if flatten_state and len(decoder_inputs[0]) > 1:
+    raise ValueError('In export mode, unroll length should not be more than 1')
+  with tf.variable_scope(scope or 'rnn_decoder'):
+    state_tuple = initial_state
+    outputs = []
+    states = []
+    batch_size = decoder_inputs[0][0].shape[0].value
+    num_sequences = len(decoder_inputs)
+    sequence_length = len(decoder_inputs[0])
+
+    for local_step in range(sequence_length):
+      for sequence_index in range(num_sequences):
+        if preprocess_fn_list is not None:
+          decoder_inputs[sequence_index][local_step] = (
+              preprocess_fn_list[sequence_index](
+                  decoder_inputs[sequence_index][local_step], state_tuple[0]))
+        if pre_bottleneck:
+          decoder_inputs[sequence_index][local_step] = cell.pre_bottleneck(
+              inputs=decoder_inputs[sequence_index][local_step],
+              state=state_tuple[1],
+              input_index=sequence_index)
+
+      action = generate_action(selection_strategy, local_step, sequence_step,
+                               [batch_size, 1, 1, 1])
+      inputs, _ = select_inputs(decoder_inputs, action, local_step)
+      # Mark base network endpoints under raw_inputs/
+      with tf.name_scope(None):
+        inputs = tf.identity(inputs, 'raw_inputs/base_endpoint')
+      output, state_tuple_out = cell(inputs, state_tuple)
+      state_tuple = select_state(state_tuple, state_tuple_out, action)
+
+      outputs.append(output)
+      states.append(state_tuple)
+  return outputs, states
+
+
+def generate_action(selection_strategy, local_step, sequence_step,
+                    action_shape):
+  """Generate current (binary) action based on selection strategy.
+
+  Args:
+    selection_strategy: Method for picking the decoder_input to use at each
+      timestep. Must be 'RANDOM', 'SKIPX' for integer X,  where X is the number
+      of times to use the second input before using the first.
+    local_step: Tensor [batch_size] of the step number within the current
+      unrolled batch.
+    sequence_step: Tensor [batch_size] of the step number of the first elements
+      in the sequence.
+    action_shape: The shape of action tensor to be generated.
+
+  Returns:
+    A tensor of shape action_shape, each element is an individual action.
+
+  Raises:
+    ValueError: if selection_strategy is not supported or if 'SKIP' is not
+      followed by numerics.
+  """
+  if selection_strategy.startswith('RANDOM'):
+    action = tf.random.uniform(action_shape, maxval=2, dtype=tf.int32)
+    action = tf.minimum(action, 1)
+
+    # First step always runs large network.
+    if local_step == 0 and sequence_step is not None:
+      action *= tf.minimum(
+          tf.reshape(tf.cast(sequence_step, tf.int32), action_shape), 1)
+  elif selection_strategy.startswith('SKIP'):
+    inter_count = int(selection_strategy[4:])
+    if local_step % (inter_count + 1) == 0:
+      action = tf.zeros(action_shape)
+    else:
+      action = tf.ones(action_shape)
+  else:
+    raise ValueError('Selection strategy %s not recognized' %
+                     selection_strategy)
+  return tf.cast(action, tf.int32)
+
+
+def select_inputs(decoder_inputs, action, local_step, get_alt_inputs=False):
+  """Selects sequence from decoder_inputs based on 1D actions.
+
+  Given multiple input batches, creates a single output batch by
+  selecting from the action[i]-ith input for the i-th batch element.
+
+  Args:
+    decoder_inputs: A 2-D list of tensor inputs.
+    action: A tensor of shape [batch_size]. Each element corresponds to an index
+      of decoder_inputs to choose.
+    step: The current timestep.
+    get_alt_inputs: Whether the non-chosen inputs should also be returned.
+
+  Returns:
+    The constructed output. Also outputs the elements that were not chosen
+    if get_alt_inputs is True, otherwise None.
+
+  Raises:
+    ValueError: if the decoder inputs contains other than two sequences.
+  """
+  num_seqs = len(decoder_inputs)
+  if not num_seqs == 2:
+    raise ValueError('Currently only supports two sets of inputs.')
+  stacked_inputs = tf.stack(
+      [decoder_inputs[seq_index][local_step] for seq_index in range(num_seqs)],
+      axis=-1)
+  action_index = tf.one_hot(action, num_seqs)
+  inputs = tf.reduce_sum(stacked_inputs * action_index, axis=-1)
+  inputs_alt = None
+  # Only works for 2 models.
+  if get_alt_inputs:
+    # Reverse of action_index.
+    action_index_alt = tf.one_hot(action, num_seqs, on_value=0.0, off_value=1.0)
+    inputs_alt = tf.reduce_sum(stacked_inputs * action_index_alt, axis=-1)
+  return inputs, inputs_alt
+
+def select_state(previous_state, new_state, action):
+  """Select state given action.
+
+  Currently only supports binary action. If action is 0, it means the state is
+  generated from the large model, and thus we will update the state. Otherwise,
+  if the action is 1, it means the state is generated from the small model, and
+  in interleaved model, we skip this state update.
+
+  Args:
+    previous_state: A state tuple representing state from previous step.
+    new_state: A state tuple representing newly computed state.
+    action: A tensor the same shape as state.
+
+  Returns:
+    A state tuple selected based on the given action.
+  """
+  action = tf.cast(action, tf.float32)
+  state_c = previous_state[0] * action + new_state[0] * (1 - action)
+  state_h = previous_state[1] * action + new_state[1] * (1 - action)
+  return (state_c, state_h)

research/lstm_object_detection/lstm/rnn_decoder_test.py

+# Copyright 2019 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.
+# ==============================================================================
+
+"""Tests for lstm_object_detection.lstm.rnn_decoder."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import numpy as np
+
+from lstm_object_detection.lstm import rnn_decoder
+
+
+class MockRnnCell(tf.contrib.rnn.RNNCell):
+
+  def __init__(self, input_size, num_units):
+    self._input_size = input_size
+    self._num_units = num_units
+    self._filter_size = [3, 3]
+
+  def __call__(self, inputs, state_tuple):
+    outputs = tf.concat([inputs, state_tuple[0]], axis=3)
+    new_state_tuple = (tf.multiply(state_tuple[0], 2), state_tuple[1])
+    return outputs, new_state_tuple
+
+  def state_size(self):
+    return self._num_units
+
+  def output_size(self):
+    return self._input_size + self._num_units
+
+  def pre_bottleneck(self, inputs, state, input_index):
+    with tf.variable_scope('bottleneck_%d' % input_index, reuse=tf.AUTO_REUSE):
+      inputs = tf.contrib.layers.separable_conv2d(
+          tf.concat([inputs, state], 3),
+          self._input_size,
+          self._filter_size,
+          depth_multiplier=1,
+          activation_fn=tf.nn.relu6,
+          normalizer_fn=None)
+    return inputs
+
+
+class RnnDecoderTest(tf.test.TestCase):
+
+  def test_rnn_decoder_single_unroll(self):
+    batch_size = 2
+    num_unroll = 1
+    num_units = 64
+    width = 8
+    height = 10
+    input_channels = 128
+
+    initial_state = tf.random_normal((batch_size, width, height, num_units))
+    inputs = tf.random_normal([batch_size, width, height, input_channels])
+
+    rnn_cell = MockRnnCell(input_channels, num_units)
+    outputs, states = rnn_decoder.rnn_decoder(
+        decoder_inputs=[inputs] * num_unroll,
+        initial_state=(initial_state, initial_state),
+        cell=rnn_cell)
+
+    self.assertEqual(len(outputs), num_unroll)
+    self.assertEqual(len(states), num_unroll)
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      results = sess.run((outputs, states, inputs, initial_state))
+      outputs_results = results[0]
+      states_results = results[1]
+      inputs_results = results[2]
+      initial_states_results = results[3]
+      self.assertEqual(outputs_results[0].shape,
+                       (batch_size, width, height, input_channels + num_units))
+      self.assertAllEqual(
+          outputs_results[0],
+          np.concatenate((inputs_results, initial_states_results), axis=3))
+      self.assertEqual(states_results[0][0].shape,
+                       (batch_size, width, height, num_units))
+      self.assertEqual(states_results[0][1].shape,
+                       (batch_size, width, height, num_units))
+      self.assertAllEqual(states_results[0][0],
+                          np.multiply(initial_states_results, 2.0))
+      self.assertAllEqual(states_results[0][1], initial_states_results)
+
+  def test_rnn_decoder_multiple_unroll(self):
+    batch_size = 2
+    num_unroll = 3
+    num_units = 64
+    width = 8
+    height = 10
+    input_channels = 128
+
+    initial_state = tf.random_normal((batch_size, width, height, num_units))
+    inputs = tf.random_normal([batch_size, width, height, input_channels])
+
+    rnn_cell = MockRnnCell(input_channels, num_units)
+    outputs, states = rnn_decoder.rnn_decoder(
+        decoder_inputs=[inputs] * num_unroll,
+        initial_state=(initial_state, initial_state),
+        cell=rnn_cell)
+
+    self.assertEqual(len(outputs), num_unroll)
+    self.assertEqual(len(states), num_unroll)
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      results = sess.run((outputs, states, inputs, initial_state))
+      outputs_results = results[0]
+      states_results = results[1]
+      inputs_results = results[2]
+      initial_states_results = results[3]
+      for i in range(num_unroll):
+        previous_state = ([initial_states_results, initial_states_results]
+                          if i == 0 else states_results[i - 1])
+        self.assertEqual(
+            outputs_results[i].shape,
+            (batch_size, width, height, input_channels + num_units))
+        self.assertAllEqual(
+            outputs_results[i],
+            np.concatenate((inputs_results, previous_state[0]), axis=3))
+        self.assertEqual(states_results[i][0].shape,
+                         (batch_size, width, height, num_units))
+        self.assertEqual(states_results[i][1].shape,
+                         (batch_size, width, height, num_units))
+        self.assertAllEqual(states_results[i][0],
+                            np.multiply(previous_state[0], 2.0))
+        self.assertAllEqual(states_results[i][1], previous_state[1])
+
+
+class MultiInputRnnDecoderTest(tf.test.TestCase):
+
+  def test_rnn_decoder_single_unroll(self):
+    batch_size = 2
+    num_unroll = 1
+    num_units = 12
+    width = 8
+    height = 10
+    input_channels_large = 24
+    input_channels_small = 12
+    bottleneck_channels = 20
+
+    initial_state_c = tf.random_normal((batch_size, width, height, num_units))
+    initial_state_h = tf.random_normal((batch_size, width, height, num_units))
+    initial_state = (initial_state_c, initial_state_h)
+    inputs_large = tf.random_normal(
+        [batch_size, width, height, input_channels_large])
+    inputs_small = tf.random_normal(
+        [batch_size, width, height, input_channels_small])
+
+    rnn_cell = MockRnnCell(bottleneck_channels, num_units)
+    outputs, states = rnn_decoder.multi_input_rnn_decoder(
+        decoder_inputs=[[inputs_large] * num_unroll,
+                        [inputs_small] * num_unroll],
+        initial_state=initial_state,
+        cell=rnn_cell,
+        sequence_step=tf.zeros([batch_size]),
+        pre_bottleneck=True)
+
+    self.assertEqual(len(outputs), num_unroll)
+    self.assertEqual(len(states), num_unroll)
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      results = sess.run(
+          (outputs, states, inputs_large, inputs_small, initial_state))
+      outputs_results = results[0]
+      states_results = results[1]
+      inputs_large_results = results[2]
+      inputs_small_results = results[3]
+      initial_states_results = results[4]
+      self.assertEqual(
+          outputs_results[0].shape,
+          (batch_size, width, height, bottleneck_channels + num_units))
+      self.assertEqual(states_results[0][0].shape,
+                       (batch_size, width, height, num_units))
+      self.assertEqual(states_results[0][1].shape,
+                       (batch_size, width, height, num_units))
+      # The first step should always update state.
+      self.assertAllEqual(states_results[0][0],
+                              np.multiply(initial_states_results[0], 2))
+      self.assertAllEqual(states_results[0][1], initial_states_results[1])
+
+  def test_rnn_decoder_multiple_unroll(self):
+    batch_size = 2
+    num_unroll = 3
+    num_units = 12
+    width = 8
+    height = 10
+    input_channels_large = 24
+    input_channels_small = 12
+    bottleneck_channels = 20
+
+    initial_state_c = tf.random_normal((batch_size, width, height, num_units))
+    initial_state_h = tf.random_normal((batch_size, width, height, num_units))
+    initial_state = (initial_state_c, initial_state_h)
+    inputs_large = tf.random_normal(
+        [batch_size, width, height, input_channels_large])
+    inputs_small = tf.random_normal(
+        [batch_size, width, height, input_channels_small])
+
+    rnn_cell = MockRnnCell(bottleneck_channels, num_units)
+    outputs, states = rnn_decoder.multi_input_rnn_decoder(
+        decoder_inputs=[[inputs_large] * num_unroll,
+                        [inputs_small] * num_unroll],
+        initial_state=initial_state,
+        cell=rnn_cell,
+        sequence_step=tf.zeros([batch_size]),
+        pre_bottleneck=True)
+
+    self.assertEqual(len(outputs), num_unroll)
+    self.assertEqual(len(states), num_unroll)
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      results = sess.run(
+          (outputs, states, inputs_large, inputs_small, initial_state))
+      outputs_results = results[0]
+      states_results = results[1]
+      inputs_large_results = results[2]
+      inputs_small_results = results[3]
+      initial_states_results = results[4]
+
+      # The first step should always update state.
+      self.assertAllEqual(states_results[0][0],
+                          np.multiply(initial_states_results[0], 2))
+      self.assertAllEqual(states_results[0][1], initial_states_results[1])
+      for i in range(num_unroll):
+        self.assertEqual(
+            outputs_results[i].shape,
+            (batch_size, width, height, bottleneck_channels + num_units))
+        self.assertEqual(states_results[i][0].shape,
+                         (batch_size, width, height, num_units))
+        self.assertEqual(states_results[i][1].shape,
+                         (batch_size, width, height, num_units))
+
+  def test_rnn_decoder_multiple_unroll_with_skip(self):
+    batch_size = 2
+    num_unroll = 5
+    num_units = 12
+    width = 8
+    height = 10
+    input_channels_large = 24
+    input_channels_small = 12
+    bottleneck_channels = 20
+    skip = 2
+
+    initial_state_c = tf.random_normal((batch_size, width, height, num_units))
+    initial_state_h = tf.random_normal((batch_size, width, height, num_units))
+    initial_state = (initial_state_c, initial_state_h)
+    inputs_large = tf.random_normal(
+        [batch_size, width, height, input_channels_large])
+    inputs_small = tf.random_normal(
+        [batch_size, width, height, input_channels_small])
+
+    rnn_cell = MockRnnCell(bottleneck_channels, num_units)
+    outputs, states = rnn_decoder.multi_input_rnn_decoder(
+        decoder_inputs=[[inputs_large] * num_unroll,
+                        [inputs_small] * num_unroll],
+        initial_state=initial_state,
+        cell=rnn_cell,
+        sequence_step=tf.zeros([batch_size]),
+        pre_bottleneck=True,
+        selection_strategy='SKIP%d' % skip)
+
+    self.assertEqual(len(outputs), num_unroll)
+    self.assertEqual(len(states), num_unroll)
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      results = sess.run(
+          (outputs, states, inputs_large, inputs_small, initial_state))
+      outputs_results = results[0]
+      states_results = results[1]
+      inputs_large_results = results[2]
+      inputs_small_results = results[3]
+      initial_states_results = results[4]
+
+      for i in range(num_unroll):
+        self.assertEqual(
+            outputs_results[i].shape,
+            (batch_size, width, height, bottleneck_channels + num_units))
+        self.assertEqual(states_results[i][0].shape,
+                         (batch_size, width, height, num_units))
+        self.assertEqual(states_results[i][1].shape,
+                         (batch_size, width, height, num_units))
+
+        previous_state = (
+            initial_states_results if i == 0 else states_results[i - 1])
+        # State only updates during key frames
+        if i % (skip + 1) == 0:
+          self.assertAllEqual(states_results[i][0],
+                              np.multiply(previous_state[0], 2))
+          self.assertAllEqual(states_results[i][1], previous_state[1])
+        else:
+          self.assertAllEqual(states_results[i][0], previous_state[0])
+          self.assertAllEqual(states_results[i][1], previous_state[1])
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/lstm_object_detection/lstm/utils.py

+# Copyright 2019 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.
+# ==============================================================================
+
+"""Quantization related ops for LSTM."""
+
+from __future__ import absolute_import
+from __future__ import division
+
+import tensorflow as tf
+from tensorflow.python.training import moving_averages
+
+
+def _quant_var(
+    name,
+    initializer_val,
+    vars_collection=tf.GraphKeys.MOVING_AVERAGE_VARIABLES,
+):
+  """Create an var for storing the min/max quantization range."""
+  return tf.contrib.framework.model_variable(
+      name,
+      shape=[],
+      initializer=tf.constant_initializer(initializer_val),
+      collections=[vars_collection],
+      trainable=False)
+
+
+def quantizable_concat(inputs,
+                       axis,
+                       is_training,
+                       is_quantized=True,
+                       default_min=0,
+                       default_max=6,
+                       ema_decay=0.999,
+                       scope='quantized_concat'):
+  """Concat replacement with quantization option.
+
+  Allows concat inputs to share the same min max ranges,
+  from experimental/gazelle/synthetic/model/tpu/utils.py.
+
+  Args:
+    inputs: list of tensors to concatenate.
+    axis: dimension along which to concatenate.
+    is_training: true if the graph is a training graph.
+    is_quantized: flag to enable/disable quantization.
+    default_min: default min value for fake quant op.
+    default_max: default max value for fake quant op.
+    ema_decay: the moving average decay for the quantization variables.
+    scope: Optional scope for variable_scope.
+
+  Returns:
+    Tensor resulting from concatenation of input tensors
+  """
+  if is_quantized:
+    with tf.variable_scope(scope):
+      tf.logging.info('inputs: {}'.format(inputs))
+      for t in inputs:
+        tf.logging.info(t)
+
+      min_var = _quant_var('min', default_min)
+      max_var = _quant_var('max', default_max)
+      if not is_training:
+        # If we are building an eval graph just use the values in the variables.
+        quant_inputs = [
+            tf.fake_quant_with_min_max_vars(t, min_var, max_var) for t in inputs
+        ]
+        tf.logging.info('min_val: {}'.format(min_var))
+        tf.logging.info('max_val: {}'.format(max_var))
+      else:
+        concat_tensors = tf.concat(inputs, axis=axis)
+        tf.logging.info('concat_tensors: {}'.format(concat_tensors))
+        # Otherwise we need to keep track of the moving averages of the min and
+        # of the elements of the input tensor max.
+        min_val = moving_averages.assign_moving_average(
+            min_var,
+            tf.reduce_min(concat_tensors),
+            ema_decay,
+            name='AssignMinEma')
+        max_val = moving_averages.assign_moving_average(
+            max_var,
+            tf.reduce_max(concat_tensors),
+            ema_decay,
+            name='AssignMaxEma')
+        tf.logging.info('min_val: {}'.format(min_val))
+        tf.logging.info('max_val: {}'.format(max_val))
+        quant_inputs = [
+            tf.fake_quant_with_min_max_vars(t, min_val, max_val) for t in inputs
+        ]
+      tf.logging.info('quant_inputs: {}'.format(quant_inputs))
+      outputs = tf.concat(quant_inputs, axis=axis)
+      tf.logging.info('outputs: {}'.format(outputs))
+  else:
+    outputs = tf.concat(inputs, axis=axis)
+  return outputs
+
+
+def quantizable_separable_conv2d(inputs,
+                                 num_outputs,
+                                 kernel_size,
+                                 is_quantized=True,
+                                 depth_multiplier=1,
+                                 stride=1,
+                                 activation_fn=tf.nn.relu6,
+                                 normalizer_fn=None,
+                                 scope=None):
+  """Quantization friendly backward compatible separable conv2d.
+
+  This op has the same API is separable_conv2d. The main difference is that an
+  additional BiasAdd is manually inserted after the depthwise conv, such that
+  the depthwise bias will not have name conflict with pointwise bias. The
+  motivation of this op is that quantization script need BiasAdd in order to
+  recognize the op, in which a native call to separable_conv2d do not create
+  for the depthwise conv.
+
+  Args:
+    inputs: A tensor of size [batch_size, height, width, channels].
+    num_outputs: The number of pointwise convolution output filters. If is
+      None, then we skip the pointwise convolution stage.
+    kernel_size: A list of length 2: [kernel_height, kernel_width] of the
+      filters. Can be an int if both values are the same.
+    is_quantized: flag to enable/disable quantization.
+    depth_multiplier: The number of depthwise convolution output channels for
+      each input channel. The total number of depthwise convolution output
+      channels will be equal to num_filters_in * depth_multiplier.
+    stride: A list of length 2: [stride_height, stride_width], specifying the
+      depthwise convolution stride. Can be an int if both strides are the same.
+    activation_fn: Activation function. The default value is a ReLU function.
+      Explicitly set it to None to skip it and maintain a linear activation.
+    normalizer_fn: Normalization function to use instead of biases.
+    scope: Optional scope for variable_scope.
+
+  Returns:
+    Tensor resulting from concatenation of input tensors
+  """
+  if is_quantized:
+    outputs = tf.contrib.layers.separable_conv2d(
+        inputs,
+        None,
+        kernel_size,
+        depth_multiplier=depth_multiplier,
+        stride=1,
+        activation_fn=None,
+        normalizer_fn=None,
+        biases_initializer=None,
+        scope=scope)
+    outputs = tf.contrib.layers.bias_add(
+        outputs, trainable=True, scope='%s_bias' % scope)
+    outputs = tf.contrib.layers.conv2d(
+        outputs,
+        num_outputs, [1, 1],
+        activation_fn=activation_fn,
+        stride=stride,
+        normalizer_fn=normalizer_fn,
+        scope=scope)
+  else:
+    outputs = tf.contrib.layers.separable_conv2d(
+        inputs,
+        num_outputs,
+        kernel_size,
+        depth_multiplier=depth_multiplier,
+        stride=stride,
+        activation_fn=activation_fn,
+        normalizer_fn=normalizer_fn,
+        scope=scope)
+  return outputs
+
+
+def quantize_op(inputs,
+                is_training=True,
+                is_quantized=True,
+                default_min=0,
+                default_max=6,
+                ema_decay=0.999,
+                scope='quant'):
+  """Inserts a fake quantization op after inputs.
+
+  Args:
+    inputs: A tensor of size [batch_size, height, width, channels].
+    is_training: true if the graph is a training graph.
+    is_quantized: flag to enable/disable quantization.
+    default_min: default min value for fake quant op.
+    default_max: default max value for fake quant op.
+    ema_decay: the moving average decay for the quantization variables.
+    scope: Optional scope for variable_scope.
+
+  Returns:
+    Tensor resulting from quantizing the input tensors.
+  """
+  if is_quantized:
+    with tf.variable_scope(scope):
+      min_var = _quant_var('min', default_min)
+      max_var = _quant_var('max', default_max)
+      if is_training:
+        min_val = moving_averages.assign_moving_average(
+            min_var, tf.reduce_min(inputs), ema_decay, name='AssignMinEma')
+        max_val = moving_averages.assign_moving_average(
+            max_var, tf.reduce_max(inputs), ema_decay, name='AssignMaxEma')
+        inputs = tf.fake_quant_with_min_max_vars(inputs, min_val, max_val)
+      else:
+        inputs = tf.fake_quant_with_min_max_vars(inputs, min_var, max_var)
+  return inputs

research/lstm_object_detection/lstm/utils_test.py

+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for lstm_object_detection.lstm.utils."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+from lstm_object_detection.lstm import utils
+
+
+class QuantizableUtilsTest(tf.test.TestCase):
+
+  def test_quantizable_concat_is_training(self):
+    inputs_1 = tf.zeros([4, 10, 10, 1], dtype=tf.float32)
+    inputs_2 = tf.ones([4, 10, 10, 2], dtype=tf.float32)
+    concat_in_train = utils.quantizable_concat([inputs_1, inputs_2],
+                                               axis=3,
+                                               is_training=True)
+    self.assertAllEqual([4, 10, 10, 3], concat_in_train.shape.as_list())
+    self._check_min_max_ema(tf.get_default_graph())
+    self._check_min_max_vars(tf.get_default_graph())
+
+  def test_quantizable_concat_inference(self):
+    inputs_1 = tf.zeros([4, 10, 10, 1], dtype=tf.float32)
+    inputs_2 = tf.ones([4, 10, 10, 2], dtype=tf.float32)
+    concat_in_train = utils.quantizable_concat([inputs_1, inputs_2],
+                                               axis=3,
+                                               is_training=False)
+    self.assertAllEqual([4, 10, 10, 3], concat_in_train.shape.as_list())
+    self._check_no_min_max_ema(tf.get_default_graph())
+    self._check_min_max_vars(tf.get_default_graph())
+
+  def test_quantizable_concat_not_quantized_is_training(self):
+    inputs_1 = tf.zeros([4, 10, 10, 1], dtype=tf.float32)
+    inputs_2 = tf.ones([4, 10, 10, 2], dtype=tf.float32)
+    concat_in_train = utils.quantizable_concat([inputs_1, inputs_2],
+                                               axis=3,
+                                               is_training=True,
+                                               is_quantized=False)
+    self.assertAllEqual([4, 10, 10, 3], concat_in_train.shape.as_list())
+    self._check_no_min_max_ema(tf.get_default_graph())
+    self._check_no_min_max_vars(tf.get_default_graph())
+
+  def test_quantizable_concat_not_quantized_inference(self):
+    inputs_1 = tf.zeros([4, 10, 10, 1], dtype=tf.float32)
+    inputs_2 = tf.ones([4, 10, 10, 2], dtype=tf.float32)
+    concat_in_train = utils.quantizable_concat([inputs_1, inputs_2],
+                                               axis=3,
+                                               is_training=False,
+                                               is_quantized=False)
+    self.assertAllEqual([4, 10, 10, 3], concat_in_train.shape.as_list())
+    self._check_no_min_max_ema(tf.get_default_graph())
+    self._check_no_min_max_vars(tf.get_default_graph())
+
+  def test_quantize_op_is_training(self):
+    inputs = tf.zeros([4, 10, 10, 128], dtype=tf.float32)
+    outputs = utils.quantize_op(inputs)
+    self.assertAllEqual(inputs.shape.as_list(), outputs.shape.as_list())
+    self._check_min_max_ema(tf.get_default_graph())
+    self._check_min_max_vars(tf.get_default_graph())
+
+  def test_quantize_op_inferene(self):
+    inputs = tf.zeros([4, 10, 10, 128], dtype=tf.float32)
+    outputs = utils.quantize_op(inputs, is_training=False)
+    self.assertAllEqual(inputs.shape.as_list(), outputs.shape.as_list())
+    self._check_no_min_max_ema(tf.get_default_graph())
+    self._check_min_max_vars(tf.get_default_graph())
+
+  def _check_min_max_vars(self, graph):
+    op_types = [op.type for op in graph.get_operations()]
+    self.assertTrue(
+        any('FakeQuantWithMinMaxVars' in op_type for op_type in op_types))
+
+  def _check_min_max_ema(self, graph):
+    op_names = [op.name for op in graph.get_operations()]
+    self.assertTrue(any('AssignMinEma' in name for name in op_names))
+    self.assertTrue(any('AssignMaxEma' in name for name in op_names))
+
+  def _check_no_min_max_vars(self, graph):
+    op_types = [op.type for op in graph.get_operations()]
+    self.assertFalse(
+        any('FakeQuantWithMinMaxVars' in op_type for op_type in op_types))
+
+  def _check_no_min_max_ema(self, graph):
+    op_names = [op.name for op in graph.get_operations()]
+    self.assertFalse(any('AssignMinEma' in name for name in op_names))
+    self.assertFalse(any('AssignMaxEma' in name for name in op_names))
+
+
+class QuantizableSeparableConv2dTest(tf.test.TestCase):
+
+  def test_quantizable_separable_conv2d(self):
+    inputs = tf.zeros([4, 10, 10, 128], dtype=tf.float32)
+    num_outputs = 64
+    kernel_size = [3, 3]
+    scope = 'QuantSeparable'
+    outputs = utils.quantizable_separable_conv2d(
+        inputs, num_outputs, kernel_size, scope=scope)
+    self.assertAllEqual([4, 10, 10, num_outputs], outputs.shape.as_list())
+    self._check_depthwise_bias_add(tf.get_default_graph(), scope)
+
+  def test_quantizable_separable_conv2d_not_quantized(self):
+    inputs = tf.zeros([4, 10, 10, 128], dtype=tf.float32)
+    num_outputs = 64
+    kernel_size = [3, 3]
+    scope = 'QuantSeparable'
+    outputs = utils.quantizable_separable_conv2d(
+        inputs, num_outputs, kernel_size, is_quantized=False, scope=scope)
+    self.assertAllEqual([4, 10, 10, num_outputs], outputs.shape.as_list())
+    self._check_no_depthwise_bias_add(tf.get_default_graph(), scope)
+
+  def _check_depthwise_bias_add(self, graph, scope):
+    op_names = [op.name for op in graph.get_operations()]
+    self.assertTrue(
+        any('%s_bias/BiasAdd' % scope in name for name in op_names))
+
+  def _check_no_depthwise_bias_add(self, graph, scope):
+    op_names = [op.name for op in graph.get_operations()]
+    self.assertFalse(
+        any('%s_bias/BiasAdd' % scope in name for name in op_names))
+
+
+if __name__ == '__main__':
+  tf.test.main()