来自生成器的张量流数据集

Question

I am using such code to recursively load images from directory and get associated labels - directory names. But when I have more images, it crashes due memory error. I would like to use generator, but I am really stuck with it. Could somebody help? The code without generator is:

import pathlib
import random

data_dir = "./images"

print(data_dir)
data_root = pathlib.Path(data_dir)
all_image_paths = list(data_root.glob('*/*'))
all_image_paths = [str(path) for path in all_image_paths]

label_names = sorted(item.name for item in data_root.glob('*/') if item.is_dir())
label_to_index = dict((name, index) for index,name in enumerate(label_names))

all_image_labels = [label_to_index[pathlib.Path(path).parent.name] for path in all_image_paths]

path_ds = tf.data.Dataset.from_tensor_slices(all_image_paths)
image_ds = path_ds.map(load_and_preprocess_image, num_parallel_calls=8)
label_ds = tf.data.Dataset.from_tensor_slices(tf.cast(all_image_labels, tf.int64))
dataset = tf.data.Dataset.zip((image_ds, label_ds))
dataset = dataset.shuffle(params.train_size)buffer
dataset = dataset.repeat(params.num_epochs)
dataset = dataset.batch(params.batch_size)
dataset = dataset.prefetch(1) to serve

return dataset

Answer 1

You don't need to use tf.data.Dataset.from_generator . Creating dataset from images with tf.data.Dataset.from_tensor_slices writes data chunks in the graph as tf.constant(), wasting memory. With a large enough dataset you can hit Tensorflow's 2GB GraphDef limit. You just need to define dataset using placeholders

features_placeholder = tf.placeholder(features.dtype, features.shape)
labels_placeholder = tf.placeholder(labels.dtype, labels.shape)

dataset = tf.data.Dataset.from_tensor_slices((features_placeholder, labels_placeholder))

As explained here https://www.tensorflow.org/guide/datasets#consuming_numpy_arrays