python 数据集输入fn使用Tensorflow 1.4 API。

Posted 2021-05-09

def dataset_input_fn(is_training, data_dir, batch_size, num_epochs=1):
    """Input function which provides batches for train or eval.
    """
    dataset = tf.data.Dataset.from_tensor_slices(_filenames(is_training,
                                                            data_dir))

    if is_training:
        dataset = dataset.shuffle(buffer_size=_FILE_SHUFFLE_BUFFER)

    dataset = dataset.flat_map(tf.data.TFRecordDataset)
    dataset = dataset.map(lambda value: _record_parser(value, is_training),
                          num_parallel_calls=5)
    dataset = dataset.prefetch(batch_size)

    if is_training:
        # When choosing shuffle buffer sizes, larger size result in better
        # randomness, while smaller size have better performance.
        dataset = dataset.shuffle(buffer_size=_SHUFFLE_BUFFER)

    # We call repeat after shuffling, rather than before, to prevent separate
    # epochs from blending together.
    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size)

    iterator = dataset.make_one_shot_iterator()
    images, labels = iterator.get_next()

    return images, labels