title: How to efficiently load large training data too big for RAM for training in tensorflow?
body: <p>I've got around 10 GB of training data in numpy array format. However, my RAM is not big enough to load the data and the tensorflow 2.0 model at the same time. I've done plenty of research into tf.data, tf.TFRecords and generators, but I'm stuck in the actual implementation of the training loop. I currently have the following standard training loop, when using a loadable subset of the data to test whether everything worked. </p>
  
  <p>The options I considered are:</p>
  
  <p>1) Splitting the 10GB of numpy files into a large number of tf.TFRecords files (shuffled) with each being the size of the BATCH_SIZE and then loading them in every train_step and epoch</p>
  
  <p>2) Splitting the 10GB of numpy files into smaller files of 1-2 GB each (shuffled) and then sequentially training the model by loading the smaller files every epoch sequentially</p>
  
  <p>Is there a best practice for this, or a better option I am not considering? I imagine this should be a trivial problem, but I can't find any good solutions online.</p>
  
  <pre><code>if __name__ == '__main__':
      # get the original_dataset
      #train_dataset, valid_dataset
      train_dataset = tf.data.Dataset.from_tensor_slices((mtr, labels))
      train_dataset = train_dataset.shuffle(buffer_size=mtr.shape[0]).batch(BATCH_SIZE)
  
      valid_dataset = tf.data.Dataset.from_tensor_slices((val_mtr, val_labels))
      valid_dataset = valid_dataset.batch(BATCH_SIZE)
  
      # create model
      model = get_model()
  
      # define loss and optimizer 
      loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
      optimizer = tf.keras.optimizers.Adam(learning_rate=1e-04, beta_1=0.9, beta_2=0.999, epsilon=1e-08,)
  
      train_loss = tf.keras.metrics.Mean(name='train_loss')
      train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
  
      valid_loss = tf.keras.metrics.Mean(name='valid_loss')
      valid_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='valid_accuracy')
  
      @tf.function
      def train_step(images, labels):
          with tf.GradientTape() as tape:
              predictions = model(images, training=True)
              loss = loss_object(y_true=labels, y_pred=predictions)
          gradients = tape.gradient(loss, model.trainable_variables)
          optimizer.apply_gradients(grads_and_vars=zip(gradients, model.trainable_variables))
          train_loss(loss)
          train_accuracy(labels, predictions)
  
      @tf.function
      def valid_step(images, labels):
          predictions = model(images, training=False)
          v_loss = loss_object(labels, predictions)
          valid_loss(v_loss)
          valid_accuracy(labels, predictions)
  
      # start training
      for epoch in range(EPOCHS):
          train_loss.reset_states()
          train_accuracy.reset_states()
          valid_loss.reset_states()
          valid_accuracy.reset_states()
          step = 0
          for imgs, lbls in train_dataset:
              step += 1
              train_step(imgs, lbls)
              print("Epoch: {}/{}, step: {}/{}, loss: {:.5f}, accuracy: {:.5f}".format(epoch + 1, EPOCHS, step, math.ceil(mtr.shape[0] / BATCH_SIZE),
                                                                                       train_loss.result(), train_accuracy.result()))
          for valid_images, valid_labels in valid_dataset:
              valid_step(valid_images, valid_labels)
  
          print("Epoch: {}/{}, train loss: {:.5f}, train accuracy: {:.5f}, "
                "valid loss: {:.5f}, valid accuracy: {:.5f}".format(epoch + 1, EPOCHS, train_loss.result(),
                                                                    train_accuracy.result(), valid_loss.result(), valid_accuracy.result()))
  
      model.save_weights(filepath=save_model_dir, save_format='tf')
  
  </code></pre>
  
  <p><strong>Based on the answer below this seemed to work:</strong></p>
  
  <pre><code>def load_data(filename: tf.Tensor):
       coll = pickle.load(open("./TrainData2016/%s" %(str(filename.numpy().decode('utf-8'))), "rb")) #coll[0] = data, coll[1] = label
       return tf.convert_to_tensor(coll[0],dtype=tf.float32), tf.convert_to_tensor(coll[1],dtype=tf.int32)
  
  filenames = os.listdir("./TrainData2016") # You can just shuffle this array and not use tf.data.Dataset.shuffle()
  train_dataset = tf.data.Dataset.from_tensor_slices(filenames)
  train_dataset = train_dataset.map(lambda x: tf.py_function(func=load_data, inp=[x], Tout=(tf.float32, tf.int32)))
  </code></pre>