tensorflow OOM on model.fit with prior validation set preprocessing on GPU

Viewed 180

I'm playing with this colab locally with an 8gb rtx 3070 on Fedora 35 and tensorflow 2.4.0: https://github.com/tensorflow/similarity/blob/master/examples/kaggle.ipynb I have the same consistent errors on windows with a nvidia gtx 1050 ti (4gb).

I tried to decouple the origin of the OOM error on model.fit and seems linked to the preprocessing phase in which I'm resizing the validation set. In this phase the GPU vram is allocated.

# load validation image in memory
x_test = []
with tf.device('/CPU:0'): #solution to avoid occupation of GPU memory and OOM in model.fit
    for p in tqdm(x_test_p):
        img = tf.io.read_file(p)
        img =  tf.io.decode_image(img,dtype=tf.dtypes.float32)
        img = tf.image.resize_with_pad(img, IMG_SIZE, IMG_SIZE)
        # if grayscale, convert to rgb
        if tf.shape(img)[2]==3:
            pass
        else:
            img = tf.image.grayscale_to_rgb(img)
        x_test.append(img)
  1. If i reduce the validation set size alot the model.fit will succeed.
  2. If i process the whole validation with CPU instead of GPU the model.fit will succeed.
  3. If i don't pass the whole preprocessed validation set to model.fit the OOM error will still be present. That's why i'm suggesting it's a problem related to the preprocessing alone occupying useful GPU Memory.

Is it possible that this preprocessing is loaded in VRAM and not released therefore limiting the model.fit GPU memory space left?

The problem is somehow similar to this question from which i took the idea of preprocessing the validation with CPU: Keras OOM for data validation using GPU

0 Answers
Related