[Python] OOM error on GPU while training a RNN in model selection phase, tensor dimensions...

I have a NLP tensor such as train22k, 170, 300) val: (2k, 170, 300), test: (25k, 170,300) where the last dim 300 are FastText embs, also I have one GPU Tesla 32GB. I'm doing model selection on a RNN untrained an the batch/buffer_size is 64 the layers are 5:

for config in param_grid:
model = self.create_model(config)
train = model.forward(embedded_training_data)
val = model.forward(embedded_val_data)
test= model.forward(embedded_test_data)


The model created is a Keras sequential where each layers are keras.layers.Bidirectional and sequence=True hence output is 3D (batch, timesteps, features), the forward method is the following and uses a batches computation called compute_states:

def compute_states(self, x):
x_train_states = []
for i, layer in enumerate(self.layers):
outputs, r, b = layer(x)
x_train_states.append(outputs) # Aggiungi alla lista

x = outputs # Update input for next layer

return tf.concat(x_train_states, axis=2) if x_train_states else None


@tf.function
def forward(self, data):
total_samples = tf.shape(data)[0]
buffer_size = tf.constant(self.buffer_size, dtype=tf.int32)
num_batches = tf.cast(tf.math.ceil(total_samples / buffer_size), tf.int32)

states_array = tf.TensorArray(dtype=tf.float32, size=num_batches)

for i in tf.range(num_batches):
start_idx = i * buffer_size
end_idx = tf.minimum((i + 1) * buffer_size, total_samples)
batch = data[start_idx:end_idx]
states = self.compute_states(batch)
states_array = states_array.write(i, states)

states = states_array.concat()
return states


These functions are good and very speed on CPU but on GPU I receive OOM error on the concatenation of batches (states_array.concat()). I'd like to know if there is any issues in my code, and therefore I could optimize it or if the tensors dimensions are intractable.

Continue reading...