GPU goes out of memory during training large dataset - tensorflow

Question

I am using a Transformer network for machine translation, during training of model the GPU runs out of memory during large dataset, it works fine with small data.
This is the self attention part, The error comes during the computation of matrices.
import tensorflow as tf
class SelfAttention(tf.keras.layers.Layer):
def __init__(self, embed_size, head):
super(SelfAttention, self).__init__()
self.head = head
self.embed_size = embed_size
self.head_dim = embed_size // head
assert (self.head_dim * head == embed_size), 'size of head_dim is not matching'
self.query = tf.keras.layers.Dense(self.head_dim, activation='linear', use_bias=False)
self.value = tf.keras.layers.Dense(self.head_dim, activation='linear', use_bias=False)
self.key = tf.keras.layers.Dense(self.head_dim, activation='linear', use_bias=False)
self.fc_layer = tf.keras.layers.Dense(self.embed_size, activation='linear')
def call(self, value, key, query, mask):
# Number of training examples
N = query.shape[0]
query_len, value_len, key_len = query.shape[1], value.shape[1], key.shape[1]
# Reshape according to the number of examples and words
query = tf.reshape(query, (N, query_len, self.head, self.head_dim))
value = tf.reshape(value, (N, value_len, self.head, self.head_dim))
key = tf.reshape(key, (N, key_len, self.head, self.head_dim))
query = self.query(query)
value = self.value(value)
key = self.key(key)
# energy shape: (N, head, query_len, key_len) try to imagine the shape in mind
energy = tf.einsum("nqhd, nkhd->nhqk", query, key)
if mask is not None:
energy = energy * mask
energy = tf.where(tf.equal(energy, 0), -1e20, energy)
attention = tf.keras.activations.softmax(energy, axis=3)
# attention shape: (N, head, query_len, key_len)
# value shape:(N, value_len, head, head_dim)
# output: (N, query_len, head, head_dim)
output = tf.reshape(tf.einsum("nhql, nlhd->nqhd", attention, value), (N, query_len, self.head*self.head_dim))
output = tf.keras.activations.linear(output)
return output
The error is
2021-09-20 11:51:49.615495: I tensorflow/core/common_runtime/bfc_allocator.cc:1036] 1 Chunks of size 35477760 totalling 33.83MiB
2021-09-20 11:51:49.615502: I tensorflow/core/common_runtime/bfc_allocator.cc:1036] 1 Chunks of size 40866304 totalling 38.97MiB
2021-09-20 11:51:49.615509: I tensorflow/core/common_runtime/bfc_allocator.cc:1036] 1 Chunks of size 47409664 totalling 45.21MiB
2021-09-20 11:51:49.615516: I tensorflow/core/common_runtime/bfc_allocator.cc:1036] 1 Chunks of size 47547136 totalling 45.34MiB
/opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
6860 message = e.message + (" name: " + name if name is not None else "")
6861 # pylint: disable=protected-access
-> 6862 six.raise_from(core._status_to_exception(e.code, message), None)
6863 # pylint: enable=protected-access
6864
/opt/conda/lib/python3.7/site-packages/six.py in raise_from(value, from_value)
ResourceExhaustedError: OOM when allocating tensor with shape[32,334,25335] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:BiasAdd]
What should I do?

Answer 1

You can use a generator to load just a part of the dataset in the GPU memory and with that you will be able to train with your model.
Here is an example of a simple generator for image classification that you need to adjust to your use for NLP:
class DataGenerator(keras.utils.Sequence):
'Generates data for Keras'
def __init__(self, list_IDs, labels, batch_size=32, dim=(32,32,32), n_channels=1,
n_classes=10, shuffle=True):
'Initialization'
self.dim = dim
self.batch_size = batch_size
self.labels = labels
self.list_IDs = list_IDs
self.n_channels = n_channels
self.n_classes = n_classes
self.shuffle = shuffle
self.on_epoch_end()
def __len__(self):
'Denotes the number of batches per epoch'
return int(np.floor(len(self.list_IDs) / self.batch_size))
def __getitem__(self, index):
'Generate one batch of data'
# Generate indexes of the batch
indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]
# Find list of IDs
list_IDs_temp = [self.list_IDs[k] for k in indexes]
# Generate data
X, y = self.__data_generation(list_IDs_temp)
return X, y
def on_epoch_end(self):
'Updates indexes after each epoch'
self.indexes = np.arange(len(self.list_IDs))
if self.shuffle == True:
np.random.shuffle(self.indexes)
def __data_generation(self, list_IDs_temp):
'Generates data containing batch_size samples' # X : (n_samples, *dim, n_channels)
# Initialization
X = np.empty((self.batch_size, *self.dim, self.n_channels))
y = np.empty((self.batch_size), dtype=int)
# Generate data
for i, ID in enumerate(list_IDs_temp):
# Store sample
X[i,] = np.load('data/' + ID + '.npy')
# Store class
y[i] = self.labels[ID]
return X, keras.utils.to_categorical(y, num_classes=self.n_classes)
And then pass it to .fit
params = {'dim': (32,32,32),
'batch_size': 64,
'n_classes': 6,
'n_channels': 1,
'shuffle': True}
# Datasets
partition = # IDs
labels = # Labels
# Generators
training_generator = DataGenerator(partition['train'], labels, **params)
validation_generator = DataGenerator(partition['validation'], labels, **params)
model.fit_generator(generator=training_generator,
validation_data=validation_generator)