Consider the following code that works with a Keras Sequential model on the CIFAR-10 data set. Background is given at the end of the post:
import tensorflow as tf
from sklearn.datasets import fetch_openml
from sklearn.utils import shuffle
data, targets = shuffle(*fetch_openml('CIFAR_10', version=1, return_X_y=True))
train_sz = 50000
X_train, X_test, y_train, y_test = data[:train_sz, :], data[train_sz:, :], np.asarray(targets[:train_sz], dtype=np.int), np.asarray(targets[train_sz:], dtype=np.int)
model = tf.keras.Sequential()
model.add(tf.keras.Input(shape=(X_train.shape[1],)))
model.add(tf.keras.layers.Dense(64, activation='relu'))
model.add(tf.keras.layers.Dense(10))
model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer='adam')
s = 0
for _ in range(500):
for i in range(100):
layers = []
for layer in model.get_weights():
layers.append(np.random.normal(0, 1, layer.shape))
model.set_weights(layers)
eval = model.evaluate(X_train, y_train)
s += eval
print(f'Done {i}')
print(s)
After about 1 (sometimes a little before that, sometimes a little after) iteration of the outer for loop, Python crashes with exit code 137, which usually means out of memory AFAIK. I have 16 GB of memory on my system, of which about 20% is used before running this. After running it, it steadily increases up to
about 80%-90% memory usage, then drops to 60%-70% (GC kicking in?), then increases again and so on for 2-3 times until it finally crashes.
I'm on a headless Ubuntu 18.04 Server machine, with Python 3.7 in Anaconda, on Tensorflow 2.2 with a Titan X GTX GPU that is not being used for anything else (so about 11GB of memory free there).
My calculations (very pessimistic, to be sure):
- I have about 12 GB free when I run this.
- Storing the data uses
60000*32*32*3floating point numbers, which is about1500 MBfor float64s. Let's put down 6 GB here because of all the copies I'm making. Regardless, it looks like this is what uses the most memory. - The layer sizes are negligible at this point:
X_train.shape[1]is 3072 (32*32*3), and 64 hidden units is nothing. model.evaluatehas a default batch size of 32, so inside it, it should use about32*32*32*3*64float64s for the output of the middle layer. That's 50 MB, let's put in 1 GB here just to be sure again.model.evaluateprobably also needs to store the predictions, so that's50000*10float64s, which is another 4 MB. Let's put in another 1 GB here for good measure.
Total: 6 + 1 + 1 = 8 GB. My memory usage should absolutely not exceed 80%, and I have overestimated the calculations by a lot.
Why is so much memory being used and can I optimize anything in how I manage the data?
I've tried forcing the X's to np.int using np.asarray, there's no point for float64s there, but that just makes it crash much faster - it's like it keeps both the float64s and the ints in memory or something.
Background
I'm working on a genetic algorithm that trains artificial neural networks. I've traced the crash to the computation of the fitnesses, which involves applying the trained weights stored in each individual to the neural network and evaluating the network (inner i loop, I have 100 individuals in my population). This is repeated for each generation (outermost for loop). A bit more memory is used there, but still very little.
That is why there is no fitting going on here, the weights are determined by my genetic algorithm and applied to the network.
This reduced code reproduces the issue.

