Allow specifying ops that the regularizer does not pass during traversal.

PiperOrigin-RevId: 248683238

Author: pkch

README.md

@@ -130,10 +130,14 @@ models with batch norm; it requires that batch norm scale is enabled.
 * *Latency* optimizes for the estimated inference latency of the network, based
 on the specific hardware characteristics.
 
-## Example: Adding a FLOPs Regularizer
+## Examples
+
+### Adding a FLOPs Regularizer
 
 The example below demonstrates how to use MorphNet to reduce the number of FLOPs
-in your model.
+in your model. In this example, the regularizer will traverse the graph
+starting with `logits`, and will not go past any op whose name matches the regex
+`/images.*`; this allows to specify the subgraph for MorphNet to optimize.
 
 ```python
 from morph_net.network_regularizers import flop_regularizer
@@ -142,7 +146,7 @@ from morph_net.tools import structure_exporter
 logits = build_model()
 
 network_regularizer = flop_regularizer.GammaFlopsRegularizer(
-    [logits.op], gamma_threshold=1e-3)
+    [logits.op], input_boundary=[images, labels], gamma_threshold=1e-3)
 regularization_strength = 1e-10
 regularizer_loss = (network_regularizer.get_regularization_term() * regularization_strength)
 

morph_net/framework/op_regularizer_manager.py

@@ -5,6 +5,7 @@
 from __future__ import print_function
 
 import collections
+
 from morph_net.framework import concat_and_slice_regularizers
 from morph_net.framework import constant_op_regularizer
 from morph_net.framework import grouping_regularizers
@@ -48,6 +49,7 @@ def __init__(
       create_grouping_regularizer=grouping_regularizers.MaxGroupingRegularizer,
       force_group=None,
       regularizer_blacklist=None,
+      input_boundary=None,
       iteration_limit=ITERATION_LIMIT):
     """Creates an instance of OpRegularizerManager.
 
@@ -84,6 +86,8 @@ def __init__(
         multiple patterns in a single regex.
       regularizer_blacklist: List of regex for ops that should not be
         regularized.
+      input_boundary: A list of ops that represent the input boundary of the
+        subgraph being regularized (input boundary is not regularized).
       iteration_limit: Integer iteration limit for OpRegularizerManager to
         finish analyzing the network.  If the limit is reached, it is assumed
         that OpRegularizerManager got stuck in a loop.
@@ -115,7 +119,7 @@ def __init__(
 
     # Start DFS from outputs to find all source ops.
     tf.logging.info('OpRegularizerManager starting analysis from: %s.', ops)
-    self._dfs_for_source_ops(ops)
+    self._dfs_for_source_ops(ops, input_boundary)
     tf.logging.info('OpRegularizerManager found %d ops and %d sources.',
                     len(self._all_ops), len(self._op_deque))
 
@@ -581,18 +585,25 @@ def _get_source_slices(self, sizes, op_slices):
       size_index += 1
     return is_source
 
-  def _dfs_for_source_ops(self, ops):
+  def _dfs_for_source_ops(self, ops, input_boundary=None):
     """Performs DFS from ops and finds source ops to process.
 
     Args:
-      ops: List of tf.Operation.
+      ops: A list of tf.Operation's.
+      input_boundary: A list of ops where traversal should terminate.
     """
+    if input_boundary:
+      input_boundary = set(input_boundary)
+    else:
+      input_boundary = set()
     to_visit = list(ops)
     visited = set()
     while to_visit:
       # Get next op and mark as visited.
       op = to_visit.pop()
       visited.add(op)
+      if op in input_boundary:
+        continue
       self._all_ops.add(op)
 
       # Check if op is a source by querying OpHandler.

morph_net/network_regularizers/activation_regularizer.py

@@ -27,6 +27,7 @@ def __init__(
       gamma_threshold,
       regularizer_decorator: Type[generic_regularizers.OpRegularizer] = None,
       decorator_parameters=None,
+      input_boundary=None,
       force_group=None,
       regularizer_blacklist=None):
     """Creates a GammaActivationRegularizer object.
@@ -42,6 +43,8 @@ def __init__(
       decorator_parameters: A dictionary of parameters to pass to the decorator
         factory. To be used only with decorators that requires parameters,
         otherwise use None.
+      input_boundary: A list of ops that represent the input boundary of the
+        subgraph being regularized (input boundary is not regularized).
       force_group: List of regex for ops that should be force-grouped.  Each
         regex corresponds to a separate group.  Use '|' operator to specify
         multiple patterns in a single regex. See op_regularizer_manager for more
@@ -63,6 +66,7 @@ def __init__(
     self._manager = orm.OpRegularizerManager(
         ops,
         op_handler_dict,
+        input_boundary=input_boundary,
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
@@ -97,6 +101,7 @@ def __init__(
       l1_fraction=0,
       regularizer_decorator: Type[generic_regularizers.OpRegularizer] = None,
       decorator_parameters=None,
+      input_boundary=None,
       force_group=None,
       regularizer_blacklist=None):
     """Creates a GroupLassoActivationRegularizer object.
@@ -113,6 +118,8 @@ def __init__(
       decorator_parameters: A dictionary of parameters to pass to the decorator
         factory. To be used only with decorators that requires parameters,
         otherwise use None.
+      input_boundary: A list of ops that represent the input boundary of the
+        subgraph being regularized (input boundary is not regularized).
       force_group: List of regex for ops that should be force-grouped.  Each
         regex corresponds to a separate group.  Use '|' operator to specify
         multiple patterns in a single regex. See op_regularizer_manager for more
@@ -145,6 +152,7 @@ def __init__(
     self._manager = orm.OpRegularizerManager(
         ops,
         op_handler_dict,
+        input_boundary=input_boundary,
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(

morph_net/network_regularizers/flop_regularizer.py

@@ -27,21 +27,24 @@ def __init__(
       gamma_threshold,
       regularizer_decorator: Type[generic_regularizers.OpRegularizer] = None,
       decorator_parameters=None,
+      input_boundary=None,
       force_group=None,
       regularizer_blacklist=None):
     """Creates a GammaFlopsRegularizer object.
 
     Args:
       ops: A list of tf.Operation. An OpRegularizer will be created for all the
         ops in `ops`, and recursively for all ops they depend on via data
-        dependency. Typically `ops` would contain a single tf.Operation, which
-        is the output of the network.
+        dependency that does not involve input ops. Typically `ops` would
+        contain a single tf.Operation, which is the output of the network.
       gamma_threshold: A float scalar, will be used as a 'gamma_threshold' for
         all instances GammaL1Regularizer created by this class.
       regularizer_decorator: A class of OpRegularizer decorator to use.
       decorator_parameters: A dictionary of parameters to pass to the decorator
         factory. To be used only with decorators that requires parameters,
         otherwise use None.
+      input_boundary: A list of ops that represent the input boundary of the
+        subgraph being regularized (input boundary is not regularized).
       force_group: List of regex for ops that should be force-grouped.  Each
         regex corresponds to a separate group.  Use '|' operator to specify
         multiple patterns in a single regex. See op_regularizer_manager for more
@@ -63,6 +66,7 @@ def __init__(
     self._manager = orm.OpRegularizerManager(
         ops,
         op_handler_dict,
+        input_boundary=input_boundary,
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
@@ -97,6 +101,7 @@ def __init__(
       l1_fraction=0,
       regularizer_decorator: Type[generic_regularizers.OpRegularizer] = None,
       decorator_parameters=None,
+      input_boundary=None,
       force_group=None,
       regularizer_blacklist=None,
       convert_to_variable=True):
@@ -114,6 +119,8 @@ def __init__(
       decorator_parameters: A dictionary of parameters to pass to the decorator
         factory. To be used only with decorators that requires parameters,
         otherwise use None.
+      input_boundary: A list of ops that represent the input boundary of the
+        subgraph being regularized (input boundary is not regularized).
       force_group: List of regex for ops that should be force-grouped.  Each
         regex corresponds to a separate group.  Use '|' operator to specify
         multiple patterns in a single regex. See op_regularizer_manager for more
@@ -149,6 +156,7 @@ def __init__(
     self._manager = orm.OpRegularizerManager(
         ops,
         op_handler_dict,
+        input_boundary=input_boundary,
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(

morph_net/network_regularizers/flop_regularizer_test.py

@@ -49,18 +49,26 @@ def BuildModel(self):
     #
     # (the model has two "outputs", conv3 and conv4).
     #
+
+    # op.name: 'Const'
     image = tf.constant(0.0, shape=[1, 17, 19, NUM_CHANNELS])
-    conv1 = slim.layers.conv2d(image, 13, [7, 5], padding='SAME', scope='conv1')
-    conv2 = slim.layers.conv2d(image, 23, [1, 1], padding='SAME', scope='conv2')
-    concat = tf.concat([conv1, conv2], 3)
+    # op.name: 'conv1/Conv2D'
+    self.conv1 = slim.layers.conv2d(
+        image, 13, [7, 5], padding='SAME', scope='conv1')
+    self.conv2 = slim.layers.conv2d(
+        image, 23, [1, 1], padding='SAME', scope='conv2')
+    self.concat = tf.concat([self.conv1, self.conv2], 3)
     self.conv3 = slim.layers.conv2d(
-        concat, 29, [3, 3], stride=2, padding='SAME', scope='conv3')
+        self.concat, 29, [3, 3], stride=2, padding='SAME', scope='conv3')
     self.conv4 = slim.layers.conv2d(
-        concat, 31, [1, 1], stride=1, padding='SAME', scope='conv4')
+        self.concat, 31, [1, 1], stride=1, padding='SAME', scope='conv4')
     self.name_to_var = {v.op.name: v for v in tf.global_variables()}
 
+  def AddRegularizer(self, input_boundary=None):
     self.gamma_flop_reg = flop_regularizer.GammaFlopsRegularizer(
-        [self.conv3.op, self.conv4.op], gamma_threshold=0.45)
+        [self.conv3.op, self.conv4.op],
+        gamma_threshold=0.45,
+        input_boundary=input_boundary)
 
   def GetConv(self, name):
     return tf.get_default_graph().get_operation_by_name(name + '/Conv2D')
@@ -84,8 +92,18 @@ def GetLoss(self, conv):
     with self.cached_session():
       return self.gamma_flop_reg.get_regularization_term(conv).eval()
 
-  def testCost(self,):
+  def GetSourceOps(self):
+    op_regularizer_manager = self.gamma_flop_reg.op_regularizer_manager
+    return [
+        op.name
+        for op in op_regularizer_manager.ops
+        if op_regularizer_manager.is_source_op(op)
+    ]
+
+  def testCost(self):
     self.BuildWithBatchNorm(fused=True)
+    self.AddRegularizer(input_boundary=None)
+
     # Conv1 has 7 gammas above 0.45, and NUM_CHANNELS inputs (from the image).
     conv = self.GetConv('conv1')
     self.assertEqual(_coeff(conv) * 7 * NUM_CHANNELS, self.GetCost([conv]))
@@ -107,8 +125,40 @@ def testCost(self,):
     self.assertEqual(
         self.GetCost(convs[:1]) + self.GetCost(convs[1:]), self.GetCost(convs))
 
+  def testInputBoundaryNone(self):
+    self.BuildWithBatchNorm(fused=True)
+    self.AddRegularizer(input_boundary=None)
+    self.assertCountEqual(self.GetSourceOps(), [
+        'conv1/BatchNorm/FusedBatchNorm', 'conv2/BatchNorm/FusedBatchNorm',
+        'conv3/BatchNorm/FusedBatchNorm', 'conv4/BatchNorm/FusedBatchNorm'
+    ])
+
+  def testInputBoundaryConv3(self):
+    # Only block one path, can still reach all other convolutions.
+    self.BuildWithBatchNorm(fused=True)
+    self.AddRegularizer(input_boundary=[self.conv3.op])
+    self.assertCountEqual(self.GetSourceOps(), [
+        'conv1/BatchNorm/FusedBatchNorm', 'conv2/BatchNorm/FusedBatchNorm',
+        'conv4/BatchNorm/FusedBatchNorm'
+    ])
+
+  def testInputBoundaryConv3And4(self):
+    # Block both paths, can no longer reach Concat and earlier convolutions.
+    self.BuildWithBatchNorm(fused=True)
+    self.AddRegularizer(input_boundary=[self.conv3.op, self.conv4.op])
+    self.assertCountEqual(self.GetSourceOps(), [])
+
+  def testInputBoundaryConcat(self):
+    # Block concat, can only see conv3 and conv4.
+    self.BuildWithBatchNorm(fused=True)
+    self.AddRegularizer(input_boundary=[self.concat.op])
+    self.assertCountEqual(
+        self.GetSourceOps(),
+        ['conv3/BatchNorm/FusedBatchNorm', 'conv4/BatchNorm/FusedBatchNorm'])
+
   def testLossDecorated(self):
     self.BuildWithBatchNorm(True)
+    self.AddRegularizer()
     # Create network regularizer with DummyDecorator op regularization.
     self.gamma_flop_reg = flop_regularizer.GammaFlopsRegularizer(
         [self.conv3.op, self.conv4.op],

morph_net/network_regularizers/latency_regularizer.py

@@ -26,6 +26,7 @@ def __init__(
       batch_size=1,
       regularizer_decorator: Type[generic_regularizers.OpRegularizer] = None,
       decorator_parameters=None,
+      input_boundary=None,
       force_group=None,
       regularizer_blacklist=None) -> None:
     """Creates a GammaLatencyRegularizer object.
@@ -49,6 +50,8 @@ def __init__(
       decorator_parameters: A dictionary of parameters to pass to the decorator
         factory. To be used only with decorators that requires parameters,
         otherwise use None.
+      input_boundary: A list of ops that represent the input boundary of the
+        subgraph being regularized (input boundary is not regularized).
       force_group: List of regex for ops that should be force-grouped.  Each
         regex corresponds to a separate group.  Use '|' operator to specify
         multiple patterns in a single regex. See op_regularizer_manager for
@@ -69,7 +72,7 @@ def __init__(
     })
 
     self._manager = orm.OpRegularizerManager(
-        ops, op_handler_dict,
+        ops, op_handler_dict, input_boundary=input_boundary,
         force_group=force_group, regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
         self._manager,

morph_net/network_regularizers/model_size_regularizer.py

@@ -24,6 +24,7 @@ def __init__(
       gamma_threshold,
       regularizer_decorator: Type[generic_regularizers.OpRegularizer] = None,
       decorator_parameters=None,
+      input_boundary=None,
       force_group=None,
       regularizer_blacklist=None):
     """Creates a GammaModelSizeRegularizer object.
@@ -41,6 +42,8 @@ def __init__(
       decorator_parameters: A dictionary of parameters to pass to the decorator
         factory. To be used only with decorators that requires parameters,
         otherwise use None.
+      input_boundary: A list of ops that represent the input boundary of the
+        subgraph being regularized (input boundary is not regularized).
       force_group: List of regex for ops that should be force-grouped.  Each
         regex corresponds to a separate group.  Use '|' operator to specify
         multiple patterns in a single regex. See op_regularizer_manager for
@@ -60,7 +63,7 @@ def __init__(
     })
 
     self._manager = orm.OpRegularizerManager(
-        ops, op_handler_dict,
+        ops, op_handler_dict, input_boundary=input_boundary,
         force_group=force_group, regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
         self._manager, resource_function.model_size_function)