import tensorflow as tf

          num_shards=FLAGS.num_shards,
          per_host_input_for_training=tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2
      ))
  return run_config


def build_image_serving_input_receiver_fn(shape,
                                          dtype=tf.float32):
  """Returns a input_receiver_fn for raw images during serving."""

  def _preprocess_image(encoded_image):
    """Preprocess a single raw image."""
    image = tf.image.decode_image(encoded_image, channels=shape[-1])
    image.set_shape(shape)
    return tf.cast(image, dtype)

  def serving_input_receiver_fn():
    image_bytes_list = tf.placeholder(
        shape=[None],
        dtype=tf.string,
    )
    images = tf.map_fn(
        _preprocess_image, image_bytes_list, back_prop=False, dtype=dtype)
    return tf.estimator.export.TensorServingInputReceiver(
        features=images, receiver_tensors=image_bytes_list)

  return serving_input_receiver_fn

