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

斯特凡

我的问题似乎很普遍。

我正在使用香草政策梯度法进行一些强化学习。环境只是一个简单的单周期游戏,其中状态和动作空间是真实的界限。该代理是具有两个输出头的神经网络,我使用Keras的密集层手动构建了该输出头,例如,我的第一个隐藏层将是

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

其中NH包含用于隐藏层的神经元数量的列表。输出是我的高斯政策的均值和标准差。这部分无关紧要,但我还是将其包括在内。

环境很简单:状态是一个普通变量,操作是一些实数标量,只有一个周期。我多次运行该策略,收集生成的批处理,并使用tf.GradientTape()中的工具在自定义损失函数的基础上更新网络。我可以毫无疑问地运行该代码数千次以了解算法的学习。

真正的问题是,我想多次运行学习过程,每次都随机重新初始化网络权重以获取奖励历史记录的分布,但是如果我将所有这些循环运行,计算机将迅速死机。显然,这是Keras和Tensorflow的一个非常普遍的问题,人们多年来一直在抱怨这个问题,但仍然是一个问题...现在,我尝试了通常的解决方案。在这里,人们建议在循环的末尾添加类似以下内容的内容,以便在我重新初始化网络之前获得清晰的结果。

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 )
胡安·卡尼亚·莫拉雷斯

您可以尝试通过调用重启后端

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()

本文收集自互联网,转载请注明来源。

如有侵权,请联系 [email protected] 删除。

编辑于
0

我来说两句

0 条评论
登录 后参与评论

相关文章