title: How to feed large PrefetchDataset to 3d CNN without exhausting memory?
body: <p>I'm reading 4D image numpy array saved into hard disk by memory mapping and feeding into tf.data.Dataset.from_tensor_slices along with the target variable. The numpy array size is more than 10 GB. When I go into 3d CNN model training, the kernel dies because of memory exhausting. How to solve this problem?</p>
  <p>Below is the partial code:</p>
  <pre><code>def load_train():
      
      X = np.load('trainX.npy', mmap_mode='c')
      y = np.load('trainY.npy', mmap_mode='c')
      
      return X, y
  
  x_train, y_train = load_train()
  
  def load_val():
      
      X = np.load('testX.npy', mmap_mode='c')
      y = np.load('testY.npy', mmap_mode='c')
      
      return X, y
  
  x_val, y_val = load_val()
  
  
  train_loader = tf.data.Dataset.from_tensor_slices((x_train, y_train))
  validation_loader = tf.data.Dataset.from_tensor_slices((x_val, y_val))
  
  batch_size = 2
  
  train_dataset = (
      train_loader.shuffle(len(x_train))
      .map(train_preprocessing)
      .batch(batch_size)
      .prefetch(2)
  )
  
  validation_dataset = (
      validation_loader.shuffle(len(x_val))
      .map(validation_preprocessing)
      .batch(batch_size)
      .prefetch(2)
  )
  
  
  def get_model(width=512, height=512, depth=645):
  
      inputs = keras.Input((width, height, depth, 1))
  
      x = layers.Conv3D(filters=64, kernel_size=3, activation=&quot;relu&quot;)(inputs)
      x = layers.MaxPool3D(pool_size=2)(x)
      x = layers.BatchNormalization()(x)
  
      x = layers.Conv3D(filters=64, kernel_size=3, activation=&quot;relu&quot;)(x)
      x = layers.MaxPool3D(pool_size=2)(x)
      x = layers.BatchNormalization()(x)
  
      x = layers.Conv3D(filters=128, kernel_size=3, activation=&quot;relu&quot;)(x)
      x = layers.MaxPool3D(pool_size=2)(x)
      x = layers.BatchNormalization()(x)
  
      x = layers.Conv3D(filters=256, kernel_size=3, activation=&quot;relu&quot;)(x)
      x = layers.MaxPool3D(pool_size=2)(x)
      x = layers.BatchNormalization()(x)
  
      x = layers.GlobalAveragePooling3D()(x)
      x = layers.Dense(units=512, activation=&quot;relu&quot;)(x)
      x = layers.Dropout(0.3)(x)
  
      outputs = layers.Dense(units=1, activation=&quot;sigmoid&quot;)(x)
  
      model = keras.Model(inputs, outputs, name=&quot;3dcnn&quot;)
      return model
  
  model = get_model(width=512, height=512, depth=645)
  
  model.compile(
      loss=&quot;binary_crossentropy&quot;,
      optimizer=keras.optimizers.Adam(learning_rate=lr_schedule),
      metrics=[&quot;acc&quot;],
  )
  
  model.fit(
      train_dataset,
      validation_data=validation_dataset,
      epochs=epochs,
      shuffle=True,
      verbose=2,
      callbacks=[checkpoint_cb, early_stopping_cb],
  )
  </code></pre>