Search code examples
pythontensorflowkerasrecurrent-neural-networkbatchsize

Keras seem to ignore my batch_size and tries to fit all data in GPU memory


I have a RNN with 2 input arrays, and 2 output arrays, the sizes are:

input1 = (339679, 90, 15)
input2 =(339679, 90, 27)
output1 = 339679,2
output2 = 339679,16

The code to create the RNN LSTM is (I will only show one of the two RNN, the other one is identical but with 16 outputs and getting input sizes from input2):

inputs = Input(shape=(n_in1.shape[1], n_in1.shape[2]), name='inputs')                   
lstmA1 = LSTM(1024, return_sequences=True, name="lstmA1") (inputs)                    
lstmA2 = LSTM(512//1, return_sequences=True, name="lstmA2")(lstmA1)                   
lstmA3 = LSTM(512//2, return_sequences=True, name="lstmA3")(lstmA2)                   
lstmA4 = LSTM(512//2, return_sequences=True, name="lstmA4")(lstmA3)                   
lstmA5 = LSTM(512//4, return_sequences=False, name="lstmA5")(lstmA4)                  
denseA1 = DenseBND(lstmA5, 512//1, "denseA1", None, False, 0.2)                       
denseA2 = DenseBND(denseA1, 512//4, "denseA2", None, False, 0.2)                      
denseA3 = DenseBND(denseA2, 512//8, "denseA3", None, False, 0.2)                      
outputsA = Dense(2, name="outputsA")(denseA3)  

Here, n_in1 is input1 I described before, so the shape given is 90,15

DenseBND is just a function that returns a Dense layer with BatchNormalization and dropout. In this case BatchNormalization is False, the activation function is None and Dropout is 0.2 So it just returns Dense layer with linear activation function and 20% Dropout.

And finally, the line to train it:

model.fit( {'inputsA': np.array(n_input1), 'inputsB': np.array(n_input2)},
      {'outputsA': np.array(n_output1), 'outputsB': np.array(n_output2)},
      validation_split=0.1, epochs=1000, batch_size=256, 
      callbacks=callbacks_list)

You can see that the validation_split is 0.1 and the batch_size is 256

Nevertheless, when I try to train it I get the following error:

ResourceExhaustedError: OOM when allocating tensor with shape[335376,90,27] and type float on /job:

As you can see, it seems it tries to fit the whole dataset into GPU memory, instead of going batch by batch. I ever set batch_size to 1, and this error persists. The first number 335376 is a 90% of my dataset (this number is different to the one above, the one above is the one that works, this one doesn't).

Shouldn't it try to allocate tensor with shape 256,90,27?


Solution

  • No, keras is not ignoring your batch size.

    You're trying to create numpy arrays with too large dimensions.

    'inputsB': np.array(n_input2) this allocates a really large numpy array, so even before the training starts, this numpy conversion is not possible due to limited memory.

    You need to use data generators which doesn't laod the full data at once into memory.

    Ref: https://keras.io/api/preprocessing/