text: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  
  # Copyright (c) 2018, NVIDIA CORPORATION. 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.
  
  import sys
  
  import tensorflow as tf
  import horovod.tensorflow as hvd
  
  from utils import image_processing
  from utils import hvd_utils
  from utils import dali_utils
  
  __all__ = ["get_synth_input_fn", "normalized_inputs"]
  
  _R_MEAN = 123.68
  _G_MEAN = 116.78
  _B_MEAN = 103.94
  _CHANNEL_MEANS = [_R_MEAN, _G_MEAN, _B_MEAN]
  
  _NUM_CHANNELS = 3
  
  
  def get_synth_input_fn(batch_size, height, width, num_channels, data_format, num_classes, dtype=tf.float32):
      """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
      tunning 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.
  
      Returns:
          An input_fn that can be used in place of a real one to return a dataset
          that can be used for iteration.
      """
  
      if data_format not in ["NHWC", "NCHW"]:
          raise ValueError("Unknown data_format: %s" % str(data_format))
  
      if data_format == "NHWC":
          input_shape = [batch_size, height, width, num_channels]
      else:
          input_shape = [batch_size, num_channels, height, width]
  
      # Convert the inputs to a Dataset.
      inputs = tf.truncated_normal(input_shape, dtype=dtype, mean=127, stddev=60, name='synthetic_inputs')
      labels = tf.random_uniform([batch_size], minval=0, maxval=num_classes - 1, dtype=tf.int32, name='synthetic_labels')
  
      data = tf.data.Dataset.from_tensors((inputs, labels))
  
      data = data.repeat()
  
      data = data.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
  
      return data
  
  
  def get_tfrecords_input_fn(filenames, batch_size, height, width, training, distort_color, num_threads, deterministic):
  
      shuffle_buffer_size = 4096
  
      if deterministic:
          if hvd_utils.is_using_hvd():
              seed = 13 * (1 + hvd.rank())
          else:
              seed = 13
      else:
          seed = None
  
      ds = tf.data.Dataset.from_tensor_slices(filenames)
  
      if hvd_utils.is_using_hvd() and training:
          ds = ds.shard(hvd.size(), hvd.rank())
  
      ds = ds.apply(
          tf.data.experimental.parallel_interleave(
              tf.data.TFRecordDataset,
              cycle_length=10,
              block_length=8,
              sloppy=not deterministic,
              prefetch_input_elements=16
          )
      )
  
      counter = tf.data.Dataset.range(sys.maxsize)
      ds = tf.data.Dataset.zip((ds, counter))
  
      def preproc_func(record, counter_):
          return image_processing.preprocess_image_record(record, height, width, _NUM_CHANNELS, training)
  
      ds = ds.cache()
      
      if training:
  
          ds = ds.apply(tf.data.experimental.shuffle_and_repeat(buffer_size=shuffle_buffer_size, seed=seed))
  
      else:
          ds = ds.repeat()
  
      ds = ds.apply(
          tf.data.experimental.map_and_batch(
              map_func=preproc_func,
              num_parallel_calls=num_threads,
              batch_size=batch_size,
              drop_remainder=True,
          )
      )
  
      ds = ds.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
  
      return ds
  
  def get_inference_input_fn(filenames, height, width, num_threads):
      
      ds = tf.data.Dataset.from_tensor_slices(filenames)
  
      counter = tf.data.Dataset.range(sys.maxsize)
      ds = tf.data.Dataset.zip((ds, counter))
  
      def preproc_func(record, counter_):
          return image_processing.preprocess_image_file(record, height, width, _NUM_CHANNELS, is_training=False)
      
      ds = ds.apply(
          tf.data.experimental.map_and_batch(
              map_func=preproc_func,
              num_parallel_calls=num_threads,
              batch_size=1
          )
      )
  
      ds = ds.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
      
      return ds
  
      
      
  def get_dali_input_fn(filenames, idx_filenames, batch_size, height, width, training, distort_color, num_threads, deterministic):
  
      if idx_filenames is None:
          raise ValueError("Must provide idx_filenames for DALI's reader")
          
      preprocessor = dali_utils.DALIPreprocessor(
          filenames,
          idx_filenames,
          height, width,
          batch_size,
          num_threads,
          dali_cpu=False,
          deterministic=deterministic,
          training=training)
      
      images, labels = preprocessor.get_device_minibatches()
      
      return (images, labels)
  
  
  def normalized_inputs(inputs):
  
      num_channels = inputs.get_shape()[-1]
  
      if inputs.get_shape().ndims != 4:
          raise ValueError('Input must be of size [batch_size, height, width, C>0]')
  
      if len(_CHANNEL_MEANS) != num_channels:
          raise ValueError('len(means) must match the number of channels')
  
      # We have a 1-D tensor of means; convert to 3-D.
      means_per_channel = tf.reshape(_CHANNEL_MEANS, [1, 1, num_channels])
      means_per_channel = tf.cast(means_per_channel, dtype=inputs.dtype)
  
      inputs = tf.subtract(inputs, means_per_channel)
  
      return tf.divide(inputs, 255.0)
  
  def get_serving_input_receiver_fn(batch_size, height, width, num_channels, data_format, dtype=tf.float32):
      
      if data_format not in ["NHWC", "NCHW"]:
          raise ValueError("Unknown data_format: %s" % str(data_format))
  
      if data_format == "NHWC":
          input_shape = [batch_size] + [height, width, num_channels]
      else:
          input_shape = [batch_size] + [num_channels, height, width]
          
      def serving_input_receiver_fn():
          features = tf.placeholder(dtype=dtype, shape=input_shape, name='input_tensor')
          return tf.estimator.export.TensorServingInputReceiver(features=features, receiver_tensors=features)
      
      return serving_input_receiver_fn
  