TensorFlow处理大尺寸视频时出现内存耗尽OOM报错如何解决

问题根因 #

你就算指定算子跑CPU依然报GPU OOM，是因为TensorFlow默认会预占几乎全部可用GPU显存，即便运算逻辑不在GPU执行，显存也会被预占满。同时循环处理每帧时，TensorFlow的中间张量没有被及时释放，累积占用显存，最终触发OOM。

可落地解决方案 #

1. 优先调整TensorFlow显存配置 #

在代码最开头（导入TensorFlow后立即执行）添加显存动态分配配置，禁止TensorFlow预占全部GPU显存：

import tensorflow as tf
# 开启显存动态分配
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

如果确认要完全切换到CPU运行，额外添加配置隐藏GPU，从根源避免TensorFlow占用GPU显存：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

2. 优化循环内资源占用 #

你当前的循环逻辑会持续累积无用的张量和缓存，可做如下调整：

• 每处理完一帧，手动清空TensorFlow的后端缓存，释放无用显存：
  单帧推理完成后添加一行 tf.keras.backend.clear_session()
• 正式运行时注释所有plt.imshow相关代码，matplotlib的图像缓存会额外占用内存/显存
• 不要在循环外留存output、panoptic_map这类临时变量，处理完直接覆盖即可

3. 你提到的分段处理+落盘合并方案完全可行，具体实现逻辑如下： #

先将视频拆分为小于200帧的分段，每处理完一个分段就写入磁盘、清空缓存，再处理下一段，最后合并所有分段即可：

cap = cv2.VideoCapture(filename)
fourcc = cv2.VideoWriter_fourcc(*'MP42')
fps = 20.0
frame_size = (1656, 949)

# 获取视频总帧数
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# 每段150帧，留足余量避免触发OOM
segment_frames = 150
seg_index = 0
segment_files = []

for seg_start in range(0, total_frames, segment_frames):
    # 跳转到当前分段的起始帧位置
    cap.set(cv2.CAP_PROP_POS_FRAMES, seg_start)
    seg_filename = f'temp_seg_{seg_index}.avi'
    segment_files.append(seg_filename)
    seg_out = cv2.VideoWriter(seg_filename, fourcc, fps, frame_size)
    
    # 处理当前分段的所有帧
    for _ in range(segment_frames):
        ret, image = cap.read()
        if not ret:
            break
        # 原有帧处理逻辑
        image = cv2.resize(image, frame_size)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  
        output = LOADED_MODEL(tf.cast(image, tf.uint8))
        panoptic_map, used_colors = color_panoptic_map(output['panoptic_pred'][0], DATASET_INFO, 60)
        dst = cv2.addWeighted(image, 0.5, panoptic_map, 1, 1)
        seg_out.write(cv2.cvtColor(dst, cv2.COLOR_RGB2BGR))
    
    seg_out.release()
    tf.keras.backend.clear_session()
    seg_index += 1

# 合并所有分段到最终输出视频
out = cv2.VideoWriter('output.avi', fourcc, fps, frame_size)
for seg_file in segment_files:
    seg_cap = cv2.VideoCapture(seg_file)
    while seg_cap.isOpened():
        ret, frame = seg_cap.read()
        if not ret:
            break
        out.write(frame)
    seg_cap.release()
    # 可选：处理完直接删除临时分段文件
    os.remove(seg_file)

out.release()
cap.release()

4. 额外可选优化 #

• 若业务允许降低输出分辨率，可将输入帧从1656949缩小到1280720，显存占用可降低40%以上
• 开启TensorFlow混合精度推理，使用float16替代float32做运算，显存占用直接减半
• 推理时指定batch size为1，避免模型自动累积批次占用显存

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