启用Eager Execution时CudnnLSTM显存不足问题求助

我之前也碰到过类似的坑，Eager模式下跑CudnnLSTM的内存管理逻辑和Graph模式确实有差异，结合你的场景（AWS p2.xlarge实例、3层双向CudnnLSTM、batch size=32），给你几个实用的解决方案：

1. 开启GPU显存按需分配 #

Eager模式下TensorFlow默认会一次性占用大量GPU显存，导致后续分配RNN workspace时空间不足。你可以通过配置让显存按需增长，避免一开始就占满：

如果是你使用的TensorFlow 1.x版本（2018年的主流版本），要在开启Eager前设置显存策略：

import tensorflow as tf

# 配置GPU显存按需分配
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
# 启用Eager Execution
tf.enable_eager_execution(config=config)

如果是升级到TF2.x后的参考配置：

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 对每个GPU设置显存按需增长
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

2. 临时调整batch size测试 #

Eager模式下单步的内存开销确实比Graph模式高一些，你可以先把batch size降到16甚至8，测试是否能正常运行。如果可以，再逐步调大找到合适的数值——p2.xlarge的K80有12GB显存，32的batch size理论上是没问题的，大概率是显存分配策略导致的冲突。

3. 升级TensorFlow版本（可选） #

2018年的TF版本（比如1.8/1.9）对Eager模式的支持还不算完善，和CuDNN的配合存在一些内存管理的小bug。如果条件允许，升级到同系列的稳定版本（比如TF1.12+），后续版本对Eager下的CudnnLSTM内存优化做了不少改进。

4. 调整模型结构（备选） #

如果上面的方法都无效，可以暂时减少LSTM层数（比如先改成2层双向）或者缩小hidden size（比如降到64），验证是否是模型本身的内存占用超出了Eager模式的承载能力——不过这是最后一步，毕竟改动模型会影响业务效果。

补充说明 #

Graph模式下TensorFlow会通过静态图做全局内存规划和复用，显存使用更高效；而Eager模式是即时执行，每个操作的显存分配更直接，容易出现碎片化或提前占满显存的情况。开启显存按需分配后，TF会根据实际需要逐步申请显存，能有效避免这类workspace分配失败的问题。

内容的提问来源于stack exchange，提问作者tangy