Repository: tensorflow/models
File: official/vision/image_classification/resnet/common.py (Python)

Code before:
  return inputs, labels


def define_pruning_flags():
  """Define flags for pruning methods."""
  flags.DEFINE_string('pruning_method', None,
                      'Pruning method.'
                      'None (no pruning) or polynomial_decay.')
  flags.DEFINE_float('pruning_initial_sparsity', 0.0,
                     'Initial sparsity for pruning.')
  flags.DEFINE_float('pruning_final_sparsity', 0.5,
                     'Final sparsity for pruning.')
  flags.DEFINE_integer('pruning_begin_step', 0,
                       'Begin step for pruning.')
  flags.DEFINE_integer('pruning_end_step', 100000,
                       'End step for pruning.')
  flags.DEFINE_integer('pruning_frequency', 100,
                       'Frequency for pruning.')

def define_clustering_flags():
  """Define flags for clustering methods."""
  flags.DEFINE_string('clustering_method', None,
                      'None (no clustering) or selective_clustering.')
  flags.DEFINE_integer('number_of_clusters', 256,
                       'Number of clusters used in each layer.')
  flags.DEFINE_string('save_files_to', None,
                      'The path to save Keras models and tflite models.')

def get_synth_input_fn(height, width, num_channels, num_classes,
                       dtype=tf.float32, drop_remainder=True):
  """Returns an input function that returns a dataset with random data.

  This input_fn returns a data set that iterates over a set of random data and
  bypasses all preprocessing, e.g. jpeg decode and copy. The host to device
  copy is still included. This used to find the upper throughput bound when
  tuning the full input pipeline.

  Args:
    height: Integer height that will be used to create a fake image tensor.
    width: Integer width that will be used to create a fake image tensor.
    num_channels: Integer depth that will be used to create a fake image tensor.
    num_classes: Number of classes that should be represented in the fake labels
      tensor
    dtype: Data type for features/images.
    drop_remainder: A boolean indicates whether to drop the remainder of the
      batches. If True, the batch dimension will be static.

  Returns:
    An input_fn that can be used in place of a real one to return a dataset
    that can be used for iteration.
  """

Review comment (suggestion):
Don't think a new `save_files_to` flag should exist that's specific to clustering. There is existing export logic in the script already. The current code doesn't create a TFLite model and that seems fine. 

Code after:

  Args:
    height: Integer height that will be used to create a fake image tensor.
    width: Integer width that will be used to create a fake image tensor.
    num_channels: Integer depth that will be used to create a fake image tensor.
    num_classes: Number of classes that should be represented in the fake labels
      tensor
    dtype: Data type for features/images.

  Returns:
    A tuple of tensors representing the inputs and labels.

  """
  # Synthetic input should be within [0, 255].
  inputs = tf.random.truncated_normal([height, width, num_channels],
                                      dtype=dtype,
                                      mean=127,
                                      stddev=60,
                                      name='synthetic_inputs')
  labels = tf.random.uniform([1],
                             minval=0,
                             maxval=num_classes - 1,
                             dtype=tf.int32,
                             name='synthetic_labels')
  return inputs, labels


def define_pruning_flags():
  """Define flags for pruning methods."""
  flags.DEFINE_string(
      'pruning_method', None, 'Pruning method.'
      'None (no pruning) or polynomial_decay.')
  flags.DEFINE_float('pruning_initial_sparsity', 0.0,
                     'Initial sparsity for pruning.')
  flags.DEFINE_float('pruning_final_sparsity', 0.5,
                     'Final sparsity for pruning.')
  flags.DEFINE_integer('pruning_begin_step', 0, 'Begin step for pruning.')
  flags.DEFINE_integer('pruning_end_step', 100000, 'End step for pruning.')
  flags.DEFINE_integer('pruning_frequency', 100, 'Frequency for pruning.')

def define_clustering_flags():
  """Define flags for clustering methods."""
  flags.DEFINE_string('clustering_method', None,
                      'None (no clustering) or selective_clustering'\
                      '(cluster last three Conv2D layers of the model).')

def get_synth_input_fn(height,
                       width,
                       num_channels,
                       num_classes,
                       dtype=tf.float32,