• Search by Module
• Search by Words
• Search Projects

• Java
• Python
• JavaScript
• TypeScript
• C++
• Scala
• Blog

• FastNN-master
  • images
    • datasets
      • cifar10.py
      • mnist.py
      • download_and_covert_data
        • download_and_convert_data.py
        • download_and_convert_cifar10.py
        • __init__.py
        • download_and_convert_mnist.py
        • utils.py
        • download_and_convert_flowers.py
      • flowers.py
      • __init__.py
      • dataset_factory.py
    • utils
      • dataset_utils.py
      • hooks_utils.py
      • cluster_utils.py
      • __init__.py
      • misc_utils.py
      • optimizer_utils.py
    • models
      • model_factory.py
      • overfeat
        • overfeat.py
        • __init__.py
      • lenet
        • lenet.py
        • __init__.py
      • vgg
        • vgg.py
        • __init__.py
      • pix2pix
        • pix2pix.py
        • __init__.py
      • inceptions
        • inception_utils.py
        • inception_v3.py
        • inception_v1.py
        • inception_resnet_v2.py
        • inception_v2.py
        • __init__.py
        • inception_v4.py
        • inception.py
      • resnet
        • resnet_v1.py
        • resnet_utils.py
        • __init__.py
        • resnet_v2.py
      • dcgan
        • dcgan.py
        • __init__.py
      • cifarnet
        • cifarnet.py
        • __init__.py
      • mobilenet
        • conv_blocks.py
        • mobilenet.py
        • mobilenet_v1.py
        • mobilenet_v2.py
        • __init__.py
      • __init__.py
      • alexnet
        • alexnet.py
        • __init__.py
      • nasnet
        • pnasnet.py
        • nasnet_utils.py
        • nasnet.py
        • __init__.py
    • flags.py
    • train_image_classifiers.py
    • __init__.py
    • preprocessing
      • vgg_preprocessing.py
      • cifarnet_preprocessing.py
      • inception_preprocessing.py
      • lenet_preprocessing.py
      • __init__.py
      • preprocessing_factory.py
  • README_en.md
  • LICENSE
  • README.md

# Copyright (C) 2016 The TensorFlow Authors.
# Copyright (C) 2019 Alibaba Group Holding Limited.
# All Rights Reserved.
# ==============================================================================
"""A factory-pattern map which returns classification dataset iterator."""
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import cifar10
import flowers
import mnist
from flags import FLAGS
from utils import dataset_utils

datasets_map = {
    'cifar10': cifar10,
    'flowers': flowers,
    'mnist': mnist,
}


def get_dataset_iterator(dataset_name, train_image_size, preprocessing_fn=None, data_sources=None, reader=None):
  with tf.device("/cpu:0"):
    if not dataset_name:
      raise ValueError('expect dataset_name not None.')
    if dataset_name == 'mock':
      return dataset_utils._create_mock_iterator(train_image_size)
    if dataset_name not in datasets_map:
      raise ValueError('Name of network unknown %s' % dataset_name)

    def parse_fn(example):
      with tf.device("/cpu:0"):
        image, label = datasets_map[dataset_name].parse_fn(example)
        if preprocessing_fn is not None:
          image = preprocessing_fn(image, train_image_size, train_image_size)
        if FLAGS.use_fp16:
          image = tf.cast(image, tf.float16)
        label -= FLAGS.labels_offset
        return image, label
    return dataset_utils._create_dataset_iterator(data_sources, parse_fn, reader)