CODEWITHSUNDEEP
tensorflowtensorflow-datasetstfrecord
Naz_Jnr
Naz_Jnr

Reputation: 580

TensorFlow - TF Record too large to be loaded into an np array at once

I am trying to train an AlexNet CNN model by following the steps in the tutorial from the TensorFlow guide site .However, the tutorial makes use of the below code to load in the training data

mnist = tf.contrib.learn.datasets.load_dataset("mnist")
train_data = mnist.train.images # Returns np.array
train_labels = np.asarray(mnist.train.labels, dtype=np.int32)
eval_data = mnist.test.images # Returns np.array
eval_labels = np.asarray(mnist.test.labels, dtype=np.int32)

For me, I wrote a script to write my dataset examples into a TFRecord file, and then during training, try to read these record back and feed it into the alexnet network. See code below:

#FUNCTION TO GET ALL DATASET DATA 
def _read_multiple_images(filenames, perform_shuffle=False, repeat_count=1, 
batch_size=1, available_record=39209, num_of_epochs=1):
    def _read_one_image(serialized):
        #Specify the fatures you want to extract
        features = {'image/shape': tf.FixedLenFeature([], tf.string),
            'image/class/label': tf.FixedLenFeature([], tf.int64),
            'image/class/text': tf.FixedLenFeature([], tf.string),
            'image/filename': tf.FixedLenFeature([], tf.string),
            'image/encoded': tf.FixedLenFeature([], tf.string)} 
        parsed_example = tf.parse_single_example(serialized, 
        features=features)

        #Finese extracted data
        image_raw = tf.decode_raw(parsed_example['image/encoded'], tf.uint8)
        shape = tf.decode_raw(parsed_example['image/shape'], tf.int32)
        label = tf.cast(parsed_example['image/class/label'], dtype=tf.int32)
        reshaped_img = tf.reshape(image_raw, shape)
        casted_img =  tf.cast(reshaped_img, tf.float32)
        label_tensor= [label]
        image_tensor = [casted_img]
        return label_tensor, image_tensor

complete_labels = np.array([])
complete_images = np.array([])

dataset = tf.data.TFRecordDataset(filenames=filenames)
dataset = dataset.map(_read_one_image)
dataset = dataset.repeat(repeat_count)      #Repeats dataset this # times
dataset = dataset.batch(batch_size)         #Batch size to use
iterator = dataset.make_initializable_iterator()
labels_tensor, images_tensor = iterator.get_next() #Get batch data
no_of_rounds = int(math.ceil(available_record/batch_size));

#Create tf session, get nest set of batches, and evelauate them in batches
sess = tf.Session()
 count=1
for _ in range(num_of_epochs):
  sess.run(iterator.initializer)

  while True:
    try:
      evaluated_label, evaluated_image = sess.run([labels_tensor, 
       images_tensor])

      #convert evaluated tensors to np array 
      label_np_array = np.asarray(evaluated_label, dtype=np.uint8)
      image_np_array = np.asarray(evaluated_image, dtype=np.uint8)

      #squeeze np array to make dimesnsions appropriate
      squeezed_label_np_array = label_np_array.squeeze()
      squeezed_image_np_array = image_np_array.squeeze()

      #add current batch to total
      complete_labels = np.append(complete_labels, squeezed_label_np_array)
      complete_images = np.append(complete_images, squeezed_image_np_array)
      except tf.errors.OutOfRangeError:
      print("End of Dataset Reached")
      break
    count=count+1

sess.close()
return complete_labels, complete_images

My main issue is that, while tring to recoverer all the 39209 images in my dataset (227x227x3) as np array so I can feed to my TF estimator. My computer runs out of memory. 

train_input_fn = tf.estimator.inputs.numpy_input_fn(x={"x": 
complete_images},y=complete_labels,batch_size=100,num_epochs=1, 
shuffle=True)
dataset_classifier.train(input_fn=train_input_fn,num_epochs=1,hooks=
[logging_hook])

Is there a way I can get my images and labels out of my TF record in batches and feed it to my TF.Estimator in batches rather than having to need to load it all into an np array as specified in this tutorial

Upvotes: 0

Views: 1461

Answers (1)

mrry
mrry

Reputation: 126154

If you can access your data as a tf.data.Dataset, there is no need to convert it to a NumPy array before passing it to an Estimator. You can simply build the Dataset directly in your input function, with something like the following:

def train_input_fn():
  dataset = tf.data.TFRecordDataset(filenames=filenames)
  dataset = dataset.map(_read_one_image)
  dataset = dataset.repeat(1)  # Because `num_epochs=1`.
  dataset = dataset.batch(100)  # Because `batch_size=1`.

  dataset = dataset.prefetch(1)  # To improve performance by overlapping execution.

  iterator = dataset.make_one_shot_iterator()  # NOTE: Use a "one-shot" iterator.
  labels_tensor, images_tensor = iterator.get_next()

  return {"x": images_tensor}, labels_tensor

dataset_classifier.train(
    input_fn=train_input_fn, num_epochs=1, hooks=[logging_hook])

This should be much more efficient than building a NumPy array, because it avoids having to materialize the whole dataset in memory at once. You can also use performance enhancements like Dataset.prefetch() and the parallel version of Dataset.map() to improve the training speed.

Upvotes: 3

Related Questions