Python Keras 代码无明显原因内存不足

Question

考虑以下与 CIFAR-10 数据集上的 Keras 顺序模型配合使用的代码。文章末尾给出了背景：

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)

在外部 for 循环大约 1 次（有时稍早，有时稍晚）迭代之后，Python 崩溃并显示 exit code 137，这通常意味着内存不足。我的系统上有 16 GB 内存，其中大约 20% 在运行之前已使用。运行后，它的内存使用率稳步增加到大约80%-90%，然后下降到60%-70%（GC启动？），然后再次增加，依此类推2-3次，直到最终崩溃。

我在一台无头 Ubuntu 18.04 服务器机器上，在 Anaconda 中使用 Python 3.7，在 Tensorflow 2.2 上使用 Titan X GTX GPU，该 GPU 没有用于其他任何用途（因此大约有 11GB 的可用内存）。

我的计算（当然非常悲观）：

1. 当我运行这个程序时，我有大约 12 GB 的可用空间。
2. 存储数据使用60000*32*32*3浮点数，这大约适用1500 MB于 float64。由于我正在制作所有副本，因此我们在此预留6 GB空间。无论如何，看起来这是使用最多内存的。
3. 此时层大小可以忽略不计：X_train.shape[1]为 3072 ( 32*32*3)，64 个隐藏单元不算什么。
4. model.evaluate默认批量大小为 32，因此在其内部，应该使用大约32*32*32*3*64float64s 作为中间层的输出。这是 50 MB，为了再次确定，我们在这里放入1 GB 。
5. model.evaluate可能还需要存储预测，所以这是50000*10float64，又是 4 MB。让我们在这里另外放入1 GB以进行更好的衡量。

总计：6 + 1 + 1 = 8 GB。我的内存使用绝对不应该超过80%，而且我高估了计算量。

为什么使用了这么多内存？我可以优化管理数据的方式吗？

我尝试过强制 X 使用np.int，np.asarray那里的 float64 没有任何意义，但这只会让它崩溃得更快 - 就像它将 float64 和 int 都保留在内存中一样。

背景

我正在研究一种训练人工神经网络的遗传算法。我将崩溃追溯到适应度的计算，其中涉及将每个个体中存储的训练权重应用到神经网络并评估网络（内部i循环，我的群体中有 100 个个体）。每一代都会重复此操作（最外层的 for 循环）。那里使用了更多的内存，但仍然很少。

这就是为什么这里没有进行拟合，权重由我的遗传算法确定并应用于网络。

此减少的代码重现了该问题。

Answer 1

我能够在具有 12GB 内存的 google colab GPU 实例上的笔记本中复制您的问题。经过 6 次迭代后，我的内存飙升至约 2.5GB，然后在第 50 次迭代时约 6GB，然后内核死亡。

通过在稳定在 ~1GB 的每个内部循环内存中调用垃圾收集器，我能够继续传递 2 个外部迭代。（然后我取消了）

我怀疑为什么会发生这种情况，因为tensorflow在迭代中创建引用的速度比垃圾收集器默认情况下收集它们的速度要快。

import gc
import tensorflow as tf
from sklearn.datasets import fetch_openml
from sklearn.utils import shuffle
import numpy as np

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):
        gc.collect()
        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} eval {eval}')
s

PS：手动正式调用垃圾收集器是您想要避免的事情，但有时它可以完成工作。