Make quantizable_separable_conv2d initializer explicitly configurable....

Make quantizable_separable_conv2d initializer explicitly configurable. (otherwise we could use slim.arg_scope to config the initializer) (#8661)

PiperOrigin-RevId: 315732759

research/lstm_object_detection/lstm/lstm_cells.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 """BottleneckConvLSTMCell implementation."""
+import functools
 
 import tensorflow.compat.v1 as tf
 import tf_slim as slim
+
 from tensorflow.contrib import rnn as contrib_rnn
 from tensorflow.contrib.framework.python.ops import variables as contrib_variables
 import lstm_object_detection.lstm.utils as lstm_utils
@@ -285,7 +287,8 @@ class GroupedConvLSTMCell(contrib_rnn.RNNCell):
                output_bottleneck=False,
                pre_bottleneck=False,
                is_quantized=False,
-               visualize_gates=False):
+               visualize_gates=False,
+               conv_op_overrides=None):
     """Initialize the basic LSTM cell.
 
     Args:
@@ -311,6 +314,10 @@ class GroupedConvLSTMCell(contrib_rnn.RNNCell):
         quantization friendly concat and separable_conv2d ops.
       visualize_gates: if True, add histogram summaries of all gates and outputs
         to tensorboard
+      conv_op_overrides: A list of convolutional operations that override the
+        'bottleneck' and 'convolution' layers before lstm gates. If None, the
+        original implementation of seperable_conv will be used. The length of
+        the list should be two.
 
     Raises:
       ValueError: when both clip_state and scale_state are enabled.
@@ -336,6 +343,10 @@ class GroupedConvLSTMCell(contrib_rnn.RNNCell):
     self._is_quantized = is_quantized
     for dim in self._output_size:
       self._param_count *= dim
+    self._conv_op_overrides = conv_op_overrides
+    if self._conv_op_overrides and len(self._conv_op_overrides) != 2:
+      raise ValueError('Bottleneck and Convolutional layer should be overriden'
+                       'together')
 
   @property
   def state_size(self):
@@ -405,23 +416,26 @@ class GroupedConvLSTMCell(contrib_rnn.RNNCell):
         if self._pre_bottleneck:
           bottleneck = inputs_list[k]
         else:
+          if self._conv_op_overrides:
+            bottleneck_fn = self._conv_op_overrides[0]
+          else:
+            bottleneck_fn = functools.partial(
+                lstm_utils.quantizable_separable_conv2d,
+                kernel_size=self._filter_size,
+                activation_fn=self._activation)
           if self._use_batch_norm:
-            b_x = lstm_utils.quantizable_separable_conv2d(
-                inputs,
-                self._num_units // self._groups,
-                self._filter_size,
+            b_x = bottleneck_fn(
+                inputs=inputs,
+                num_outputs=self._num_units // self._groups,
                 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,
+            b_h = bottleneck_fn(
+                inputs=h_list[k],
+                num_outputs=self._num_units // self._groups,
                 is_quantized=self._is_quantized,
                 depth_multiplier=1,
-                activation_fn=None,
                 normalizer_fn=None,
                 scope='bottleneck_%d_h' % k)
             b_x = slim.batch_norm(
@@ -445,24 +459,26 @@ class GroupedConvLSTMCell(contrib_rnn.RNNCell):
                 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,
+            bottleneck = bottleneck_fn(
+                inputs=bottleneck_concat,
+                num_outputs=self._num_units // self._groups,
                 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,
+        if self._conv_op_overrides:
+          conv_fn = self._conv_op_overrides[1]
+        else:
+          conv_fn = functools.partial(
+              lstm_utils.quantizable_separable_conv2d,
+              kernel_size=self._filter_size,
+              activation_fn=None)
+        concat = conv_fn(
+            inputs=bottleneck,
+            num_outputs=4 * self._num_units // self._groups,
             is_quantized=self._is_quantized,
             depth_multiplier=1,
-            activation_fn=None,
             normalizer_fn=None,
             scope='concat_conv_%d' % k)
 
@@ -490,14 +506,6 @@ class GroupedConvLSTMCell(contrib_rnn.RNNCell):
             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.fixed_quantize_op(
-            f_act,
-            fixed_min=0.0,
-            fixed_max=1.0,
-            is_quantized=self._is_quantized,
-            scope='forget_gate_%d/act_quant' % k)
 
         a = c_list[k] * f_act
         a = lstm_utils.quantize_op(
@@ -507,14 +515,6 @@ class GroupedConvLSTMCell(contrib_rnn.RNNCell):
             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.fixed_quantize_op(
-            i_act,
-            fixed_min=0.0,
-            fixed_max=1.0,
-            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
@@ -567,14 +567,6 @@ class GroupedConvLSTMCell(contrib_rnn.RNNCell):
             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.fixed_quantize_op(
-            o_act,
-            fixed_min=0.0,
-            fixed_max=1.0,
-            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.

research/lstm_object_detection/lstm/utils.py

@@ -120,6 +120,8 @@ def quantizable_separable_conv2d(inputs,
                                  stride=1,
                                  activation_fn=tf.nn.relu6,
                                  normalizer_fn=None,
+                                 weights_initializer=None,
+                                 pointwise_initializer=None,
                                  scope=None):
   """Quantization friendly backward compatible separable conv2d.
 
@@ -145,6 +147,8 @@ def quantizable_separable_conv2d(inputs,
     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.
+    weights_initializer: An initializer for the depthwise weights.
+    pointwise_initializer: An initializer for the pointwise weights.
     scope: Optional scope for variable_scope.
 
   Returns:
@@ -160,6 +164,8 @@ def quantizable_separable_conv2d(inputs,
         activation_fn=None,
         normalizer_fn=None,
         biases_initializer=None,
+        weights_initializer=weights_initializer,
+        pointwise_initializer=None,
         scope=scope)
     outputs = contrib_layers.bias_add(
         outputs, trainable=True, scope='%s_bias' % scope)
@@ -169,6 +175,7 @@ def quantizable_separable_conv2d(inputs,
         activation_fn=activation_fn,
         stride=stride,
         normalizer_fn=normalizer_fn,
+        weights_initializer=pointwise_initializer,
         scope=scope)
   else:
     outputs = contrib_layers.separable_conv2d(
@@ -179,6 +186,8 @@ def quantizable_separable_conv2d(inputs,
         stride=stride,
         activation_fn=activation_fn,
         normalizer_fn=normalizer_fn,
+        weights_initializer=weights_initializer,
+        pointwise_initializer=pointwise_initializer,
         scope=scope)
   return outputs
 

research/lstm_object_detection/meta_architectures/lstm_ssd_meta_arch_test.py

@@ -82,7 +82,7 @@ class FakeLSTMFeatureExtractor(
         min_depth=self._min_depth,
         insert_1x1_conv=True,
         image_features=image_features)
-    return feature_maps.values()
+    return list(feature_maps.values())
 
 
 class FakeLSTMInterleavedFeatureExtractor(
@@ -141,7 +141,7 @@ class FakeLSTMInterleavedFeatureExtractor(
         min_depth=self._min_depth,
         insert_1x1_conv=True,
         image_features=image_features)
-    return feature_maps.values()
+    return list(feature_maps.values())
 
 
 class MockAnchorGenerator2x2(anchor_generator.AnchorGenerator):