将 Keras 模型应用于符号张量的 TF2.0 内存泄漏答案

【问题标题】：TF2.0 Memory Leak From Applying Keras Model to Symbolic Tensor将 Keras 模型应用于符号张量的 TF2.0 内存泄漏
【发布时间】：2020-03-13 21:06:10
【问题描述】：

tldr：我的实现的内存使用显然随着通过它的样本数量而增长，但在网络/样本馈送中应该没有什么关心如何到目前为止已经通过了很多样本。

当通过函数式 API 创建的自定义 Keras 模型传递大量高维数据时，我观察到我认为 GPU 内存使用量持续增长，观察到的实例数量不断增加.以下是通过网络传递样本的过程的最小示例：

sequence_length = 100
batch_size = 128

env = gym.make("ShadowHand-v1")
_, _, joint = build_shadow_brain(env, bs=batch_size)
optimizer: tf.keras.optimizers.Optimizer = tf.keras.optimizers.SGD()

start_time = time.time()
for t in tqdm(range(sequence_length), disable=False):
    sample_batch = (
        tf.random.normal([batch_size, 1, 200, 200, 3]),
        tf.random.normal([batch_size, 1, 48]),
        tf.random.normal([batch_size, 1, 92]),
        tf.random.normal([batch_size, 1, 7])
    )

    with tf.GradientTape() as tape:
        out, v = joint(sample_batch)
        loss = tf.reduce_mean(out - v)

    grads = tape.gradient(loss, joint.trainable_variables)
    optimizer.apply_gradients(zip(grads, joint.trainable_variables))
    joint.reset_states()

print(f"Execution Time: {time.time() - start_time}")

我知道考虑到批量大小，这是一个很大的样本，但是如果它实际上对于我的 GPU 来说太大并且我还假设实际上有 6GB 的 VRAM，我会期望它会立即出现 OOM 错误够了。那是因为只有在 33 次之后才会发生 OOM 错误，这让我怀疑内存使用量在增加。

在下面查看我的模型的 Keras 摘要：

__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
visual_input (InputLayer)       [(32, None, 200, 200 0                                            
__________________________________________________________________________________________________
proprioceptive_input (InputLaye [(32, None, 48)]     0                                            
__________________________________________________________________________________________________
somatosensory_input (InputLayer [(32, None, 92)]     0                                            
__________________________________________________________________________________________________
time_distributed (TimeDistribut (None, None, 64)     272032      visual_input[0][0]               
__________________________________________________________________________________________________
time_distributed_1 (TimeDistrib (None, None, 8)      848         proprioceptive_input[0][0]       
__________________________________________________________________________________________________
time_distributed_2 (TimeDistrib (None, None, 8)      3032        somatosensory_input[0][0]        
__________________________________________________________________________________________________
concatenate (Concatenate)       (None, None, 80)     0           time_distributed[0][0]           
                                                                 time_distributed_1[0][0]         
                                                                 time_distributed_2[0][0]         
__________________________________________________________________________________________________
time_distributed_3 (TimeDistrib (None, None, 48)     3888        concatenate[0][0]                
__________________________________________________________________________________________________
time_distributed_4 (TimeDistrib (None, None, 48)     0           time_distributed_3[0][0]         
__________________________________________________________________________________________________
time_distributed_5 (TimeDistrib (None, None, 32)     1568        time_distributed_4[0][0]         
__________________________________________________________________________________________________
time_distributed_6 (TimeDistrib (None, None, 32)     0           time_distributed_5[0][0]         
__________________________________________________________________________________________________
goal_input (InputLayer)         [(32, None, 7)]      0                                            
__________________________________________________________________________________________________
concatenate_1 (Concatenate)     (32, None, 39)       0           time_distributed_6[0][0]         
                                                                 goal_input[0][0]                 
__________________________________________________________________________________________________
lstm (LSTM)                     (32, 32)             9216        concatenate_1[0][0]              
__________________________________________________________________________________________________
dense_10 (Dense)                (32, 20)             660         lstm[0][0]                       
__________________________________________________________________________________________________
dense_11 (Dense)                (32, 20)             660         lstm[0][0]                       
__________________________________________________________________________________________________
activation (Activation)         (32, 20)             0           dense_10[0][0]                   
__________________________________________________________________________________________________
activation_1 (Activation)       (32, 20)             0           dense_11[0][0]                   
__________________________________________________________________________________________________
concatenate_2 (Concatenate)     (32, 40)             0           activation[0][0]                 
                                                                 activation_1[0][0]               
__________________________________________________________________________________________________
dense_12 (Dense)                (32, 1)              33          lstm[0][0]                       
==================================================================================================
Total params: 291,937
Trainable params: 291,937
Non-trainable params: 0
__________________________________________________________________________________________________

如您所见，该网络中有一个 LSTM 层。它通常应该是有状态的，但是我已经关闭了它，因为我认为问题出在那儿。实际上我已经尝试了以下方法，但没有消除问题

状态转换
完全删除 LSTM
不计算任何梯度
每次实例后重建模型

现在我对问题潜在原因的想法已经结束。

我还将进程强制到 CPU 上 并检查了标准内存（这里没有发生 OOM，因为我的 RAM 比 VRAM 多得多）。有趣的是，内存使用率会上下波动，但呈上升趋势。对于每个实例，大约会占用 2GB 内存，但在执行下一个采样之前释放内存时，只释放了比占用内存少大约 200MB 的内存。

编辑 1： 如 cmets 中所述，问题可能是在输入上调用模型会添加到计算图上。但是我不能使用joint.predict()，因为我需要计算梯度。

编辑 2： 我更密切地监测了内存的增长，实际上每次迭代都会保留一些内存，如您在此处看到的前 9 个步骤：

0: 8744054784
1: 8885506048
2: 9015111680
3: 9143611392
4: 9272619008
5: 9405591552
6: 9516531712
7: 9647988736
8: 9785032704

这是在批量大小为 32 的情况下完成的。一个 sample_batch 的大小为 256 * (200 * 200 * 3 + 48 + 92 + 7) * 32 = 984244224 位（精度为 float32），这或多或少表明确实存在问题，当将样本通过网络，样本被添加到图中，因为它是符号的，正如@MatiasValdenegro 所建议的那样。所以我想现在的问题可以归结为“如何使张量成为非符号化的”，如果这是一件事的话。

免责声明：我知道您无法使用给定代码重现该问题，因为缺少组件，但我无法在此处提供完整的项目代码。

【问题讨论】：

问题是你给模型（tf.random.normal）提供了符号张量，这给图形增加了操作，所以图形的大小不断增加。您应该使用 model.predict 而不是将张量传递给模型。
@MatiasValdenegro 虽然确实使用 predict 可以防止 OOM 错误发生，但我不能在这里使用它，因为我最终想训练网络，为此我（认为）我需要直接调用（它也是tensorflow 在他们的指南中使用了什么：tensorflow.org/guide/keras/…)。我将编辑问题以澄清我的意图。

标签： python tensorflow keras memory-leaks out-of-memory

【解决方案1】：

我花了一段时间，但我现在已经解决了这个问题。正如我之前编辑过的问题：问题是 Keras 的功能 API 似乎将每个样本添加到计算图中，而没有删除迭代后我们不再需要的输入。似乎没有简单的方法可以显式删除它，但是tf.function 装饰器可以解决问题。

以我上面的代码示例为例，可以如下应用：

sequence_length = 100
batch_size = 256

env = gym.make("ShadowHand-v1")
_, _, joint = build_shadow_brain(env, bs=batch_size)
plot_model(joint, to_file="model.png")
optimizer: tf.keras.optimizers.Optimizer = tf.keras.optimizers.SGD()

@tf.function
def _train():
    start_time = time.time()

    for _ in tqdm(range(sequence_length), disable=False):
        sample_batch = (tf.convert_to_tensor(tf.random.normal([batch_size, 4, 224, 224, 3])),
                        tf.convert_to_tensor(tf.random.normal([batch_size, 4, 48])),
                        tf.convert_to_tensor(tf.random.normal([batch_size, 4, 92])),
                        tf.convert_to_tensor(tf.random.normal([batch_size, 4, 7])))

        with tf.GradientTape() as tape:
            out, v = joint(sample_batch, training=True)
            loss = tf.reduce_mean(out - v)

        grads = tape.gradient(loss, joint.trainable_variables)
        optimizer.apply_gradients(zip(grads, joint.trainable_variables))

    print(f"Execution Time: {time.time() - start_time}")

_train()

也就是说，训练循环可以在带有 tf.function 装饰器的函数中提供。这意味着训练将以图形模式执行，并且由于某种原因，这消除了问题，很可能是因为函数结束后图形将被转储。有关tf.function 的更多信息，请参阅该主题的TF2.0 Guide。

【讨论】：