Merge pull request #4334 from gariel-google/master

Added the MorphNet library

research/morph_net/network_regularizers/model_size_regularizer.py

+# Copyright 2018 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.
+# ==============================================================================
+"""A NetworkRegularizer that targets the number of weights in the model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from morph_net.framework import op_regularizer_manager
+from morph_net.network_regularizers import bilinear_cost_utils
+from morph_net.op_regularizers import gamma_l1_regularizer
+
+
+# TODO: Add unit tests. This class is very similar to
+# GammaFlopsRegularizer in flop_regularizer, so we have some indirect testing,
+# but more is needed.
+class GammaModelSizeRegularizer(
+    bilinear_cost_utils.BilinearNetworkRegularizer):
+  """A NetworkRegularizer that targets model size using Gamma L1 OpReg."""
+
+  def __init__(self, ops, gamma_threshold):
+    gamma_l1_reg_factory = gamma_l1_regularizer.GammaL1RegularizerFactory(
+        gamma_threshold)
+    opreg_manager = op_regularizer_manager.OpRegularizerManager(
+        ops, {'Conv2D': gamma_l1_reg_factory.create_regularizer,
+              'DepthwiseConv2dNative': gamma_l1_reg_factory.create_regularizer})
+    super(GammaModelSizeRegularizer, self).__init__(
+        opreg_manager, bilinear_cost_utils.num_weights_coeff)

research/morph_net/op_regularizers/conv_group_lasso_regularizer.py

+# Copyright 2018 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.
+# ==============================================================================
+"""A regularizer for convolutions, based on group-lasso.
+
+All the weights that are related to a single output are grouped into one LASSO
+group (https://arxiv.org/pdf/1611.06321.pdf).
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from morph_net.framework import generic_regularizers
+
+
+class ConvGroupLassoRegularizer(generic_regularizers.OpRegularizer):
+  """A regularizer for convolutions, based on group-lasso.
+
+  Supported ops: Conv2D and Conv2DBackpropInput (transposed Conv2D).
+  are supported. The grouping is done according to the formula:
+
+  (1 - l1_fraction) * L2(weights) / sqrt(dim) + l1_fraction * L1(weights) / dim,
+
+  where `dim` is the number of weights associated with an activation, L2 and L1
+  are the respective norms, and l1_fraction controls the balance between L1 and
+  L2 grouping. The paper cited above experiments with 0.0 and 0.5 for
+  l1_fraction.
+  """
+
+  def __init__(self, op, threshold, l1_fraction=0.0):
+    """Creates an instance.
+
+    Args:
+      op: A tf.Operation object of type Conv2D or Conv2DBackpropInput.
+      threshold: A float. When the norm of the group associated with an
+        activation is below the threshold, it will be considered dead.
+      l1_fraction: A float, controls the balance between L1 and L2 grouping
+        (see above).
+
+    Raises:
+      ValueError: `op` is not of type 'Conv2D' or 'Conv2DBackpropInput', or
+        l1_fraction is outside interval [0.0, 1.0].
+    """
+    if op.type not in ('Conv2D', 'Conv2DBackpropInput'):
+      raise ValueError('The given op is not Conv2D or Conv2DBackpropInput.')
+    if l1_fraction < 0.0 or l1_fraction > 1.0:
+      raise ValueError(
+          'l1_fraction should be in [0.0, 1.0], not %e.' % l1_fraction)
+
+    self._threshold = threshold
+    conv_weights = op.inputs[1]
+    # For a Conv2D (Conv2DBackpropInput) the output dimension of the weight
+    # matrix is 3 (2). We thus reduce over all other dimensions.
+
+    l2_norm = tf.sqrt(
+        tf.reduce_mean(tf.square(conv_weights), axis=_get_reduce_dims(op)))
+    if l1_fraction > 0.0:
+      l1_norm = tf.reduce_mean(tf.abs(conv_weights), axis=_get_reduce_dims(op))
+      norm = l1_fraction * l1_norm + (1.0 - l1_fraction) * l2_norm
+    else:
+      norm = l2_norm
+    # Sanity check: Output dimension of 'op' should match that of 'norm':
+    assert op.outputs[0].shape.ndims == 4
+    assert norm.shape.ndims == 1
+    op.outputs[0].shape.dims[3].assert_is_compatible_with(norm.shape.dims[0])
+    self._regularization_vector = norm
+    self._alive_vector = norm > threshold
+
+  @property
+  def regularization_vector(self):
+    return self._regularization_vector
+
+  @property
+  def alive_vector(self):
+    return self._alive_vector
+
+
+class ConvGroupLassoRegularizerFactory(object):
+  """A class for creating a ConvGroupLassoRegularizer for convolutions."""
+
+  def __init__(self, threshold, l1_fraction=0.0):
+    """Creates an instance.
+
+    Args:
+      threshold: A float scalar, will be used as a threshold for all
+        ConvGroupLassoRegularizer-s created by this class.
+      l1_fraction: A float scalar, will be passed as l1_fraction to all
+        ConvGroupLassoRegularizer-s created by this class.
+    """
+    self._threshold = threshold
+    self._l1_fraction = l1_fraction
+
+  def create_regularizer(self, op, opreg_manager=None):
+    """Creates a ConvGroupLassoRegularizer for `op`.
+
+    Args:
+      op: A tf.Operation of type 'Conv2D'.
+      opreg_manager: unused
+
+    Returns:
+      a ConvGroupLassoRegularizer that corresponds to `op`.
+    """
+    del opreg_manager  # unused
+    return ConvGroupLassoRegularizer(op, self._threshold, self._l1_fraction)
+
+
+def _get_reduce_dims(op):
+  """Returns the reduction dimensions for grouping weights of various ops."""
+  type_to_dims = {'Conv2D': (0, 1, 2), 'Conv2DBackpropInput': (0, 1, 3)}
+  try:
+    return type_to_dims[op.type]
+  except KeyError:
+    raise ValueError('Reduce dims are unknown for op type %s' % op.type)

research/morph_net/op_regularizers/conv_group_lasso_regularizer_test.py

+# Copyright 2018 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 op_regularizers.conv_group_lasso_regularizer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+from morph_net.op_regularizers import conv_group_lasso_regularizer
+
+layers = tf.contrib.layers
+
+ALIVE_THRESHOLD = 1.0
+
+
+def assert_not_all_are_alive_or_dead(alive_vector):
+  assert not all(alive_vector), (
+      'All activations are alive, test case is trivial. Increase threshold')
+  assert any(alive_vector), (
+      'All activations are dead, test case is trivial. Decrease threshold')
+
+
+class GroupLassoRegularizerTest(parameterized.TestCase, tf.test.TestCase):
+
+  def setUp(self):
+    tf.reset_default_graph()
+    tf.set_random_seed(7907)
+    with tf.contrib.framework.arg_scope(
+        [layers.conv2d, layers.conv2d_transpose],
+        weights_initializer=tf.random_normal_initializer):
+      self.BuildModel()
+    with self.test_session():
+      tf.global_variables_initializer().run()
+
+  def BuildModel(self):
+    image = tf.constant(0.0, shape=[1, 17, 19, 3])
+    conv = layers.conv2d(image, 13, [7, 5], padding='SAME', scope='conv')
+    layers.conv2d_transpose(conv, 11, [5, 5], scope='convt')
+
+  # For Conv2D (Conv2DBackpropInput, aka conv2d transpose), the reduction
+  # indices for group lasso are (0, 1, 2) ((0, 1, 3)).
+  @parameterized.named_parameters(
+      ('_regular_conv', 'conv/Conv2D', (0, 1, 2), 0.0),
+      ('_transpose_conv', 'convt/conv2d_transpose', (0, 1, 3), 0.0),
+      ('_regular_conv_l10.5', 'conv/Conv2D', (0, 1, 2), 0.5))
+  def testOp(self, op_name, axis, l1_fraction):
+    op = tf.get_default_graph().get_operation_by_name(op_name)
+    with self.test_session():
+      weights = op.inputs[1].eval()
+    l1_reg_vector = np.mean(np.abs(weights), axis=axis)
+    l2_reg_vector = np.sqrt(np.mean(weights**2, axis=axis))
+    expected_reg_vector = (
+        l1_fraction * l1_reg_vector + (1.0 - l1_fraction) * l2_reg_vector)
+
+    # We choose the threshold at the expectation value, so that some activations
+    # end up above threshold and others end up below. The weights are normally
+    # distributed, so the L2 norm is 1.0, and the L1 norm is sqrt(2/pi).
+    # With a general l1_fraction, we compute a weighted average of the two:
+    threshold = (1.0 - l1_fraction) + l1_fraction * np.sqrt(2 / np.pi)
+    expected_alive = expected_reg_vector > threshold
+    assert_not_all_are_alive_or_dead(expected_alive)
+
+    conv_reg = (
+        conv_group_lasso_regularizer.ConvGroupLassoRegularizer(
+            op, threshold=threshold, l1_fraction=l1_fraction))
+
+    with self.test_session():
+      actual_reg_vector = conv_reg.regularization_vector.eval()
+      actual_alive = conv_reg.alive_vector.eval()
+
+    self.assertAllClose(expected_reg_vector, actual_reg_vector)
+    self.assertAllEqual(expected_alive, actual_alive)
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/morph_net/op_regularizers/gamma_l1_regularizer.py

+# Copyright 2018 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.
+# ==============================================================================
+"""An OpRegularizer that applies L1 regularization on batch-norm gammas."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from morph_net.framework import generic_regularizers
+from morph_net.op_regularizers import gamma_mapper
+
+
+class GammaL1Regularizer(generic_regularizers.OpRegularizer):
+  """An OpRegularizer that L1-regularizes batch-norm gamma."""
+
+  def __init__(self, gamma, gamma_threshold):
+    """Creates an instance.
+
+    Args:
+      gamma: a tf.Tensor of rank 1 with the gammas.
+      gamma_threshold: A float scalar, the threshold above which a gamma is
+        considered 'alive'.
+    """
+    self._gamma = gamma
+    self._gamma_threshold = gamma_threshold
+    abs_gamma = tf.abs(gamma)
+    self._alive_vector = abs_gamma > gamma_threshold
+    self._regularization_vector = abs_gamma
+
+  @property
+  def regularization_vector(self):
+    return self._regularization_vector
+
+  @property
+  def alive_vector(self):
+    return self._alive_vector
+
+
+class GammaL1RegularizerFactory(object):
+  """A class for creating a GammaL1Regularizer for convolutions."""
+
+  def __init__(self, gamma_threshold):
+    """Creates an instance.
+
+    Args:
+      gamma_threshold: A float scalar, will be used as a 'gamma_threshold' for
+        all the GammaL1Regularizer-s created by this class.
+    """
+    self._gamma_conv_mapper = gamma_mapper.ConvGammaMapperByName()
+    self._gamma_threshold = gamma_threshold
+
+  def create_regularizer(self, op, opreg_manager):
+    """Creates a GammaL1Regularizer for `op`.
+
+    Args:
+      op: A tf.Operation of type 'Conv2D' or 'DepthwiseConv2dNative'.
+      opreg_manager: An OpRegularizerManager object that will host the created
+        OpRegularizer object.
+
+    Returns:
+      a GammaL1Regularizer that corresponds to `op`.
+
+    Raises:
+      ValueError: If `op` does not have a Gamma that corresponds to it.
+    """
+    gamma = self._gamma_conv_mapper.get_gamma(op)
+    if gamma is None:
+      regularizer = None
+    else:
+      regularizer = GammaL1Regularizer(gamma, self._gamma_threshold)
+
+    if op.type == 'DepthwiseConv2dNative':
+      regularizer = _group_depthwise_conv_regularizer(op, regularizer,
+                                                      opreg_manager)
+    return regularizer
+
+
+def _group_depthwise_conv_regularizer(op, regularizer, opreg_manager):
+  """Groups the regularizer of a depthwise convolution if needed."""
+  # If its first input doesn't have regularizers, return None. While pruning
+  # the depthwise_conv still effectively shut down input channels, fluid_net
+  # currently has not implemented a way to interpret it. In particular, the
+  # union of active channels of dw's input and output will always return all
+  # channels.
+  # TODO: update the interpretation to discover channels that are
+  # effectively pruned by dw.
+  input_reg = opreg_manager.get_regularizer(op.inputs[0].op)
+  if input_reg is None:
+    return None
+  # Check for cases where the depthwise convolution has a multiplier that
+  # is not one. Do not regularize this case.
+  # TODO: add support for depthwise with multiplier != 1.
+  if (op.inputs[0].shape.as_list()[-1] !=
+      op.outputs[0].shape.as_list()[-1]):
+    return None
+  # If the op is not regularized, return the regularizer of its input.
+  # This applies to cases in separable convolutions where the pointwise
+  # is batchnormed but the depthwise is not.
+  if regularizer is None:
+    return input_reg
+  # If both the input and depthwise have regularizers, we group them.
+  # This applies to Mobilenets where both 1x1 and depthwise have
+  # batchnorms.
+  else:
+    return opreg_manager.group_and_replace_regularizers(
+        [regularizer, input_reg])

research/morph_net/op_regularizers/gamma_mapper.py

+# Copyright 2018 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.
+# ==============================================================================
+"""Classes for mapping convolutions to their batch-norm gammas."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import abc
+import collections
+
+import tensorflow as tf
+
+from morph_net.framework import op_regularizer_manager
+
+
+class GenericConvGammaMapper(object):
+  """An interface for mapping convolutions to their batch-norm gammas."""
+
+  __metaclass__ = abc.ABCMeta
+
+  @abc.abstractmethod
+  def get_gamma(self, conv_op):
+    """Returns the BatchNorm gamma tensor associated with `conv_op`, or None.
+
+    Args:
+      conv_op: A tf.Operation of type Conv2D.
+
+    Returns:
+      A tf.Tensor containing the BatchNorm gamma associated with `conv_op`, or
+      None if `conv_op` has no BatchNorm gamma.
+
+    Raises:
+      ValueError: `conv_op` is not a tf.Operation of type `Conv2D`.
+      KeyError: `conv_op` is not in the graph that was used to construct `self`
+    """
+
+  @abc.abstractproperty
+  def all_conv_ops(self):
+    """Return all Conv2D ops that were in the graph when `self` was created."""
+    pass
+
+
+def _get_existing_variable(name):
+  """Fetches a variable by name (like tf.get_variable with reuse=True).
+
+  The reason why we can't simply use tf.get_variable with reuse=True is that
+  when variable partitioner is used, tf.get_variable requires knowing the shape
+  of the variable (even though it knows it and thus shouldn't require it). This
+  helper is a convenience function to solve this problem.
+
+  Args:
+    name: A string, the name of the variable.
+
+  Returns:
+    A tf.Tensor which is the result of convert_to_tensor of the variable, or
+    None if the variable does not exist.
+  """
+  try:
+    op = tf.get_default_graph().get_operation_by_name(name)
+  except KeyError:
+    return None
+
+  # Among all cases (partitioned variable, resource variable, or regular one),
+  # we assume that there's either a shape attribute to the op or to its output.
+  try:
+    shape = tf.TensorShape(op.get_attr('shape'))
+  except ValueError:
+    shape = op.outputs[0].shape
+
+  with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+    try:
+      # tf.Variable and tf.PartitionedVariable are not polymorphic, but
+      # both support convert_to_tensor. The result is thus always a
+      # tf.Tensor.
+      return tf.convert_to_tensor(tf.get_variable(name, shape=shape))
+    except ValueError as e:
+      if 'Variable %s does not exist' % name in str(e):
+        return None
+      else:
+        raise e  # pass through any other exceptions.
+
+
+class ConvGammaMapperByName(GenericConvGammaMapper):
+  """Maps a convolution to its BatchNorm gamma.
+
+  Assumes that the convolutions and their respective gammas conform to the
+  naming convention of tf.contrib.layers: A convolution's name ends with
+  `<BASE_NAME>/Conv2D`, and the respective batch-norm gamma ends with
+  `<BASE_NAME>/BatchNorm/gamma`
+  """
+
+  def __init__(self):
+    """Constructs an instance. Builds mapping from Conv2D ops to their Gamma."""
+    self._conv_to_gamma = {}
+    # We use get_variable under a reuse=True scope because this is a way to
+    # capture both a regular tf.Variable and a PartitionedVariable.
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      for op in tf.get_default_graph().get_operations():
+        if op.type != 'Conv2D' and op.type != 'DepthwiseConv2dNative':
+          continue
+        base_name = op.name.rsplit('/', 1)[0]
+        self._conv_to_gamma[op] = _get_existing_variable(base_name +
+                                                         '/BatchNorm/gamma')
+
+  def get_gamma(self, conv_op):
+    _raise_if_not_conv(conv_op)
+    return self._conv_to_gamma[conv_op]
+
+  @property
+  def all_conv_ops(self):
+    return self._conv_to_gamma.keys()
+
+
+class ConvGammaMapperByConnectivity(GenericConvGammaMapper):
+  """Maps a convolution to its BatchNorm gammas based on graph connectivity.
+
+  Given a batch-norm gamma, propagates along the graph to find the convolutions
+  that are batch-nomalized by this gamma. It can me more than one convolution
+  that are normalized by the same batch-norm gamma in ResNet-s, where
+  un-normalized convolutions are first summed and then their sum is normalized.
+  The converse is also true - a single convolution can be connected (through
+  residual connections) to multiple batch-norms.
+
+  Only fused batch-norm is supported: there seems to be significant variability
+  in the way non-fused batch-norm manifests in the tensorflow graph.
+  """
+
+  def __init__(self):
+    """Constructs an instance. Builds mapping from Conv2D ops to their Gamma."""
+    self._conv_to_gamma = collections.defaultdict(set)
+    for op in tf.get_default_graph().get_operations():
+      if op.type != 'FusedBatchNorm':
+        continue
+
+      convs = _dfs(op)
+      for conv in convs:
+        if conv.type == 'Conv2D':
+          self._conv_to_gamma[conv].add(op.inputs[1])  # Input #1 is gamma.
+
+    for op in tf.get_default_graph().get_operations():
+      if op.type == 'Conv2D' and op not in self._conv_to_gamma:
+        self._conv_to_gamma[op] = None
+
+  def get_gamma(self, conv_op):
+    _raise_if_not_conv(conv_op)
+    if conv_op not in self._conv_to_gamma:
+      raise KeyError
+    gammas = self._conv_to_gamma[conv_op]
+    if gammas and len(gammas) == 1:
+      # For a single element, return the element itself, to conform with
+      # ConvGammaMapperByName.
+      return list(gammas)[0]
+    return gammas
+
+  @property
+  def all_conv_ops(self):
+    return self._conv_to_gamma.keys()
+
+
+def _dfs(op, visited=None):
+  """Perform DFS on a graph.
+
+  Args:
+    op: A tf.Operation, the root node for the DFS.
+    visited: A set, used in the recursion.
+
+  Returns:
+    A list of the tf.Operations of type Conv2D that were encountered.
+  """
+  visited = visited or set()
+  ret = []
+  for child in op.inputs:
+    if child.op in visited:
+      return ret
+    visited.add(child.op)
+    if child.op.type not in op_regularizer_manager.NON_PASS_THROUGH_OPS:
+      ret.extend(_dfs(child.op, visited))
+    if child.op.type in ('Conv2D',):  # TODO: support depthwise conv.
+      ret.append(child.op)
+  return ret
+
+
+def _raise_if_not_conv(op):
+  if not isinstance(op, tf.Operation):
+    raise ValueError('conv_op must be a tf.Operation, not %s' % type(op))
+  if op.type != 'Conv2D' and op.type != 'DepthwiseConv2dNative':
+    raise ValueError('conv_op must be a Conv2D or DepthwiseConv2dNative,'
+                     'not %s' % op.type)

research/morph_net/op_regularizers/gamma_mapper_test.py

+# Copyright 2018 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 gamma_mapper."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.contrib.slim.nets import resnet_v1
+from tensorflow.contrib.slim.nets import resnet_v2
+from tensorflow.python.platform import flags
+from morph_net.op_regularizers import gamma_mapper
+
+
+FLAGS = flags.FLAGS
+
+
+layers = tf.contrib.layers
+arg_scope = tf.contrib.framework.arg_scope
+
+
+NUM_CHANNELS = 3
+
+
+def get_op(name):
+  return tf.get_default_graph().get_operation_by_name(name)
+
+
+CONV1_GAMMA = [0.1 * x for x in range(13)]
+SEP_CONV_GAMMA = [0.07 * x for x in range(23)]
+CKPT_FILE_NAME = 'ckpt'
+
+
+def build_model():
+  image = tf.constant(0.0, shape=[1, 17, 19, 3])
+  conv1 = layers.conv2d(image, 13, (3, 3), padding='SAME', scope='conv1')
+  layers.separable_conv2d(conv1, 23, (3, 3), 1, scope='sep_conv')
+
+
+def setUpModule():
+  """Save a model for later loading it.
+
+  This is the only way we're aware of for assigning values to variables
+  irrespectively of their type (regular or partitioned), since partitioned
+  variables do not support assignment.
+  """
+  with tf.Graph().as_default():
+    params = {
+        'normalizer_fn': layers.batch_norm,
+        'normalizer_params': {
+            'scale': True,
+        }
+    }
+    with tf.contrib.framework.arg_scope(
+        [layers.conv2d, layers.separable_conv2d], **params):
+      build_model()
+
+    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+      conv_gamma = tf.get_variable('conv1/BatchNorm/gamma')
+      sep_gamma = tf.get_variable('sep_conv/BatchNorm/gamma')
+    s = tf.Session()
+    s.run(tf.global_variables_initializer())
+    s.run([conv_gamma.assign(CONV1_GAMMA), sep_gamma.assign(SEP_CONV_GAMMA)])
+    saver = tf.train.Saver()
+    saver.save(s, os.path.join(FLAGS.test_tmpdir, CKPT_FILE_NAME))
+
+
+class ConvGammaMapperTest(parameterized.TestCase, tf.test.TestCase):
+
+  def createMapper(self, connectivity):
+    if connectivity:
+      return gamma_mapper.ConvGammaMapperByConnectivity()
+    return gamma_mapper.ConvGammaMapperByName()
+
+  def setUp(self):
+    tf.reset_default_graph()
+
+  def TestSuccess(self, connectivity, partitioning, fused, use_resource):
+    params = {
+        'trainable': True,
+        'normalizer_fn': layers.batch_norm,
+        'normalizer_params': {
+            'scale': True,
+            'fused': fused
+        }
+    }
+
+    partitioner = tf.fixed_size_partitioner(2) if partitioning else None
+    with tf.variable_scope(
+        tf.get_variable_scope(),
+        partitioner=partitioner,
+        use_resource=use_resource):
+      with tf.contrib.framework.arg_scope(
+          [layers.conv2d, layers.separable_conv2d], **params):
+        build_model()
+
+    sess = tf.Session()
+    saver = tf.train.Saver()
+    saver.restore(sess, os.path.join(FLAGS.test_tmpdir, CKPT_FILE_NAME))
+    mapper = self.createMapper(connectivity)
+    conv = get_op('conv1/Conv2D')
+    sep_conv = get_op('sep_conv/separable_conv2d')
+    with sess.as_default():
+      self.assertAllClose(CONV1_GAMMA, mapper.get_gamma(conv).eval())
+      self.assertAllClose(SEP_CONV_GAMMA, mapper.get_gamma(sep_conv).eval())
+
+  def testSuccess(self):
+    for connectivity in (False, True):
+      for partitioning in (False, True):
+        for fused in (False, True):
+          if connectivity and not fused:  # This combination is not supported
+            continue
+          for use_resource in (False, True):
+            tf.reset_default_graph()
+            self.TestSuccess(connectivity, partitioning, fused, use_resource)
+
+  @parameterized.named_parameters(
+      ('_name_nopart', False, False), ('_name_part', False, True),
+      ('_conn_nopart', True, False), ('_conn_part', True, True))
+  def testNoBatchNorm(self, connectivity, partitioning):
+    partitioner = tf.fixed_size_partitioner(2) if partitioning else None
+    with tf.variable_scope(
+        tf.get_variable_scope(), partitioner=partitioner):
+      build_model()
+    mapper = self.createMapper(connectivity)
+    conv = get_op('conv1/Conv2D')
+    self.assertEqual(None, mapper.get_gamma(conv))
+
+  @parameterized.named_parameters(('_name_nopart', False),
+                                  ('_conn_nopart', True))
+  def testNotAConv(self, connectivity):
+    build_model()
+    mapper = self.createMapper(connectivity)
+    bias_add = get_op('conv1/BiasAdd')
+    with self.assertRaises(ValueError):
+      mapper.get_gamma(bias_add)
+
+  @parameterized.named_parameters(('_name_nopart', False),
+                                  ('_conn_nopart', True))
+  def testNotAnOpButATensor(self, connectivity):
+    build_model()
+    mapper = self.createMapper(connectivity)
+    conv = get_op('conv1/Conv2D')
+    with self.assertRaises(ValueError):
+      mapper.get_gamma(conv.outputs[0])
+
+  @parameterized.named_parameters(('_name_nopart', False),
+                                  ('_conn_nopart', True))
+  def testNotInGraph(self, connectivity):
+    mapper = self.createMapper(connectivity)
+    # Graph is built after the mapper
+    build_model()
+    conv = get_op('conv1/Conv2D')
+    with self.assertRaises(KeyError):
+      mapper.get_gamma(conv)
+
+
+def build_resnet(block_fn, resnet_fn):
+  params = {
+      'trainable': True,
+      'normalizer_fn': layers.batch_norm,
+      'normalizer_params': {
+          'is_training': True,
+          'scale': True,
+          'fused': True
+      }
+  }
+
+  with arg_scope([layers.conv2d], **params):
+    with arg_scope([layers.batch_norm], **(params['normalizer_params'])):
+      # Each block looks like:
+      # Image --> unit_1/shortcut
+      # Image --> unit_1/conv1 --> unit_1/conv2 --> unit_1/conv3
+      #
+      # unit_1/shortcut + unit_1/conv3 --> unit_1 (residual connection)
+      #
+      # unit_1 --> unit_2/conv1  -> unit_2/conv2 --> unit_2/conv3
+      #
+      # unit_1 + unit_2/conv3 --> unit_2 (residual connection)
+      #
+      # In between, there are strided convolutions and pooling ops, but these
+      # should not affect the regularizer.
+      blocks = [
+          block_fn('block1', base_depth=7, num_units=2, stride=2),
+          block_fn('block2', base_depth=13, num_units=2, stride=2),
+      ]
+      image = tf.constant(0.0, shape=[1, 2, 2, NUM_CHANNELS])
+      return resnet_fn(
+          image, blocks, include_root_block=False, is_training=False)[0]
+
+
+class ConvGammaMapperByConnectivityResnetTest(parameterized.TestCase,
+                                              tf.test.TestCase):
+
+  def assertGammaMatchesConv(self, mapper, prefix):
+    conv = get_op(prefix + '/Conv2D')
+    gamma = mapper.get_gamma(conv)
+    self.assertTrue(gamma.op.name.startswith(prefix + '/BatchNorm/gamma'))
+
+  def assertConvsConnectedToGammas(self, conv_names, gamma_prefixes, mapper):
+    """Asserts that each convolution is connected to each gamma.
+
+    Args:
+      conv_names: A list of strings representing names of Conv2D operations.
+      gamma_prefixes: A list of strings representing name prefixes of gamma
+        variables (we only verify prefixes because suffixes may depend on
+        whether we have partitioning or no).
+      mapper: a ConvGammaMapperByConnectivity object
+    """
+    def make_set(item):
+      return item if isinstance(item, set) else set([item,])
+
+    convs = [get_op(conv_name) for conv_name in conv_names]
+    gamma_sets = [make_set(mapper.get_gamma(conv)) for conv in convs]
+    if len(gamma_sets) > 1:
+      for i in range(1, len(gamma_sets)):
+        self.assertEqual(gamma_sets[i], gamma_sets[0])
+
+    actual_gamma_names = sorted([g.op.name for g in gamma_sets[0]])
+    gamma_prefixes = sorted(gamma_prefixes)
+    for expected, actual in zip(gamma_prefixes, actual_gamma_names):
+      self.assertTrue(actual.startswith(expected))
+
+  def testSuccessResnetV2(self):
+    build_resnet(resnet_v2.resnet_v2_block, resnet_v2.resnet_v2)
+    mapper = gamma_mapper.ConvGammaMapperByConnectivity()
+    # Check all "regular" convs, that are connected to their own batch norm,
+    # without residual connecitons involved.
+    for block in (1, 2):
+      for unit in (1, 2):
+        for conv in (1, 2):
+          self.assertGammaMatchesConv(
+              mapper, 'resnet_v2/block%d/unit_%d/bottleneck_v2/conv%d' %
+              (block, unit, conv))
+
+    # This diagram depicts all the convs and the batch-norm that don't have a
+    # one to one mapping:
+    #
+    #                  CONVS                        BATCH-NORMS
+    #
+    #           block1/unit_1/shortcut --+
+    #                                    |
+    #           block1/unit_1/conv3  ----+-->  block1/unit_2/preact
+    #                                    |
+    #           block1/unit_2/conv3  ----+-->  block2/unit_1/preact
+    #
+    #
+    #           block2/unit_1/shortcut --+
+    #                                    |
+    #           block2/unit_1/conv3  ----+-->  block2/unit_1/preact
+    #                                    |
+    #           block2/unit_2/conv3  ----+-->  postnorm
+    #
+    # This connectivity is tested below.
+
+    self.assertConvsConnectedToGammas([
+        'resnet_v2/block1/unit_1/bottleneck_v2/shortcut/Conv2D',
+        'resnet_v2/block1/unit_1/bottleneck_v2/conv3/Conv2D'
+    ], [
+        'resnet_v2/block1/unit_2/bottleneck_v2/preact/gamma',
+        'resnet_v2/block2/unit_1/bottleneck_v2/preact/gamma'
+    ], mapper)
+
+    self.assertConvsConnectedToGammas([
+        'resnet_v2/block1/unit_2/bottleneck_v2/conv3/Conv2D',
+    ], [
+        'resnet_v2/block2/unit_1/bottleneck_v2/preact/gamma',
+    ], mapper)
+
+    self.assertConvsConnectedToGammas([
+        'resnet_v2/block2/unit_1/bottleneck_v2/shortcut/Conv2D',
+        'resnet_v2/block2/unit_1/bottleneck_v2/conv3/Conv2D'
+    ], [
+        'resnet_v2/block2/unit_2/bottleneck_v2/preact/gamma',
+        'resnet_v2/postnorm/gamma'
+    ], mapper)
+
+    self.assertConvsConnectedToGammas([
+        'resnet_v2/block2/unit_2/bottleneck_v2/conv3/Conv2D',
+    ], [
+        'resnet_v2/postnorm/gamma',
+    ], mapper)
+
+  def testSuccessResnetV1(self):
+    build_resnet(resnet_v1.resnet_v1_block, resnet_v1.resnet_v1)
+    mapper = gamma_mapper.ConvGammaMapperByConnectivity()
+    # Here the mapping between convolutions and batch-norms is simple one to
+    # one.
+    for block in (1, 2):
+      self.assertGammaMatchesConv(
+          mapper, 'resnet_v1/block%d/unit_1/bottleneck_v1/shortcut' % block)
+
+      for unit in (1, 2):
+        for conv in (1, 2, 3):
+          self.assertGammaMatchesConv(
+              mapper, 'resnet_v1/block%d/unit_%d/bottleneck_v1/conv%d' %
+              (block, unit, conv))
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/morph_net/testing/op_regularizer_stub.py

+# Copyright 2018 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.
+# ==============================================================================
+r"""Helpers for testing the regularizers framework.
+
+Contains:
+
+- Code to build a simple convolutional model with concatenation and a residual
+  connection.
+
+- Logic for creating Stubs for OpRegularizers for the convolutions in the model.
+
+- Helpers that calculate the expected values of the alive and regularization
+  vectors.
+
+
+The model is:
+
+             -> conv1 --+     -> conv3 -->  conv4 --
+            /           |    /                      \
+      image          [concat]                      (add) --> output
+            \           |    \                      /
+             -> conv2 --+     -> -------------------
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from morph_net.framework import generic_regularizers
+
+
+layers = tf.contrib.layers
+
+
+class OpRegularizerStub(generic_regularizers.OpRegularizer):
+  """A stub that exponses a constant regularization_vector and alive_vector."""
+
+  def __init__(self, regularization_vector, alive_vector):
+    self._regularization_vector = tf.constant(
+        regularization_vector, dtype=tf.float32)
+    self._alive_vector = tf.constant(alive_vector, dtype=tf.bool)
+
+  @property
+  def regularization_vector(self):
+    return self._regularization_vector
+
+  @property
+  def alive_vector(self):
+    return self._alive_vector
+
+
+ALIVE_STUB = {
+    'conv1': [False, True, True, False, True, False, True],
+    'conv2': [True, False, True, False, False],
+    'conv3': [False, False, True, True],
+    'conv4': [False, True, False, True, True, False, True,
+              False, False, True, False, False],
+    'conv5': [False, True, False]
+}
+
+
+REG_STUB = {
+    'conv1': [0.1, 0.3, 0.6, 0.2, 0.4, 0.0, 0.8],
+    'conv2': [0.15, 0.25, 0.05, 0.55, 0.45],
+    'conv3': [0.07, 0.27, 0.17, 0.37],
+    'conv4': [
+        0.07, 0.27, 0.17, 0.37, 0.28, 0.32, 0.12, 0.22, 0.19, 0.11, 0.02, 0.06
+    ],
+    'conv5': [0.24, 0.34, 0.29]
+}
+
+
+def _create_stub(key):
+  return OpRegularizerStub(REG_STUB[key], ALIVE_STUB[key])
+
+
+def build_model():
+  image = tf.constant(0.0, shape=[1, 17, 19, 3])
+  conv1 = layers.conv2d(image, 7, [7, 5], padding='SAME', scope='conv1')
+  conv2 = layers.conv2d(image, 5, [1, 1], padding='SAME', scope='conv2')
+  concat = tf.concat([conv1, conv2], 3)
+  conv3 = layers.conv2d(concat, 4, [1, 1], padding='SAME', scope='conv3')
+  conv4 = layers.conv2d(conv3, 12, [3, 3], padding='SAME', scope='conv4')
+  conv5 = layers.conv2d(
+      concat + conv4, 3, [3, 3], stride=2, padding='SAME', scope='conv5')
+  return conv5.op
+
+
+def _create_conv2d_regularizer(conv_op, manager=None):
+  del manager  # unused
+  for key in REG_STUB:
+    if conv_op.name.startswith(key):
+      return _create_stub(key)
+  raise ValueError('No regularizer for %s' % conv_op.name)
+
+
+MOCK_REG_DICT = {'Conv2D': _create_conv2d_regularizer}
+
+
+def expected_regularization():
+  """Build the expected alive vectors applying the rules of concat and group."""
+  concat = REG_STUB['conv1'] + REG_STUB['conv2']
+  # Grouping: Activation is alive after grouping if one of the constituents is
+  # alive.
+  grouped = [max(a, b) for a, b in zip(concat, REG_STUB['conv4'])]
+  conv1_length = len(REG_STUB['conv1'])
+  return {
+      'conv1': grouped[:conv1_length],
+      'conv2': grouped[conv1_length:],
+      'conv3': REG_STUB['conv3'],
+      'conv4': grouped,
+      'conv5': REG_STUB['conv5'],
+      'add': grouped,
+      'concat': grouped
+  }
+
+
+def expected_alive():
+  """Build the expected alive vectors applying the rules of concat and group."""
+  concat = ALIVE_STUB['conv1'] + ALIVE_STUB['conv2']
+  # Grouping: Activation is alive after grouping if one of the constituents is
+  # alive.
+  grouped = [a or b for a, b in zip(concat, ALIVE_STUB['conv4'])]
+  conv1_length = len(ALIVE_STUB['conv1'])
+  return {
+      'conv1': grouped[:conv1_length],
+      'conv2': grouped[conv1_length:],
+      'conv3': ALIVE_STUB['conv3'],
+      'conv4': grouped,
+      'conv5': ALIVE_STUB['conv5'],
+      'add': grouped,
+      'concat': grouped
+  }