So I'm initializing a model as:
model = tf.keras.utils.multi_gpu_model(model, gpus=NUM_GPUS) and when I do model.compile() it runs perfectly fine.
But when I do history = model.fit(tf.cast(X_train, tf.float32), tf.cast(Y_train, tf.float32), validation_split=0.25, batch_size = 16, verbose=1, epochs=100), it gives me error:
OOM when allocating tensor with shape[4760,256,256,3] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:Cast] name: Cast/
This code worked perfectly fine previously but not anymore with Tensorflow 2.0. I have 4760 samples in my training set. I don't know why it's taking the entire set instead of the batch size.
model.compile() only does configure the model for training and it doesn't have any memory allocation.
Your bug is self-explained, you directly feed a large numpy array into the model. I would suggest coding a new data generator or keras.utils.Sequence to feed your input data. If so, you do not need specify the batch_size in fit method again, because your own generator or Sequence will generate batches.