text: # coding=utf-8
  # Copyright 2019 Google LLC.
  #
  # 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.
  
  """Computes Receptive Field (RF) information for different models.
  
  The receptive field (and related parameters) for the different models are
  printed to stdout, and may also optionally be written to a CSV file.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import argparse
  import csv
  import sys
  
  from absl import app
  import receptive_field as rf
  import tensorflow as tf
  import tf_slim as slim
  
  from nets import alexnet
  from nets import inception
  from nets import mobilenet_v1
  from nets import resnet_v1
  from nets import resnet_v2
  from nets import vgg
  
  cmd_args = None
  
  # Input node name for all architectures.
  _INPUT_NODE = 'input_image'
  
  # Variants of different network architectures.
  
  # - resnet: different versions and sizes.
  _SUPPORTED_RESNET_VARIANTS = [
      'resnet_v1_50', 'resnet_v1_101', 'resnet_v1_152', 'resnet_v1_200',
      'resnet_v2_50', 'resnet_v2_101', 'resnet_v2_152', 'resnet_v2_200'
  ]
  
  # - inception_resnet_v2: default, and version with SAME padding.
  _SUPPORTED_INCEPTIONRESNETV2_VARIANTS = [
      'inception_resnet_v2', 'inception_resnet_v2-same'
  ]
  
  # - inception_v2: default, and version with no separable conv.
  _SUPPORTED_INCEPTIONV2_VARIANTS = [
      'inception_v2', 'inception_v2-no-separable-conv'
  ]
  
  # - inception_v3: default version.
  _SUPPORTED_INCEPTIONV3_VARIANTS = ['inception_v3']
  
  # - inception_v4: default version.
  _SUPPORTED_INCEPTIONV4_VARIANTS = ['inception_v4']
  
  # - alexnet_v2: default version.
  _SUPPORTED_ALEXNETV2_VARIANTS = ['alexnet_v2']
  
  # - vgg: vgg_a (with 11 layers) and vgg_16 (version D).
  _SUPPORTED_VGG_VARIANTS = ['vgg_a', 'vgg_16']
  
  # - mobilenet_v1: 100% and 75%.
  _SUPPORTED_MOBILENETV1_VARIANTS = ['mobilenet_v1', 'mobilenet_v1_075']
  
  
  def _construct_model(model_type='resnet_v1_50', placeholder_resolution=None):
    """Constructs model for the desired type of CNN.
  
    Args:
      model_type: Type of model to be used.
      placeholder_resolution: Placeholder image resolution to use.
  
    Returns:
      end_points: A dictionary from components of the network to the corresponding
        activations.
  
    Raises:
      ValueError: If the model_type is not supported.
    """
    # Placeholder input.
    images = tf.placeholder(
        tf.float32,
        shape=(1, placeholder_resolution, placeholder_resolution, 3),
        name=_INPUT_NODE)
  
    # Construct model.
    if model_type == 'inception_resnet_v2':
      _, end_points = inception.inception_resnet_v2_base(images)
    elif model_type == 'inception_resnet_v2-same':
      _, end_points = inception.inception_resnet_v2_base(
          images, align_feature_maps=True)
    elif model_type == 'inception_v2':
      _, end_points = inception.inception_v2_base(images)
    elif model_type == 'inception_v2-no-separable-conv':
      _, end_points = inception.inception_v2_base(
          images, use_separable_conv=False)
    elif model_type == 'inception_v3':
      _, end_points = inception.inception_v3_base(images)
    elif model_type == 'inception_v4':
      _, end_points = inception.inception_v4_base(images)
    elif model_type == 'alexnet_v2':
      _, end_points = alexnet.alexnet_v2(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'vgg_a':
      _, end_points = vgg.vgg_a(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'vgg_16':
      _, end_points = vgg.vgg_16(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'mobilenet_v1':
      _, end_points = mobilenet_v1.mobilenet_v1_base(images)
    elif model_type == 'mobilenet_v1_075':
      _, end_points = mobilenet_v1.mobilenet_v1_base(
          images, depth_multiplier=0.75)
    elif model_type == 'resnet_v1_50':
      _, end_points = resnet_v1.resnet_v1_50(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'resnet_v1_101':
      _, end_points = resnet_v1.resnet_v1_101(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'resnet_v1_152':
      _, end_points = resnet_v1.resnet_v1_152(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'resnet_v1_200':
      _, end_points = resnet_v1.resnet_v1_200(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'resnet_v2_50':
      _, end_points = resnet_v2.resnet_v2_50(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'resnet_v2_101':
      _, end_points = resnet_v2.resnet_v2_101(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'resnet_v2_152':
      _, end_points = resnet_v2.resnet_v2_152(
          images, num_classes=None, is_training=False, global_pool=False)
    elif model_type == 'resnet_v2_200':
      _, end_points = resnet_v2.resnet_v2_200(
          images, num_classes=None, is_training=False, global_pool=False)
    else:
      raise ValueError('Unsupported model_type %s.' % model_type)
  
    return end_points
  
  
  def _get_desired_end_point_keys(model_type='resnet_v1_50'):
    """Gets list of desired end point keys for a type of CNN.
  
    Args:
      model_type: Type of model to be used.
  
    Returns:
      desired_end_point_types: A list containing the desired end-points.
  
    Raises:
      ValueError: If the model_type is not supported.
    """
    if model_type in _SUPPORTED_RESNET_VARIANTS:
      blocks = ['block1', 'block2', 'block3', 'block4']
      desired_end_point_keys = ['%s/%s' % (model_type, i) for i in blocks]
    elif model_type in _SUPPORTED_INCEPTIONRESNETV2_VARIANTS:
      desired_end_point_keys = [
          'Conv2d_1a_3x3', 'Conv2d_2a_3x3', 'Conv2d_2b_3x3', 'MaxPool_3a_3x3',
          'Conv2d_3b_1x1', 'Conv2d_4a_3x3', 'MaxPool_5a_3x3', 'Mixed_5b',
          'Mixed_6a', 'PreAuxLogits', 'Mixed_7a', 'Conv2d_7b_1x1'
      ]
    elif model_type in _SUPPORTED_INCEPTIONV2_VARIANTS:
      desired_end_point_keys = [
          'Conv2d_1a_7x7', 'MaxPool_2a_3x3', 'Conv2d_2b_1x1', 'Conv2d_2c_3x3',
          'MaxPool_3a_3x3', 'Mixed_3b', 'Mixed_3c', 'Mixed_4a', 'Mixed_4b',
          'Mixed_4c', 'Mixed_4d', 'Mixed_4e', 'Mixed_5a', 'Mixed_5b', 'Mixed_5c'
      ]
    elif model_type in _SUPPORTED_INCEPTIONV3_VARIANTS:
      desired_end_point_keys = [
          'Conv2d_1a_3x3', 'Conv2d_2a_3x3', 'Conv2d_2b_3x3', 'MaxPool_3a_3x3',
          'Conv2d_3b_1x1', 'Conv2d_4a_3x3', 'MaxPool_5a_3x3', 'Mixed_5b',
          'Mixed_5c', 'Mixed_5d', 'Mixed_6a', 'Mixed_6b', 'Mixed_6c', 'Mixed_6d',
          'Mixed_6e', 'Mixed_7a', 'Mixed_7b', 'Mixed_7c'
      ]
    elif model_type in _SUPPORTED_INCEPTIONV4_VARIANTS:
      desired_end_point_keys = [
          'Conv2d_1a_3x3', 'Conv2d_2a_3x3', 'Conv2d_2b_3x3', 'Mixed_3a',
          'Mixed_4a', 'Mixed_5a', 'Mixed_5b', 'Mixed_5c', 'Mixed_5d', 'Mixed_5e',
          'Mixed_6a', 'Mixed_6b', 'Mixed_6c', 'Mixed_6d', 'Mixed_6e', 'Mixed_6f',
          'Mixed_6g', 'Mixed_6h', 'Mixed_7a', 'Mixed_7b', 'Mixed_7c', 'Mixed_7d'
      ]
    elif model_type in _SUPPORTED_ALEXNETV2_VARIANTS:
      ep = ['conv1', 'pool1', 'conv2', 'conv3', 'conv4', 'conv5', 'pool5']
      desired_end_point_keys = ['%s/%s' % (model_type, i) for i in ep]
    elif model_type in _SUPPORTED_VGG_VARIANTS:
      ep = [
          'conv1/conv1_1', 'pool1', 'conv2/conv2_1', 'pool2', 'conv3/conv3_1',
          'conv3/conv3_2', 'pool3', 'conv4/conv4_1', 'conv4/conv4_2', 'pool4',
          'conv5/conv5_1', 'conv5/conv5_2', 'pool5'
      ]
      desired_end_point_keys = ['%s/%s' % (model_type, i) for i in ep]
    elif model_type in _SUPPORTED_MOBILENETV1_VARIANTS:
      desired_end_point_keys = [
          'Conv2d_0', 'Conv2d_1_pointwise', 'Conv2d_2_pointwise',
          'Conv2d_3_pointwise', 'Conv2d_4_pointwise', 'Conv2d_5_pointwise',
          'Conv2d_6_pointwise', 'Conv2d_7_pointwise', 'Conv2d_8_pointwise',
          'Conv2d_9_pointwise', 'Conv2d_10_pointwise', 'Conv2d_11_pointwise',
          'Conv2d_12_pointwise', 'Conv2d_13_pointwise'
      ]
    else:
      raise ValueError('Unsupported model_type %s.' % model_type)
  
    return desired_end_point_keys
  
  
  def _model_graph_def(model_type,
                       desired_end_point_keys,
                       placeholder_resolution=None,
                       arg_sc=None):
    """Constructs a model graph, returning GraphDef's and end-points.
  
    Args:
      model_type: Type of model to be used.
      desired_end_point_keys: List of desired end points for which receptive field
        information will be computed.
      placeholder_resolution: Placeholder resolution to use when constructing the
        graph.
      arg_sc: Optional arg scope to use in constructing the graph.
  
    Returns:
      graphdefs: List of GraphDef's, one per desired end point.
      end_points: A dictionary from components of the network to the corresponding
        activations.
    """
    if arg_sc is None:
      arg_sc = {}
  
    g = tf.Graph()
    sess = tf.Session(graph=g)
    with g.as_default():
      with tf.contrib.framework.arg_scope(arg_sc):
        end_points = _construct_model(model_type, placeholder_resolution)
      sess.run(tf.global_variables_initializer())
  
    # Produce a graphdef for each desired end point. While this is not required
    # for receptive field computation, it helps to provide a better estimate of
    # the number of floating point operations, since it removes initialization
    # layers.
    graphdefs = []
    for desired_end_point_key in desired_end_point_keys:
      end_point_node_name = end_points[desired_end_point_key].name.split(':')[0]
      graphdefs.append(
          tf.graph_util.convert_variables_to_constants(sess, g.as_graph_def(),
                                                       [end_point_node_name]))
  
    return graphdefs, end_points
  
  
  def _model_rf_and_flops(graphdefs,
                          end_points,
                          desired_end_point_keys,
                          model_type='resnet_v1_50',
                          csv_writer=None,
                          input_resolution=None):
    """Computes receptive field and FLOPs for a given CNN model.
  
    The information will be printed to stdout. If the RF parameters are the same
    for the horizontal and vertical directions, it will be printed only once.
    Otherwise, they are printed once for the horizontal and once for the vertical
    directions.
  
    Args:
      graphdefs: List of GraphDef's, one per desired end point.
      end_points: A dictionary from components of the model to the corresponding
        activations.
      desired_end_point_keys: List of desired end points for which receptive field
        information will be computed.
      model_type: Type of model to be used, used only for printing purposes.
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
      input_resolution: Input resolution to use when computing RF parameters. This
        is important for the case where padding can only be defined if the input
        resolution is known, which may happen if using SAME padding. This is
        assumed the resolution for both height and width. If None, we consider the
        resolution is unknown.
    """
    # Configuration of profiler. Avoid verbose output.
    profiler_options = tf.profiler.ProfileOptionBuilder.float_operation()
    profiler_options['output'] = 'file:outfile=/dev/null'
  
    for i, desired_end_point_key in enumerate(desired_end_point_keys):
      print('- %s:' % desired_end_point_key)
      output_node_with_colon = end_points[desired_end_point_key].name
      pos = output_node_with_colon.rfind(':')
      output_node = output_node_with_colon[:pos]
      try:
        # Compute receptive field parameters.
        (receptive_field_x, receptive_field_y, effective_stride_x,
         effective_stride_y, effective_padding_x,
         effective_padding_y) = rf.compute_receptive_field_from_graph_def(
             graphdefs[i],
             _INPUT_NODE,
             output_node,
             input_resolution=input_resolution)
  
        # Compute FLOPs. Can only be done if input resolution is known.
        if input_resolution is None:
          billion_flops_str = 'None'
        else:
          g = tf.Graph()
          with g.as_default():
            tf.graph_util.import_graph_def(graphdefs[i], name='')
          flops = tf.profiler.profile(g, options=profiler_options)
          billion_flops = flops.total_float_ops / 1e9
          billion_flops_str = '%.3f' % billion_flops
  
        # If values are the same in horizontal/vertical directions, just report
        # one of them. Otherwise, report both.
        if (receptive_field_x == receptive_field_y) and (
            effective_stride_x == effective_stride_y) and (
                effective_padding_x == effective_padding_y):
          print('Receptive field size = %5s, effective stride = %5s, effective '
                'padding = %5s, FLOPs (Billion) = %7s' %
                (str(receptive_field_x), str(effective_stride_x),
                 str(effective_padding_x), billion_flops_str))
        else:
          print('Receptive field size: horizontal = %5s, vertical = %5s. '
                'Effective stride: horizontal = %5s, vertical = %5s. Effective '
                'padding: horizontal = %5s, vertical = %5s, '
                'FLOPs (Billion) = %7s' %
                (str(receptive_field_x), str(receptive_field_y),
                 str(effective_stride_x), str(effective_stride_y),
                 str(effective_padding_x), str(effective_padding_y),
                 billion_flops_str))
        if csv_writer is not None:
          csv_writer.writerow({
              'CNN':
                  model_type,
              'input resolution':
                  str(input_resolution[0])
                  if input_resolution is not None else 'None',
              'end_point':
                  desired_end_point_key,
              'FLOPs (Billion)':
                  billion_flops_str,
              'RF size hor':
                  str(receptive_field_x),
              'RF size ver':
                  str(receptive_field_y),
              'effective stride hor':
                  str(effective_stride_x),
              'effective stride ver':
                  str(effective_stride_y),
              'effective padding hor':
                  str(effective_padding_x),
              'effective padding ver':
                  str(effective_padding_y)
          })
      except ValueError as e:
        print('---->ERROR: Computing RF parameters for model %s with final end '
              'point %s and input resolution %s did not work' %
              (model_type, desired_end_point_key, input_resolution))
        print('---->The returned error is: %s' % e)
        if csv_writer is not None:
          csv_writer.writerow({
              'CNN':
                  model_type,
              'input resolution':
                  str(input_resolution[0])
                  if input_resolution is not None else 'None',
              'end_point':
                  desired_end_point_key,
              'FLOPs':
                  'None',
              'RF size hor':
                  'None',
              'RF size ver':
                  'None',
              'effective stride hor':
                  'None',
              'effective stride ver':
                  'None',
              'effective padding hor':
                  'None',
              'effective padding ver':
                  'None'
          })
  
  
  def _process_model_rf_and_flops(model_type='resnet_v1_50',
                                  csv_writer=None,
                                  arg_sc=None,
                                  input_resolutions=None):
    """Contructs model graph and desired end-points, and compute RF.
  
    The computed RF parameters are printed to stdout by the _model_rf_and_flops
    function.
  
    Args:
      model_type: Type of model to be used.
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
      arg_sc: Optional arg scope to use in constructing the graph.
      input_resolutions: List of 1D input resolutions to use when computing RF
        parameters. This is important for the case where padding can only be
        defined if the input resolution is known, which may happen if using SAME
        padding. The entries in the list are assumed the resolution for both
        height and width. If one of the elements in the list is None, we consider
        it to mean that the resolution is unknown. If the list itself is None, we
        use the default list [None, 224, 321].
    """
    # Process default value for this list.
    if input_resolutions is None:
      input_resolutions = [None, 224, 321]
  
    desired_end_point_keys = _get_desired_end_point_keys(model_type)
    for n in input_resolutions:
      print('********************%s, input resolution = %s' % (model_type, n))
      graphdefs, end_points = _model_graph_def(model_type, desired_end_point_keys,
                                               n, arg_sc)
      _model_rf_and_flops(
          graphdefs,
          end_points,
          desired_end_point_keys,
          model_type,
          csv_writer,
          input_resolution=[n, n] if n is not None else None)
  
  
  def _resnet_rf(csv_writer=None):
    """Computes RF and associated parameters for resnet models.
  
    The computed values are written to stdout.
  
    Args:
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
    """
    for model_type in _SUPPORTED_RESNET_VARIANTS:
      arg_sc = resnet_v1.resnet_arg_scope()
      _process_model_rf_and_flops(model_type, csv_writer, arg_sc)
  
  
  def _inception_resnet_v2_rf(csv_writer=None):
    """Computes RF and associated parameters for the inception_resnet_v2 model.
  
    The computed values are written to stdout.
  
    Args:
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
    """
    for model_type in _SUPPORTED_INCEPTIONRESNETV2_VARIANTS:
      _process_model_rf_and_flops(model_type, csv_writer)
  
  
  def _inception_v2_rf(csv_writer=None):
    """Computes RF and associated parameters for the inception_v2 model.
  
    The computed values are written to stdout.
  
    Args:
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
    """
    for model_type in _SUPPORTED_INCEPTIONV2_VARIANTS:
      _process_model_rf_and_flops(model_type, csv_writer)
  
  
  def _inception_v3_rf(csv_writer=None):
    """Computes RF and associated parameters for the inception_v3 model.
  
    The computed values are written to stdout.
  
    Args:
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
    """
    for model_type in _SUPPORTED_INCEPTIONV3_VARIANTS:
      _process_model_rf_and_flops(model_type, csv_writer)
  
  
  def _inception_v4_rf(csv_writer=None):
    """Computes RF and associated parameters for the inception_v4 model.
  
    The computed values are written to stdout.
  
    Args:
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
    """
    for model_type in _SUPPORTED_INCEPTIONV4_VARIANTS:
      _process_model_rf_and_flops(model_type, csv_writer)
  
  
  def _alexnet_v2_rf(csv_writer=None):
    """Computes RF and associated parameters for the alexnet_v2 model.
  
    The computed values are written to stdout.
  
    Args:
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
    """
    for model_type in _SUPPORTED_ALEXNETV2_VARIANTS:
      _process_model_rf_and_flops(model_type, csv_writer)
  
  
  def _vgg_rf(csv_writer=None):
    """Computes RF and associated parameters for the vgg model.
  
    The computed values are written to stdout.
  
    Args:
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
    """
    for model_type in _SUPPORTED_VGG_VARIANTS:
      _process_model_rf_and_flops(model_type, csv_writer)
  
  
  def _mobilenet_v1_rf(csv_writer=None):
    """Computes RF and associated parameters for the mobilenet_v1 model.
  
    The computed values are written to stdout.
  
    Args:
      csv_writer: A CSV writer for RF parameters, which is used if it is not None.
    """
    for model_type in _SUPPORTED_MOBILENETV1_VARIANTS:
      with slim.arg_scope([slim.batch_norm, slim.dropout],
                          is_training=False) as arg_sc:
        _process_model_rf_and_flops(model_type, csv_writer, arg_sc)
  
  
  def main(unused_argv):
    # Configure CSV file which will be written, if desired.
    if cmd_args.csv_path:
      csv_file = open(cmd_args.csv_path, 'w')
      field_names = [
          'CNN', 'input resolution', 'end_point', 'FLOPs (Billion)',
          'RF size hor', 'RF size ver', 'effective stride hor',
          'effective stride ver', 'effective padding hor', 'effective padding ver'
      ]
      rf_writer = csv.DictWriter(csv_file, fieldnames=field_names)
      rf_writer.writeheader()
    else:
      rf_writer = None
  
    # Compute RF parameters for each network architecture.
    _alexnet_v2_rf(rf_writer)
    _vgg_rf(rf_writer)
    _inception_v2_rf(rf_writer)
    _inception_v3_rf(rf_writer)
    _inception_v4_rf(rf_writer)
    _inception_resnet_v2_rf(rf_writer)
    _mobilenet_v1_rf(rf_writer)
    _resnet_rf(rf_writer)
  
    # Close CSV file, if it was opened.
    if cmd_args.csv_path:
      csv_file.close()
  
  
  if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.register('type', 'bool', lambda v: v.lower() == 'true')
    parser.add_argument(
        '--csv_path',
        type=str,
        default='',
        help="""\
        Path to CSV file that will be written with RF parameters.If empty, no
        file will be written.\
        """)
    cmd_args, unparsed = parser.parse_known_args()
    app.run(main=main, argv=[sys.argv[0]] + unparsed)
  