TensorFlow2.2搭配tensorflow_hub多GPU运行USE遇OOM问题求解

双GPU调度实现方法 #

• 修复MirroredStrategy兼容问题：TF2.2版本的hub.load接口和分布式策略存在上下文适配缺陷，加载模型时不要直接在策略作用域内调用hub.load，先将模型下载到本地缓存后，用tf.keras.models.load_model接口在策略作用域内加载，再将拆分后的数据块用experimental_distribute_dataset接口封装为分布式数据集，自动分配到双GPU运行。
• 手动设备分配兼容性方案：如果分布式策略仍然报错，可以跳过策略封装，手动将每个待处理数据块平均拆分2份，分别通过with tf.device('/GPU:0'):和with tf.device('/GPU:1'):指定设备运行推理，再将两个GPU的输出结果合并，该方案不需要适配框架兼容问题，实现门槛最低。

OOM问题规避方案 #

• 显存主动释放：每次单块数据推理完成后，立即调用tf.Tensor.numpy()将结果转为CPU侧numpy数组，不要在GPU显存中留存历史结果，每处理完2个块手动调用tf.keras.backend.clear_session()清理显存碎片。
• 调整数据块大小：将原来的200个数据块进一步拆分到500~1000个，控制单块推理的峰值显存占用不超过单GPU显存的60%，避免批次过大导致的显存溢出。
• 开启显存动态分配：在代码初始化阶段添加如下配置，禁止TensorFlow预占全量GPU显存，减少无效显存占用：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

• 磁盘增量存储：不要在内存中拼接全量的结果向量，每处理完3~5个数据块，就将结果追加写入本地npz格式文件，所有块处理完成后再从磁盘读取合并全量结果，大幅降低内存和显存的常驻占用。

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