Tensorflow out of memory error while running autoencoders

Question

I am running a google colab to explore autoencoders to remove noisy data and I keep getting an OOM error from TensorFlow. My code is as follows:

def denoising_autoencoder():
  model = Sequential()
  model.add(keras.layers.Input(shape=(420, 540, 1)))
  model.add(keras.layers.Conv2D(48, (3, 3), activation='relu', padding='same'))
  model.add(keras.layers.Conv2D(72, (3, 3), activation='relu', padding='same'))
  model.add(keras.layers.Conv2D(144, (3, 3), activation='relu', padding='same'))
  model.add(keras.layers.BatchNormalization())
  model.add(keras.layers.MaxPooling2D((2, 2), padding='same'))
  model.add(keras.layers.Dropout(0.5))
  model.add(keras.layers.Conv2D(144, (3, 3), activation='relu', padding='same'))
  model.add(keras.layers.Conv2D(72, (3, 3), activation='relu', padding='same'))
  model.add(keras.layers.Conv2D(48, (3, 3), activation='relu', padding='same'))
  model.add(keras.layers.BatchNormalization())
  model.add(keras.layers.UpSampling2D((2, 2)))
  model.add(keras.layers.Conv2D(1, (3, 3), activation='sigmoid', padding='same'))
  return model

model = denoising_autoencoder()
model.compile(optimizer='adam', loss='mean_squared_error', metrics=['mean_absolute_error'])
callback = EarlyStopping(monitor='loss', patience=30)
history = model.fit(X_train, Y_train, validation_data = (X_val, Y_val), epochs=30, batch_size=24, verbose=0, callbacks=[callback])

And the error I get is as follows:

ResourceExhaustedError:  OOM when allocating tensor with shape[24,48,420,540] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
     [[{{node gradient_tape/sequential/conv2d_6/Conv2D/Conv2DBackpropFilter-0-TransposeNHWCToNCHW-LayoutOptimizer}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
 [Op:__inference_train_function_1289]

I have also tried the following (based on reading some other posts regarding this issue):

gpu_options = tf.compat.v1.GPUOptions(allow_growth=True)
session = tf.compat.v1.InteractiveSession(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options))

Answer 1

It could because the numbers are floats, that can consume lots of memory. You can try quantization, which is a technique to mitigate this problem. Here's the documentation .