从TensorFlow 1.x迁移至2.x后Google AI Platform训练变慢且报ResourceExhaustedError

• TF2.x默认Eager执行模式开销：TF1.14默认基于静态图执行，训练前会做完整的图优化，而TF2.x默认开启Eager执行模式，运行时动态计算图，额外开销极高，对LSTM这类循环层的性能影响尤为明显，同时会带来更高的显存占用，容易触发ResourceExhaustedError。
• LSTM层CuDNN加速未正确启用：TF1.x中通常使用CuDNNLSTM层调用GPU加速实现，迁移到TF2.x后如果直接替换为通用LSTM层且没有匹配激活函数、参数配置（比如使用会禁用CuDNN加速的正则、约束参数），会 fallback 到无加速的实现，直接导致速度暴跌5-10倍、显存占用翻倍。
• Google AI Platform TF2.5 runtime默认配置缺陷：TF2.5 runtime默认的显存分配策略是贪心占满所有GPU显存，同时默认关闭了部分图优化开关，相比TF1.14 runtime的优化配置会产生额外的显存浪费和性能损耗。
• Keras模型编译/训练参数未适配TF2.x：未启用XLA编译、混合精度训练，或者tf.data数据流水线没有做prefetch、cache等适配，都会进一步放大性能差距。

• 启用静态图执行：在模型训练代码最开头加入tf.compat.v1.disable_eager_execution()，或者用@tf.function装饰器包装训练步，完全复现TF1.x的静态图优化效果，显存占用和执行速度都会回到接近TF1.14的水平。
• 调整LSTM层配置启用CuDNN加速：确保所有LSTM层满足以下参数要求，即可自动调用CuDNN加速实现：
  • 激活函数设置为activation='tanh'、recurrent_activation='sigmoid'
  • 不设置recurrent_dropout、kernel_constraint、recurrent_constraint、bias_constraint参数
    若必须使用recurrent dropout等特性，可以额外指定LSTM(..., implementation=2)强制使用高性能GPU实现。
• 优化显存分配：在代码开头加入以下配置，避免无效显存占用：

  gpus = tf.config.list_physical_devices('GPU')
  if gpus:
    try:
      for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
      print(e)
  

• 开启XLA编译和混合精度训练：编译模型时指定jit_compile=True，同时开启混合精度训练，可额外提升30%-50%的训练速度，进一步降低显存占用：

  # 开启混合精度
  policy = tf.keras.mixed_precision.Policy('mixed_float16')
  tf.keras.mixed_precision.set_global_policy(policy)
  # 编译模型时开启XLA
  model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'], jit_compile=True)
  

• 优化tf.data数据流水线：确保数据读取部分加入prefetch(tf.data.AUTOTUNE)，如果数据集不大可额外加入cache()，避免数据读取成为训练瓶颈。

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