如何用Python threading/multiprocessing高效并行运行深度学习模型做视频检测？

问题核心分析 #

之前使用多线程/多进程时推理时间过长，通常是以下原因导致：

• Python GIL限制：CPU密集型推理任务下，多线程无法真正并行；
• GPU资源竞争：多个模型抢占同一GPU显存或计算资源，导致上下文切换开销；
• 重复加载模型：每个线程/进程重复加载模型，浪费内存和初始化时间；
• CPU线程冲突：深度学习框架默认启用大量CPU线程，多个模型同时运行时会抢占CPU资源，拖慢整体速度。

多进程方案（推荐用于多GPU或CPU+GPU混合场景） #

多进程可以绕过GIL限制，每个进程拥有独立的Python解释器和内存空间，适合模型推理这种计算密集型任务。关键是让每个进程独立加载一个模型，并分配专属的计算资源。

实现步骤 #

1. 为每个进程分配独立的GPU（如果有多GPU），或限制单GPU下的显存使用；
2. 每个进程独立加载模型，避免跨进程共享模型（会导致序列化开销）；
3. 每个进程处理一个视频流，独立完成读取、推理、输出。

代码示例（PyTorch为例） #

import multiprocessing
import cv2
import torch

def run_model(model_path, video_path, gpu_id):
    # 设置当前进程使用的GPU
    torch.cuda.set_device(gpu_id)
    # 限制显存使用，避免抢占
    torch.cuda.set_per_process_memory_fraction(0.45, device=gpu_id)
    
    # 加载模型
    model = torch.load(model_path).to(gpu_id).eval()
    
    # 读取视频
    cap = cv2.VideoCapture(video_path)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # 预处理帧
        frame_tensor = torch.from_numpy(frame).permute(2,0,1).unsqueeze(0).float().to(gpu_id) / 255.0
        
        # 推理（关闭梯度计算加速）
        with torch.no_grad():
            outputs = model(frame_tensor)
        
        # 处理输出（此处省略具体检测逻辑）
        print(f"Model {model_path.split('/')[-1]} processed frame, GPU usage: {torch.cuda.memory_allocated(gpu_id)/1024**2:.2f}MB")
    
    cap.release()

if __name__ == "__main__":
    # 两个模型路径和视频路径
    task1 = ("model1.pt", "video1.mp4", 0)
    task2 = ("model2.pt", "video2.mp4", 1)  # 如果只有单GPU，这里改为0
    
    # 创建进程
    p1 = multiprocessing.Process(target=run_model, args=task1)
    p2 = multiprocessing.Process(target=run_model, args=task2)
    
    # 启动进程
    p1.start()
    p2.start()
    
    # 等待进程结束
    p1.join()
    p2.join()

优化要点 #

• 多GPU场景：给每个进程指定不同的gpu_id，避免资源竞争；
• 单GPU场景：通过set_per_process_memory_fraction限制每个进程的显存占用，确保两个模型都能正常加载；
• 模型加载：每个进程独立加载模型，不要在主进程加载后传递给子进程（会触发模型序列化，耗时且易出错）。

多线程方案（适合IO密集型或GPU异步推理场景） #

如果模型推理主要依赖GPU（GPU操作会释放GIL），多线程可以有效利用GPU的异步计算能力，同时处理视频读取的IO操作。但需注意模型的线程安全性。

实现步骤 #

1. 在主线程提前加载两个模型（确保模型支持线程安全，如PyTorch的eval()模式下大部分模型是线程安全的）；
2. 使用线程池或threading.Thread分配视频处理任务；
3. 限制每个模型的CPU线程数，避免抢占资源。

代码示例（TensorFlow为例） #

import threading
import cv2
import tensorflow as tf

# 限制TensorFlow的CPU线程数，避免冲突
tf.config.threading.set_intra_op_parallelism_threads(2)
tf.config.threading.set_inter_op_parallelism_threads(2)

# 加载两个模型（主线程加载，线程共享）
model1 = tf.keras.models.load_model("model1.h5")
model2 = tf.keras.models.load_model("model2.h5")

def run_detection(model, video_path):
    cap = cv2.VideoCapture(video_path)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # 预处理帧
        frame_resized = cv2.resize(frame, (224, 224))
        frame_tensor = tf.convert_to_tensor(frame_resized, dtype=tf.float32) / 255.0
        frame_tensor = tf.expand_dims(frame_tensor, axis=0)
        
        # 推理
        outputs = model.predict(frame_tensor, verbose=0)
        
        # 处理输出
        print(f"Model processed frame, output shape: {outputs.shape}")
    
    cap.release()

if __name__ == "__main__":
    # 创建线程
    t1 = threading.Thread(target=run_detection, args=(model1, "video1.mp4"))
    t2 = threading.Thread(target=run_detection, args=(model2, "video2.mp4"))
    
    # 启动线程
    t1.start()
    t2.start()
    
    # 等待线程结束
    t1.join()
    t2.join()

优化要点 #

• 线程安全：确保模型在eval()模式下运行，避免训练时的非线程安全操作（如BatchNorm的更新）；
• CPU线程限制：通过框架API限制每个模型的CPU线程数，防止多个模型抢占CPU资源；
• 异步推理：TensorFlow的predict()默认支持异步，PyTorch可配合torch.jit.trace优化模型，提升推理速度。

通用性能优化措施 #

1. 视频读取优化：使用cv2.VideoCapture.set(cv2.CAP_PROP_BUFFERSIZE, 1)减少视频缓存，或用多线程单独读取视频帧，避免推理线程等待IO；
2. 模型量化：将模型转换为FP16或INT8格式（如PyTorch的torch.quantization，TensorFlow的tf.lite），减少显存占用和推理时间；
3. 避免不必要的计算：推理时关闭梯度计算（PyTorch的torch.no_grad()，TensorFlow的tf.stop_gradient()）；
4. 显存动态分配：启用框架的显存动态增长（TensorFlow的tf.config.experimental.set_memory_growth(gpu, True)，PyTorch的torch.cuda.empty_cache()定期释放无用显存）。

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