Google Colab中使用cv2批量处理视频RAM溢出问题如何解决

核心溢出原因 #

你代码的内存占用过高来自四个典型问题：

1. 全局变量x通过tf.concat持续拼接所有时序样本，全量数据长期驻留内存，数据量达到阈值后直接占满RAM
2. 每个视频迭代周期生成的tf.Variable、临时张量、numpy数组、poseDetector实例没有主动释放，TensorFlow默认内存管理不会即时回收废弃张量，会产生内存堆积
3. 使用cap.set跳帧读取视频的逻辑会导致OpenCV缓存大量未读帧，产生额外内存开销
4. 全量标签y也全程存储在内存中，样本量较大时会进一步加剧内存占用

可落地优化方案 #

• 取消全局拼接大张量的逻辑，改用Python列表存储中间样本，最后再统一转张量，tf.Variable的内存开销远高于普通numpy数组/列表
• 每个视频处理完成后立刻主动释放中间变量，触发垃圾回收和TensorFlow缓存清空
• 将poseDetector初始化移到循环外复用，避免每次迭代新建实例产生的内存开销
• 优化视频读取逻辑：逐帧读取后间隔丢弃，代替cap.set跳帧，减少OpenCV的帧缓存占用
• 数据量极大的场景下，不要攒全量数据再存盘，改成每处理5-10个视频就追加写入一次磁盘，处理完成后再合并文件

优化后代码示例 #

import gc
import tensorflow as tf
import cv2
import numpy as np
from tqdm import tqdm
from tensorflow.keras.utils import to_categorical

# make dataset only for one person
def make_dataset(df, type_, seq_len=6):
    uses_keypoints = [0, 11, 12, 13, 14, 15, 16, 23, 24, 25, 26, 27, 28]
    # 改用列表存储中间数据，不要用tf.Variable
    x = []
    y = []
    # detector移到循环外复用
    detector = poseDetector()
    
    for k in tqdm(range(df.shape[0])):
        cap = cv2.VideoCapture(df.iloc[k].file_name)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        keypoints = []
        frame_idx = 0
        
        # 逐帧读取代替cap.set跳帧，减少缓存占用
        while cap.isOpened():
            success, img = cap.read()
            if not success:
                break
            # 每6帧处理1次
            if frame_idx % 6 == 0:
                lmList = detector.findPose(img)
                if lmList:
                    lmList = np.array(lmList)[uses_keypoints].astype('float32')
                    keypoints.append(lmList)
            frame_idx += 1
            # 及时释放帧内存
            del img
        
        cap.release()
        keypoints = np.array(keypoints)
        
        # 生成时序样本
        if keypoints.shape[0] >= seq_len:
            sequences = [keypoints[i-seq_len:i] for i in range(seq_len, keypoints.shape[0])]
            label = df.iloc[k].person_id
            x.extend(sequences)
            y.extend([label]*len(sequences))
        
        # 主动释放当前视频所有中间变量，清空缓存
        del keypoints, sequences, label
        gc.collect()
        tf.keras.backend.clear_session()
    
    # 最后统一转张量和独热编码
    x = tf.convert_to_tensor(x, dtype='float32')
    y = to_categorical(y)
    path_to_load = f'/some_path'
    save_dataset(x, y, path_to_load)
    # 处理完后释放全量数据内存
    del x, y
    gc.collect()

内容的提问来源于stack exchange，提问作者Макс Байдин