在HPC（Argon）集群运行Keras ConvLSTM时GPU内存不足如何解决？

嘿，我太懂这种挫败感了——CPU上跑ConvLSTM顺风顺水，一扔GPU队列就爆显存，尤其是看到“Allocator (GPU_0_bfc) ran out of memory”这种报错，简直头大。其实这很正常，GPU的显存通常比CPU的系统内存小不少，而且ConvLSTM因为要处理序列式的空间数据，本身对内存（显存）的需求就比普通CNN高。给你几个实操性强的解决方案：

• 先砍批量大小试试
  这是最快见效的办法。GPU显存里要同时装下整个批次的输入数据、模型参数以及中间计算结果，把batch_size从原来的数值（比如32）直接减半甚至减到8，大概率能缓解显存压力。你只需要在模型训练的fit方法里调整参数就行：

  model.fit(X_train, y_train, batch_size=8, epochs=10)
  

• 给模型“瘦个身”
  看看你的ConvLSTM层是不是堆了太多过滤器？比如把ConvLSTM2D(filters=64)改成filters=32，或者减少一层ConvLSTM，都能大幅降低显存占用。另外，检查下模型里有没有冗余的全连接层，能删就删，毕竟全连接层的参数数量可是实打实的内存大户。

• 让TensorFlow按需用显存
  默认情况下TensorFlow会直接占满整个GPU的显存，哪怕模型根本用不上这么多。你可以在代码开头加一段配置，让它按需分配显存：

  import tensorflow as tf
  from keras.backend.tensorflow_backend import set_session
  
  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True  # 开启显存动态分配
  sess = tf.Session(config=config)
  set_session(sess)
  

  这样TensorFlow只会在需要的时候申请显存，不会一开始就把GPU显存吃光。

• 缩小输入数据的“体积”
  ConvLSTM处理的一般是序列图像数据（比如(时间步, 高度, 宽度, 通道)），如果你的输入图像分辨率太高（比如256x256）、时间步太长，显存肯定扛不住。试试把图像缩到128x128，或者减少序列的时间步数量，比如从30步降到15步，显存占用会直接跟着降下来。

• 试试混合精度训练
  如果你的GPU是NVIDIA Turing架构及以上（比如RTX 20系列及以后），可以开启混合精度训练，把部分模型参数从float32换成float16，显存占用直接砍半。在Keras里设置也很简单：

  from keras.mixed_precision import set_global_policy
  set_global_policy('mixed_float16')
  

  注意要确保模型的最后一层输出是float32，避免出现精度损失的问题。

对了，顺便说一句为什么CPU能跑——CPU的系统内存通常比单块GPU的显存大得多，而且CPU的内存管理机制更灵活，所以能扛住更大的内存需求，但GPU就没这么宽松啦。

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