TensorFlow Q-learning 训练速度骤降的根源与解决方案

聖光之護 2026-01-08 00:00:00 次阅读

在 tensorflow 中实现 q-learning 时，若在训练循环中反复构建或保存模型却未清理计算图状态，会导致内存泄漏和计算图持续膨胀，从而引发后续轮次训练显著变慢；调用 `tf.keras.backend.clear_session()` 可有效释放全局状态、恢复性能。

该问题并非源于算法逻辑或超参数设置，而是典型的 TensorFlow 运行时状态管理疏漏。TensorFlow 2.x 默认启用 Eager Execution，但其底层仍维护一个全局的 Keras 后端会话（session）和计算图缓存机制。每当调用 model.save()（尤其是 HDF5 格式 .h5），Keras 会将模型结构、权重及关联的计算图元信息注册到全局状态中；若未显式清理，这些历史模型对象将持续驻留内存，并导致新训练步骤的图构建、梯度追踪和自动微分开销逐轮递增——表现为每轮 episode 的 train() 耗时明显上升。

根本解决方法：在每次模型保存后立即调用 tf.keras.backend.clear_session()

for episode in range(MAX_EPISODES):
    obs = env.reset()
    while True:
        left_action = env.left_ball.q_agent.act(np.reshape(obs, [1, *env.state_size]))
        next_obs, rewards, done, _ = env.step(left_action, right_action)

        left_state = np.reshape(obs, [1, *env.state_size])
        left_next_state = np.reshape(next_obs, [1, *env.state_size])
        env.left_ball.q_agent.train(left_state, left_action, rewards[0], left_next_state, done)
        obs = next_obs

        if done:
            # ✅ 关键修复：保存模型后立即清除后端会话
            env.left_ball.q_agent.save_model("left_trained_agent.h5")
            tf.keras.backend.clear_session()  # ← 此行必不可少
            break

注意事项与最佳实践：

clear_session() 会销毁当前所有模型、层、优化器等全局对象，因此不可在单次训练过程中频繁调用（例如每个 batch 后调用），仅适用于“阶段性保存 + 重置环境”的场景（如每 episode 结束）；
若需在训练中动态创建多个模型（如双网络 DQN 中的 target network 更新），建议统一管理模型生命周期，避免重复 save() + 忘记 clear_session()；
替代方案：改用 tf.keras.models.save_model(..., save_format='tf') 保存 SavedModel 格式，其对会话依赖更低，但仍建议配合 clear_session() 使用以确保稳定性；
验证效果：可在循环内添加计时日志（如 time.time()），观察 clear_session() 加入前后各 episode 的 train() 平均耗时是否回归稳定。