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

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,3]),tf.random.normal([batch_size,48]),92]),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来说太大,我会期望立即出现OOM错误,并且我还假定实际上有6GB的VRAM满足。那是因为仅在33个实例之后才发生OOM错误,这使我怀疑内存使用量正在不断增加。

请参见以下我的模型的 Keras摘要

__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
visual_input (InputLayer)       [(32,None,200 0                                            
__________________________________________________________________________________________________
proprioceptive_input (InputLaye [(32,48)]     0                                            
__________________________________________________________________________________________________
somatosensory_input (InputLayer [(32,92)]     0                                            
__________________________________________________________________________________________________
time_distributed (TimeDistribut (None,64)     272032      visual_input[0][0]               
__________________________________________________________________________________________________
time_distributed_1 (TimeDistrib (None,8)      848         proprioceptive_input[0][0]       
__________________________________________________________________________________________________
time_distributed_2 (TimeDistrib (None,8)      3032        somatosensory_input[0][0]        
__________________________________________________________________________________________________
concatenate (concatenate)       (None,80)     0           time_distributed[0][0]           
                                                                 time_distributed_1[0][0]         
                                                                 time_distributed_2[0][0]         
__________________________________________________________________________________________________
time_distributed_3 (TimeDistrib (None,48)     3888        concatenate[0][0]                
__________________________________________________________________________________________________
time_distributed_4 (TimeDistrib (None,48)     0           time_distributed_3[0][0]         
__________________________________________________________________________________________________
time_distributed_5 (TimeDistrib (None,32)     1568        time_distributed_4[0][0]         
__________________________________________________________________________________________________
time_distributed_6 (TimeDistrib (None,32)     0           time_distributed_5[0][0]         
__________________________________________________________________________________________________
goal_input (InputLayer)         [(32,7)]      0                                            
__________________________________________________________________________________________________
concatenate_1 (concatenate)     (32,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:如评论中所述,问题可能是在输入上调用模型会增加计算图的事实。但是我不能使用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所建议的那样,当样本通过网络传递时,由于它是象征性的,因此样本被添加到图形中。因此,我想这个问题现在归结为“即使是一件事,如何使张量变得非符号化”。

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

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

我花了一段时间,但现在已经解决了这个问题。正如我之前在“问题”中所做的编辑一样:问题是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,3])),tf.convert_to_tensor(tf.random.normal([batch_size,48])),92])),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

本文链接:https://www.f2er.com/3080767.html

大家都在问