Latency regularizer based on the LogisticSigmoid approach outlined in

https://arxiv.org/abs/2006.09581

PiperOrigin-RevId: 331889063

Author: yairmov

morph_net/network_regularizers/BUILD

@@ -113,6 +113,7 @@ pytype_library(
     srcs = ["latency_regularizer.py"],
     deps = [
         ":cost_calculator",
+        ":logistic_sigmoid_regularizer",
         ":resource_function",
         "//third_party/py/morph_net/framework:batch_norm_source_op_handler",
         "//third_party/py/morph_net/framework:conv2d_transpose_source_op_handler",

morph_net/network_regularizers/latency_regularizer.py

@@ -1,9 +1,6 @@
 """A NetworkRegularizer that targets inference latency."""
 
-from __future__ import absolute_import
-from __future__ import division
-# [internal] enable type annotations
-from __future__ import print_function
+from typing import Type, List
 
 from morph_net.framework import batch_norm_source_op_handler
 from morph_net.framework import conv2d_transpose_source_op_handler as conv2d_transpose_handler
@@ -15,9 +12,86 @@
 from morph_net.framework import op_handlers
 from morph_net.framework import op_regularizer_manager as orm
 from morph_net.network_regularizers import cost_calculator
+from morph_net.network_regularizers import logistic_sigmoid_regularizer
 from morph_net.network_regularizers import resource_function
 import tensorflow.compat.v1 as tf
-from typing import Type, List
+
+
+class LogisticSigmoidLatencyRegularizer(
+    logistic_sigmoid_regularizer.LogisticSigmoidRegularizer):
+  """A LogisticSigmoidRegularizer that targets Latency.
+
+    Args:
+      output_boundary: An OpRegularizer will be created for all these
+        operations, and recursively for all ops they depend on via data
+        dependency that does not involve ops from input_boundary.
+      batch_size: Integer batch size to calculate cost/loss for.
+      regularize_on_mask: Bool. If True uses the binary mask as the
+        regularization vector. Else uses the probability vector.
+      alive_threshold: Float. Threshold below which values are considered dead.
+        This can be used both when mask_as_alive_vector is True and then the
+        threshold is used to binarize the sampled values and
+        when mask_as_alive_vector is False, and then the threshold is on the
+        channel probability.
+      mask_as_alive_vector: Bool. If True use the thresholded sampled mask
+        as the alive vector. Else, use thresholded probabilities from the
+        logits.
+      regularizer_decorator: A string, the name of the regularizer decorators to
+        use. Supported decorators are listed in
+        op_regularizer_decorator.SUPPORTED_DECORATORS.
+      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
+        detail.
+      regularizer_blacklist: List of regex for ops that should not be
+        regularized. See op_regularizer_manager for more detail.
+  """
+
+  def __init__(
+      self,
+      output_boundary: List[tf.Operation],
+      hardware,
+      batch_size=1,
+      regularize_on_mask=True,
+      alive_threshold=0.1,
+      mask_as_alive_vector=True,
+      regularizer_decorator: Type[generic_regularizers.OpRegularizer] = None,
+      decorator_parameters=None,
+      input_boundary: List[tf.Operation] = None,
+      force_group=None,
+      regularizer_blacklist=None):
+
+    self._hardware = hardware
+    self._batch_size = batch_size
+
+    super().__init__(
+        output_boundary=output_boundary,
+        regularize_on_mask=regularize_on_mask,
+        alive_threshold=alive_threshold,
+        mask_as_alive_vector=mask_as_alive_vector,
+        regularizer_decorator=regularizer_decorator,
+        decorator_parameters=decorator_parameters,
+        input_boundary=input_boundary,
+        force_group=force_group,
+        regularizer_blacklist=regularizer_blacklist)
+
+  def get_calculator(self):
+    return cost_calculator.CostCalculator(
+        self._manager, resource_function.latency_function_factory(
+            self._hardware, self._batch_size))
+
+  @property
+  def name(self):
+    return 'LogisticSigmoidLatency'
+
+  @property
+  def cost_name(self):
+    return self._hardware + ' Latency'
 
 
 class GammaLatencyRegularizer(generic_regularizers.NetworkRegularizer):