在循环中使用Keras时，如何防止内存使用量激增

Question

我的问题似乎很常见。

我正在使用一种普通的策略梯度方法进行强化学习。环境只是一个简单的单周期博弈，其中状态和动作空间是真正的线。代理是一个具有两个输出头的神经网络，我使用Keras的密集层手动构建，例如我的第一个隐藏层是

layers.Dense(NH[0], activation ="relu", \
             kernel_initializer=initializers.GlorotNormal())(inputs)

其中NH包含隐藏层的神经元数量的列表。输出是我的高斯策略的平均值和标准差。我不知道这部分是否重要，但我还是把它包括进来了。

环境很简单:状态是一个正常变量，动作是一些真实的标量，并且只有一个句号。我多次运行该策略，收集结果批处理，并使用tf.GradientTape()中的工具根据自定义损失函数更新网络。我可以毫不费力地运行该代码数千次，以观察算法的学习情况。

真正的问题是，我想要多次运行学习过程，每次都随机重新初始化网络权重，以获得奖励历史的分布，但如果我在一个循环中运行所有这些，计算机很快就会死机。显然，这是Keras和Tensorflow的一个非常常见的问题，人们多年来一直在抱怨这个问题，现在仍然是一个问题……现在，我已经尝试了常用的解决方案。Here，人们建议在循环的末尾添加类似以下内容的内容，以便在我重新初始化网络之前可以从头开始。

keras.backend.clear_session()
gc.collect()
del actor

这并不能解决问题。然后，我看到有人给出了一个更进一步的函数

def reset_keras(model):

# Clear model, if possible
try:
    del model
except:
    pass

# Garbage collection
gc.collect()

# Clear and close tensorflow session
session = K.get_session() # Get session
K.clear_session()         # Clear session
session.close()           # Close session

# Reset all tensorflow graphs
tf.compat.v1.reset_default_graph()

这也不起作用。我也试着移动前三个命令的顺序，但也不起作用……

谁知道怎么解决这个问题？了解以下信息也会很有用：为什么这种情况就会发生。我还想知道如何在这里分析内存使用情况，这样我就不必等待4个小时才能了解到使用新解决方案时计算机再次冻结。

事实上，如果你有一个最小的工作示例，你可以证明代码不会导致爆炸性的内存使用，我将非常倾向于从头开始重新编写整个该死的东西来阻止问题。顺便说一句，为什么开发人员还没有解决这个问题？这是R和Python上唯一一个在我身上发生过这种情况的包……

按要求编辑，我提供了这个问题的最小工作示例。我做了一个简单的游戏:这是一个移动的目标，最优动作是玩状态值的一些倍数，得到的奖励是0。

我写下了一个演员类，并使用一个简单的线性回归作为评论者，这可能会被关闭。如果您查看内存使用情况，它正在攀升...这个游戏不会让我的电脑崩溃，除非我玩得更多，但它显示内存使用量增加了。

import numpy      as np
import psutil

import tensorflow                    as tf
import tensorflow.keras              as keras
import tensorflow.keras.layers       as layers
import tensorflow.keras.initializers as initializers

import tensorflow.python.keras.backend as kb
import matplotlib.pyplot as plt

BATCH    = 10
MC_DRAWS = 2000
M        = 10

# Training options
LR = 0.01
def display_memory():
    print( f'{round(psutil.virtual_memory().used/2**30, 2)} GB' )

class Actor:

    def __init__(self):
        self.nn    = self.make_actor()
        self.batch = BATCH
        self.opt   = keras.optimizers.Adam( learning_rate = LR )

    def make_actor(self):
        inputs = layers.Input( shape=(1) )
        hidden = layers.Dense(5, activation='relu', 
                              kernel_initializer=initializers.GlorotNormal() )(inputs)
        mu     = layers.Dense(1, activation='linear',
                          kernel_initializer=initializers.GlorotNormal() )(hidden)
        sigma  = layers.Dense(1, activation='softplus',
                          kernel_initializer=initializers.GlorotNormal() )(hidden)
    
        nn  = keras.Model(inputs=inputs, outputs=[mu, sigma])
    
        return nn

    def update_weights(self, state, action, reward):

        # Get proper format
        state  = tf.constant(state,  dtype='float32', shape=(self.batch,1))
        action = tf.constant(action, dtype='float32', shape=(self.batch,1))
        reward = tf.constant(reward, dtype='float32', shape=(self.batch,1))
    
        # Update Policy Network Parameters
        with tf.GradientTape() as tape:   
            # Compute Gaussian loss
            loss_value = self.custom_loss(state, action, reward)
            loss_value = tf.math.reduce_mean( loss_value, keepdims=True )
        
            # Compute gradients
            grads = tape.gradient(loss_value, self.nn.trainable_variables)
 
            # Apply gradients to update network weights
            self.opt.apply_gradients(zip(grads, self.nn.trainable_variables))
        
    def custom_loss(self, state, action, reward):
        # Obtain mean and standard deviation
        nn_mu, nn_sigma = self.nn(state)
    
        # Gaussian pdf
        pdf_value = tf.exp(-0.5 *((action - nn_mu) / (nn_sigma))**2) *\
                    1/(nn_sigma*tf.sqrt(2 *np.pi))
                    
        # Log probabilities
        log_prob  = tf.math.log( pdf_value + 1e-5 )
    
        # Compute loss
        loss_actor = -reward * log_prob
    
        return loss_actor

class moving_target_game:

    def __init__(self):
        self.action_range = [-np.inf, np.inf]
        self.state_range  = [1, 2]
        self.reward_range = [-np.inf, 0]

    def draw(self):
        return np.random.ranint(low  = self.state_range[0],
                            high = self.state_range[1])

    def get_reward(self, action, state):
        return -(5*state - action)**2

class Critic:  
    def __init__(self):
    
        self.order      = 3
        self.projection = None

    def predict(self, state, reward):
    
        # Enforce proper format
        x = np.array( state ).reshape(-1,1)
        y = np.array( reward ).reshape(-1,1)
    
        # Make regression matrix
        X = np.ones( shape = x.shape )
        for i in range( self.order ):
            X = np.hstack( (X, x**(i+1)) )
        
        # Prediction
        xt = x.transpose()
        P  = x @ np.linalg.inv( xt @ x  ) @ xt
        Py = P @ y
    
        self.projection = P
    
        return Py

#%% Moving Target Game with Actor and Actor-Critic

do_actor_critic = True

display_memory()

history    = np.zeros( shape=(MC_DRAWS, M) )
env        = moving_target_game()

for m in range(M):

    # New Actor Network
    actor  = Actor()

    if do_actor_critic:
        critic = Critic()

    for i in range(MC_DRAWS):
    
        state_tape  = []
        action_tape = []
        reward_tape = []
    
        for j in range(BATCH):
        
            # Draw state
            state = env.draw()
            s     = tf.constant([state], dtype='float32')
        
            # Take action
            mu, sigma = actor.nn( s )
            a         = tf.random.normal([1], mean=mu, stddev=sigma)
        
            # Reward
            r = env.get_reward( state, a )
        
            # Collect results
            action_tape.append( float(a)     )
            reward_tape.append( float(r)     )
            state_tape.append(  float(state) )
        
            del (s, a, mu, sigma)
    
        # Update network weights
        history[i,m] = np.mean( reward_tape )
    
        if do_actor_critic:
            # Update critic
            value = critic.predict(state_tape, reward_tape)
            # Benchmark reward
            mod = np.array(reward_tape).reshape(-1,1) - value
            # Update actor
            actor.update_weights(state_tape, action_tape, mod)
        else:
            actor.update_weights(state_tape, action_tape, reward_tape)

    del actor
    kb.clear_session()

    if do_actor_critic:
        del critic
    
    print( f'Average Reward on last: {np.mean(reward_tape)} ' )
    display_memory()

plt.plot( history )

Answer 1

您可以通过调用以下命令尝试重新启动后端

reset_tensorflow_keras_backend()

在每个模型估计之后，函数定义如下：

def reset_tensorflow_keras_backend():
    # to be further investigated, but this seems to be enough
    import tensorflow as tf
    import tensorflow.keras as keras
    tf.keras.backend.clear_session()
    tf.reset_default_graph()
    _ = gc.collect()