PyTorch结合ProcessPoolExecutor多进程：避免Forkbomb及非固定形状数据处理

核心问题分析 #

你遇到的类Forkbomb问题，大概率是因为主进程提前初始化了PyTorch相关资源（比如模型、张量、数据集对象）后再创建子进程，而ProcessPoolExecutor基于fork机制，会复制主进程的全部内存空间，包括PyTorch的CUDA上下文或全局状态，导致子进程触发资源初始化时出现递归创建进程，或是直接耗尽系统资源。

可行解决方案 #

1. 避免主进程提前初始化PyTorch资源 #

把数据集初始化、PyTorch相关操作全放到子进程内部执行，主进程只传递索引、文件路径这类轻量参数，绝对不要直接传递整个数据集对象：

from concurrent.futures import ProcessPoolExecutor
import torch

def process_single_data(idx, data_dir):
    # 子进程内单独初始化数据集、加载数据
    dataset = CustomDataset(data_dir)
    data = dataset[idx]
    # 处理非固定形状数据的自定义逻辑
    processed_data = custom_preprocess(data)
    # 直接保存到磁盘，无需返回主进程
    torch.save(processed_data, f"./processed/{idx}.pt")

if __name__ == "__main__":
    data_dir = "./raw_data"
    total_samples = 1000
    # 必须在__name__ == "__main__"块内启动进程池，避免fork时重复执行模块代码
    with ProcessPoolExecutor(max_workers=4) as executor:
        executor.map(process_single_data, range(total_samples), [data_dir]*total_samples)

2. 改用PyTorch官方多进程工具 #

PyTorch的torch.multiprocessing针对自身多进程场景做了优化，能规避fork带来的CUDA上下文复制问题：

import torch.multiprocessing as mp
import torch

def process_single_data(idx, data_dir):
    dataset = CustomDataset(data_dir)
    data = dataset[idx]
    processed_data = custom_preprocess(data)
    torch.save(processed_data, f"./processed/{idx}.pt")

if __name__ == "__main__":
    data_dir = "./raw_data"
    total_samples = 1000
    with mp.Pool(processes=4) as pool:
        pool.starmap(process_single_data, [(idx, data_dir) for idx in range(total_samples)])

如果需要更灵活的进程控制，也可以用mp.Process手动创建进程，核心原则依然是子进程内单独加载数据和PyTorch资源。

3. 超大数据集的分批处理优化 #

如果数据集过大，单进程单条加载效率低，可以让每个子进程处理一个小批次的索引，减少进程创建销毁的开销：

def process_batch(batch_idx, data_dir, batch_size):
    dataset = CustomDataset(data_dir)
    start_idx = batch_idx * batch_size
    end_idx = min(start_idx + batch_size, len(dataset))
    for idx in range(start_idx, end_idx):
        data = dataset[idx]
        processed_data = custom_preprocess(data)
        torch.save(processed_data, f"./processed/{idx}.pt")

if __name__ == "__main__":
    data_dir = "./raw_data"
    total_samples = 1000
    batch_size = 50
    total_batches = (total_samples + batch_size - 1) // batch_size
    with mp.Pool(processes=4) as pool:
        pool.starmap(process_batch, [(b_idx, data_dir, batch_size) for b_idx in range(total_batches)])

关键注意事项 #

• 绝对不要在主进程中创建PyTorch张量、模型或加载数据集后再fork子进程，这是触发资源异常的核心原因。
• 所有涉及PyTorch的操作必须放在子进程内部完成，主进程只负责任务分发。
• 控制进程池的max_workers数量，建议不超过CPU核心数的1.5倍，避免系统资源耗尽。

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