data_utils.py

#!/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

from utils import image_processing
from utils import dali_utils
from utils import hvd_wrapper as hvd

__all__ = ["get_synth_input_fn", "normalized_inputs"]

_R_MEAN = 123.68
_G_MEAN = 116.28
_B_MEAN = 103.53
_CHANNEL_MEANS = [_R_MEAN, _G_MEAN, _B_MEAN]
_CHANNEL_STDS = [58.395, 57.120, 57.385]

_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.data.experimental.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:
        seed = 13 * hvd.rank()
    else:
        seed = None

    ds = tf.data.Dataset.from_tensor_slices(filenames)

    if hvd.size() > 1 and training:
        ds = ds.shard(hvd.size(), hvd.rank())

    ds = ds.interleave(tf.data.TFRecordDataset, cycle_length=10, block_length=8)

    def preproc_func(record):
        return image_processing.preprocess_image_record(record, height, width, _NUM_CHANNELS, training)

    if training:
        ds = ds.shuffle(buffer_size=shuffle_buffer_size, seed=seed)

    ds = ds.repeat().map(preproc_func, num_parallel_calls=num_threads)
    ds = ds.batch(batch_size=batch_size, drop_remainder=True)
    ds = ds.prefetch(buffer_size=tf.data.experimental.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.data.experimental.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)

    stds_per_channel = tf.reshape(_CHANNEL_STDS, [1, 1, num_channels])
    stds_per_channel = tf.cast(stds_per_channel, dtype=inputs.dtype)

    inputs = tf.subtract(inputs, means_per_channel)
    return tf.divide(inputs, stds_per_channel)


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