多进程共享内存+PyTorch DataLoader的CUDA启动及数据读取问题

问题概述 #

实现生产者-消费者架构：生产者读取图像numpy数组写入共享内存，消费者读取数组执行PyTorch推理，遇到以下问题：

• 遍历DataLoader时触发CUDA初始化错误：RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
• 主程序设置set_start_method('spawn')后，消费者读取的共享内存数组全为0
• 设置num_workers:0时触发ValueError: persistent_workers option needs num_workers > 0
• 改用torch.multiprocessing仍存在CUDA启动错误

解决方案 #

1. 适配spawn启动方式，修正共享内存视图创建逻辑 #

spawn启动的子进程会重新初始化Python解释器，无法继承父进程的numpy共享内存视图，必须在子进程中重新构建视图：

• 主进程共享内存创建代码：

import multiprocessing as multi_processing
def create_shared_memory(self):
    type_code = "I"
    size = int(np.prod(self.image_frame_shape))
    frame_lock = multi_processing.Lock()
    
    shared_memory_array = multi_processing.Array(typecode_or_type=type_code, size_or_initializer=size, lock=frame_lock)
    # 主进程仅传递共享内存对象及必要参数，不传递numpy视图
    return shared_memory_array, self.image_frame_shape, type_code

• 消费者进程初始化共享内存：

def init_consumer_shared(self, shared_memory_array, frame_shape, type_code):
    # 在子进程中重新生成numpy共享内存视图
    self.buffered_array = np.frombuffer(shared_memory_array.get_obj(), dtype=type_code).reshape(frame_shape)

2. 修复DataLoader参数冲突 #

当num_workers=0时，必须关闭persistent_workers参数，否则触发参数不兼容错误：

# 动态调整DataLoader参数，避免冲突
loader_params = self.model_params.copy()
if loader_params.get("num_workers", 0) == 0:
    loader_params["persistent_workers"] = False

inference_data_loader = create_loader(
    InferCustomDataset(frame_list, self.validation_transforms, input_size=self.model_params['input_size'][1:]),
    **loader_params
)

3. 调整CUDA初始化时机 #

spawn子进程无法继承父进程的CUDA上下文，必须在消费者进程内部初始化CUDA和模型：

def run(self):
    # 先初始化共享内存视图
    self.init_consumer_shared(self.shared_memory_array, self.frame_shape, self.type_code)
    # 再加载模型到CUDA
    self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    self.model = self.load_model().to(self.device)
    self.model.eval()
    # 执行推理逻辑
    self.infer_on_frame_list(...)

4. 简化数据加载逻辑（避免嵌套多进程） #

消费者本身已是子进程，DataLoader启用多进程会导致嵌套多进程触发CUDA问题，建议：

• 设置num_workers=0，让消费者单线程处理数据加载
• 或直接绕过DataLoader，手动处理图像预处理：

def infer_on_frame(self):
    with torch.no_grad():
        # 加锁确保读取到完整的图像数据
        with self.shared_memory_array.get_lock():
            img_np = self.buffered_array.copy()
        # 执行预处理
        img = Image.fromarray(img_np).convert('RGB')
        img = img.resize(self.input_size)
        if self.transform:
            img = self.transform(img).unsqueeze(0).to(self.device)
        # 模型推理
        output = self.model(img)
        # 处理输出结果

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