【Bug已解决】PyTorch RuntimeError: DataLoader worker (pid(s) 15332) exited unexpectedly 解决方案

【Bug已解决】PyTorch RuntimeError: DataLoader worker (pid(s) 15332) exited unexpectedly 解决方案问题描述在使用 PyTorch 的DataLoader进行多进程数据加载时开发者经常会遇到一个令人头疼的错误RuntimeError: DataLoader worker (pid(s) 15332) exited unexpectedly这个错误通常发生在训练过程中DataLoader 的子进程worker突然崩溃退出导致主进程无法继续获取数据。错误信息中的pid是子进程的进程 ID每次运行可能不同。该问题在以下场景中尤为常见使用num_workers 0的多进程数据加载在 Windows 或 macOS 系统上运行尤其是 macOS 的 fork/spawn 问题自定义 Dataset 中使用了不可序列化的对象数据集中存在损坏的文件或异常数据系统内存不足导致子进程被操作系统杀死使用了 IPC进程间通信相关的共享内存限制这个错误的棘手之处在于它往往不是每次都出现而是间歇性发生且错误信息本身并不直接告诉你子进程为什么崩溃需要进一步排查。错误复现以下代码演示了几种常见的触发DataLoader worker exited unexpectedly错误的场景import torch from torch.utils.data import Dataset, DataLoader import numpy as np import os # 场景1Dataset 中包含不可序列化的对象 class BadDataset1(Dataset): 包含不可 pickle 序列化对象的 Dataset def __init__(self): self.data list(range(100)) # 文件句柄无法被 pickle 序列化 self.file_handle open(/dev/null, r) def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] # 尝试用多进程加载 try: dataset BadDataset1() # num_workers 0 会触发多进程 dataloader DataLoader(dataset, batch_size4, num_workers2) for batch in dataloader: pass except Exception as e: print(f场景1错误: {type(e).__name__}: {e}) # 可能输出: RuntimeError: DataLoader worker (pid(s) xxx) exited unexpectedly # 场景2数据文件损坏导致子进程崩溃 class BadDataset2(Dataset): 读取损坏文件的 Dataset def __init__(self, file_list): self.file_list file_list def __len__(self): return len(self.file_list) def __getitem__(self, idx): # 尝试读取可能损坏的文件 file_path self.file_list[idx] try: # 如果文件损坏np.load 会抛出异常 data np.load(file_path) return data except Exception as e: # 在子进程中未捕获的异常会导致 worker 崩溃 raise RuntimeError(f读取文件 {file_path} 失败: {e}) # 场景3内存不足导致子进程被杀 class MemoryHogDataset(Dataset): 每个 worker 分配大量内存的 Dataset def __init__(self, size1000): self.size size # 在 __init__ 中预分配大数组 # 每个 worker 都会复制这份内存fork 模式下 self.huge_array np.random.randn(10000, 10000) # ~800MB def __len__(self): return self.size def __getitem__(self, idx): return self.huge_array[idx % 10000] # 场景4macOS 上的 fork 问题 # macOS 默认使用 spawn 模式如果代码中使用了全局变量 # 或在模块顶层创建了对象spawn 模式下会重新导入模块 # 可能导致重复初始化或死锁 # global_model SomeModel() # 模块级全局变量 # # class GlobalVarDataset(Dataset): # def __getitem__(self, idx): # return global_model.predict(idx) # 子进程无法访问父进程的全局变量根因分析1. 多进程序列化Pickle失败DataLoader 使用 Python 的multiprocessing模块创建子进程。在 fork 模式下子进程继承父进程的内存空间但在 spawn 模式下Windows 和 macOS Python 3.8 默认子进程需要重新导入模块并重新初始化 Dataset 对象。如果 Dataset 中包含不可 pickle 序列化的对象如文件句柄、数据库连接、锁、CUDA 张量等子进程启动时就会失败。# 不可序列化的对象类型 # - 打开的文件句柄 (open()) # - 数据库连接 (sqlite3.Connection, psycopg2.Connection) # - 线程锁 (threading.Lock) # - CUDA 张量 (在多进程间传递 CUDA 张量需要特殊处理) # - lambda 函数 # - 内嵌函数/闭包2. 子进程中的未捕获异常当__getitem__方法中抛出异常且未被正确处理时子进程会崩溃。常见原因包括文件路径不存在或权限不足数据格式不匹配如期望 PNG 但实际是 JPEG数组索引越界数据预处理中的数值错误如除以零、NaN3. 共享内存限制在 Linux 系统上DataLoader 使用/dev/shm共享内存进行进程间通信。如果/dev/shm的大小限制太小默认可能是 64MB当传输的数据量超过限制时子进程会崩溃。4. 内存不足OOM每个 worker 都会复制 Dataset 对象在 fork 模式下是写时复制在 spawn 模式下是完全重建。如果 Dataset 在__init__中加载了大量数据到内存多个 worker 同时运行时可能导致系统内存不足操作系统会杀死子进程OOM Killer。5. macOS 的启动方式问题从 Python 3.8 开始macOS 默认使用spawn而非fork作为多进程启动方式。spawn模式会重新导入主模块如果主模块中有顶层代码如创建 Dataset、初始化模型等会导致重复执行。此外spawn模式要求所有通过DataLoader传递的对象都必须可 pickle。6. CUDA 与多进程的冲突在子进程中初始化 CUDA 上下文会导致问题。CUDA 上下文是进程级别的不能在 fork 的子进程中继承父进程的 CUDA 上下文。如果在 Dataset 的__getitem__中将数据移到 GPU或者使用了 CUDA 加速的数据增强库如kornia可能会导致子进程崩溃。解决方案方案一设置num_workers0进行调试首先将num_workers设为 0使用单进程加载数据。这样可以获得完整的错误堆栈信息便于定位问题from torch.utils.data import DataLoader # 使用 num_workers0 排查问题 dataloader DataLoader( dataset, batch_size32, num_workers0, # 单进程模式错误信息更清晰 shuffleTrue ) # 在单进程模式下运行观察是否有异常 for i, batch in enumerate(dataloader): print(fBatch {i}: {batch.shape}) if i 10: break方案二修复 Dataset 的序列化问题确保 Dataset 中不包含不可序列化的对象。对于文件句柄和数据库连接应在__getitem__中按需打开和关闭import torch from torch.utils.data import Dataset, DataLoader import numpy as np import os from PIL import Image class CorrectImageDataset(Dataset): 正确的图像数据集实现。 避免在 __init__ 中打开文件句柄。 def __init__(self, image_dir, transformNone): self.image_dir image_dir self.transform transform # 只存储文件路径列表不打开文件 self.image_files [ f for f in os.listdir(image_dir) if f.endswith((.png, .jpg, .jpeg)) ] # 按文件名排序确保顺序一致 self.image_files.sort() def __len__(self): return len(self.image_files) def __getitem__(self, idx): # 在 __getitem__ 中按需打开文件 img_path os.path.join(self.image_dir, self.image_files[idx]) try: image Image.open(img_path).convert(RGB) except Exception as e: print(f警告: 无法加载图像 {img_path}: {e}) # 返回一个默认图像或跳过 image Image.new(RGB, (224, 224), colorgray) if self.transform: image self.transform(image) return image, idx # 返回图像和索引 class CorrectNumpyDataset(Dataset): 正确的 NumPy 数据集实现。 延迟加载 npy 文件不在 __init__ 中全部加载。 def __init__(self, data_dir): self.data_dir data_dir # 只存储文件路径 self.file_list [ os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith(.npy) ] self.file_list.sort() def __len__(self): return len(self.file_list) def __getitem__(self, idx): # 每次按需加载 data np.load(self.file_list[idx]) return torch.from_numpy(data).float()方案三增加错误处理和容错机制在__getitem__中添加完善的错误处理避免单个数据异常导致整个 worker 崩溃import torch from torch.utils.data import Dataset, DataLoader import numpy as np import os from PIL import Image import random class RobustDataset(Dataset): 具有容错机制的数据集。 单个样本加载失败时返回替代数据而非崩溃。 def __init__(self, file_list, default_shape(3, 224, 224)): self.file_list file_list self.default_shape default_shape # 预检查文件是否存在 self.valid_files [f for f in file_list if os.path.exists(f)] if len(self.valid_files) len(self.file_list): print(f警告: {len(self.file_list) - len(self.valid_files)} 个文件不存在) def __len__(self): return len(self.valid_files) def __getitem__(self, idx): try: file_path self.valid_files[idx] # 根据文件类型加载 if file_path.endswith(.npy): data np.load(file_path) tensor torch.from_numpy(data).float() elif file_path.endswith((.png, .jpg, .jpeg)): img Image.open(file_path).convert(RGB) # 转为 tensor arr np.array(img, dtypenp.float32) / 255.0 tensor torch.from_numpy(arr).permute(2, 0, 1) else: raise ValueError(f不支持的文件类型: {file_path}) # 数据有效性检查 if torch.isnan(tensor).any() or torch.isinf(tensor).any(): raise ValueError(f数据包含 NaN 或 Inf: {file_path}) return tensor except Exception as e: # 记录错误但返回默认数据 print(f警告: 加载索引 {idx} 失败 ({e})返回默认数据) return torch.zeros(self.default_shape) def collate_fn_skip_errors(batch): 自定义 collate_fn跳过 None 样本。 # 过滤掉 None batch [item for item in batch if item is not None] if len(batch) 0: return None return torch.stack(batch) # 使用示例 def create_robust_dataloader(file_list, batch_size32, num_workers4): dataset RobustDataset(file_list) dataloader DataLoader( dataset, batch_sizebatch_size, num_workersnum_workers, shuffleTrue, collate_fncollate_fn_skip_errors, # 以下参数有助于稳定性 pin_memoryTrue, # CUDA 预分配 persistent_workersTrue if num_workers 0 else False, # 保持 worker 存活 prefetch_factor4 if num_workers 0 else None, # 预取批次  ) return dataloader方案四调整共享内存和系统限制在 Linux 上如果共享内存不足可以调整/dev/shm大小或使用文件系统作为临时目录import torch from torch.utils.data import DataLoader import tempfile import os # 方法1增大共享内存需要 root 权限 # 在 Docker 中运行时启动时添加 --shm-size 参数 # docker run --shm-size8g ... # 方法2使用文件系统作为多进程通信的临时目录 # 设置 TMPDIR 环境变量 temp_dir tempfile.mkdtemp() os.environ[TMPDIR] temp_dir # 方法3在代码中设置 torch.multiprocessing.set_sharing_strategy(file_system) # file_system 策略使用文件而非共享内存更稳定但稍慢 dataloader DataLoader( dataset, batch_size32, num_workers4, shuffleTrue, # 使用 file_system 共享策略 ) # 方法4限制每个 worker 的内存使用 # 通过在 Dataset 中使用内存映射而非全量加载 class MMapDataset(torch.utils.data.Dataset): 使用内存映射加载大文件 def __init__(self, npy_file): # mmap_moder 使用内存映射不实际加载到内存 self.data np.load(npy_file, mmap_moder) def __len__(self): return len(self.data) def __getitem__(self, idx): # 按需读取 item self.data[idx] return torch.from_numpy(np.array(item)).float()方案五处理 macOS 的 spawn 模式问题在 macOS 上确保代码符合 spawn 模式的要求import torch from torch.utils.data import Dataset, DataLoader import torch.multiprocessing as mp # 设置多进程启动方式在代码最前面 # macOS 上推荐使用 spawn mp.set_start_method(spawn, forceTrue) class SafeDataset(Dataset): macOS 安全的 Dataset 实现 def __init__(self, data_path): # 只存储路径和元数据不存储大对象 self.data_path data_path # 可以存储小的、可序列化的配置 self.length self._get_length() def _get_length(self): 获取数据长度不加载全部数据 # 根据实际情况实现 return 1000 def __len__(self): return self.length def __getitem__(self, idx): # 按需加载数据 data self._load_item(idx) return data def _load_item(self, idx): 加载单个数据项 # 实际的数据加载逻辑 return torch.randn(3, 224, 224) # 关键所有代码必须在 if __name__ __main__: 下 if __name__ __main__: dataset SafeDataset(/path/to/data) dataloader DataLoader( dataset, batch_size32, num_workers4, shuffleTrue, persistent_workersTrue, # 避免每个 epoch 重新创建 worker ) for epoch in range(10): for batch in dataloader: # 训练代码 pass完整修复代码以下是一个完整的、生产级别的 DataLoader 配置包含了所有最佳实践import torch from torch.utils.data import Dataset, DataLoader import numpy as np import os from PIL import Image import torch.multiprocessing as mp import signal import warnings import logging # 配置日志 logging.basicConfig(levellogging.INFO, format%(asctime)s - %(levelname)s - %(message)s) logger logging.getLogger(__name__) class ProductionDataset(Dataset): 生产级别的 Dataset 实现。 包含完整的错误处理、延迟加载和资源管理。 def __init__(self, data_list, transformNone, cache_modeno_cache): Args: data_list: 数据文件路径列表 transform: 数据变换 cache_mode: 缓存模式 (no_cache, memory, disk) self.data_list data_list self.transform transform self.cache_mode cache_mode self._cache {} # 内存缓存字典 # 验证文件存在性 self._validate_files() # 不在 __init__ 中打开任何文件句柄 # 不存储不可序列化的对象 logger.info(fDataset 初始化完成共 {len(self.data_list)} 个样本) def _validate_files(self): 验证文件存在性过滤无效文件 valid_files [] for f in self.data_list: if os.path.exists(f): valid_files.append(f) else: logger.warning(f文件不存在: {f}) self.data_list valid_files def __len__(self): return len(self.data_list) def __getitem__(self, idx): # 检查缓存 if self.cache_mode memory and idx in self._cache: return self._cache[idx] try: # 加载数据 data self._load_data(idx) # 应用变换 if self.transform: data self.transform(data) # 缓存 if self.cache_mode memory: self._cache[idx] data return data except Exception as e: logger.error(f加载样本 {idx} 失败: {e}) # 返回零张量作为替代 return torch.zeros(3, 224, 224) def _load_data(self, idx): 加载单个数据样本 file_path self.data_list[idx] if file_path.endswith(.npy): data np.load(file_path) return torch.from_numpy(data).float() elif file_path.endswith((.png, .jpg, .jpeg)): img Image.open(file_path).convert(RGB) arr np.array(img, dtypenp.float32) / 255.0 return torch.from_numpy(arr).permute(2, 0, 1) elif file_path.endswith(.pt): return torch.load(file_path, map_locationcpu) else: raise ValueError(f不支持的文件类型: {file_path}) class WorkerInit: Worker 初始化函数在每个子进程启动时调用 staticmethod def init_worker(worker_id): 设置 worker 的随机种子和信号处理 # 设置随机种子确保每个 worker 不同 worker_seed torch.initial_seed() % 2**32 np.random.seed(worker_seed) # 设置信号处理优雅退出 def handler(signum, frame): logger.info(fWorker {worker_id} 收到信号 {signum}正在退出...) exit(0) signal.signal(signal.SIGTERM, handler) # 忽略 SIGPIPE 信号 try: signal.signal(signal.SIGPIPE, signal.SIG_IGN) except AttributeError: pass # Windows 没有 SIGPIPE logger.info(fWorker {worker_id} 初始化完成随机种子: {worker_seed}) def create_production_dataloader( dataset, batch_size32, num_workers4, shuffleTrue, pin_memoryTrue, drop_lastFalse ): 创建生产级别的 DataLoader。 Args: dataset: Dataset 实例 batch_size: 批次大小 num_workers: 工作进程数 shuffle: 是否打乱 pin_memory: 是否预分配 CUDA 内存 drop_last: 是否丢弃最后不完整的批次 Returns: DataLoader 实例 # 设置多进程共享策略 try: torch.multiprocessing.set_sharing_strategy(file_system) except RuntimeError: pass # 可能已经设置过 dataloader DataLoader( dataset, batch_sizebatch_size, num_workersnum_workers, shuffleshuffle, pin_memorypin_memory, drop_lastdrop_last, # 保持 worker 存活避免每个 epoch 重新创建 persistent_workersTrue if num_workers 0 else False, # 预取因子 prefetch_factor4 if num_workers 0 else None, # Worker 初始化函数 worker_init_fnWorkerInit.init_worker, # 超时设置秒 timeout60 if num_workers 0 else 0, ) return dataloader def safe_train_loop(dataloader, model, optimizer, criterion, device, epochs10): 安全的训练循环包含 DataLoader 错误恢复机制。 model model.to(device) for epoch in range(epochs): logger.info(fEpoch {epoch1}/{epochs}) # 使用迭代器并支持错误恢复 data_iter iter(dataloader) batch_idx 0 consecutive_errors 0 max_consecutive_errors 5 while batch_idx len(dataloader): try: batch next(data_iter) if batch is None: logger.warning(fBatch {batch_idx} 为 None跳过) batch_idx 1 continue # 训练步骤 data, target batch data, target data.to(device), target.to(device) optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() if batch_idx % 100 0: logger.info(f Batch {batch_idx}/{len(dataloader)}, Loss: {loss.item():.4f}) consecutive_errors 0 batch_idx 1 except StopIteration: break except RuntimeError as e: if exited unexpectedly in str(e): consecutive_errors 1 logger.error(fWorker 崩溃 (连续 {consecutive_errors} 次): {e}) if consecutive_errors max_consecutive_errors: logger.error(连续错误次数过多降低 worker 数量重试) # 降低 num_workers 重试 dataloader DataLoader( dataloader.dataset, batch_sizedataloader.batch_size, num_workers0, # 降级为单进程 shuffleTrue, ) data_iter iter(dataloader) consecutive_errors 0 continue else: raise e except Exception as e: logger.error(f训练异常: {e}) raise logger.info(fEpoch {epoch1} 完成) # 主程序入口 if __name__ __main__: # 设置多进程启动方式 try: mp.set_start_method(spawn, forceTrue) except RuntimeError: pass # 创建模拟数据文件 os.makedirs(/tmp/demo_data, exist_okTrue) for i in range(100): data np.random.randn(3, 224, 224).astype(np.float32) np.save(f/tmp/demo_data/sample_{i:04d}.npy, data) # 创建文件列表 file_list [ f/tmp/demo_data/sample_{i:04d}.npy for i in range(100) ] # 创建 Dataset dataset ProductionDataset(file_list) # 创建 DataLoader dataloader create_production_dataloader( dataset, batch_size16, num_workers4, shuffleTrue, ) # 创建简单模型 model torch.nn.Sequential( torch.nn.Flatten(), torch.nn.Linear(3 * 224 * 224, 10), ) optimizer torch.optim.Adam(model.parameters(), lr1e-3) criterion torch.nn.CrossEntropyLoss() # 运行训练 device torch.device(cuda if torch.cuda.is_available() else cpu) # 生成模拟标签 class LabeledDataset(Dataset): def __init__(self, base_dataset, num_classes10): self.base base_dataset self.num_classes num_classes def __len__(self): return len(self.base) def __getitem__(self, idx): data self.base[idx] label idx % self.num_classes return data, label labeled_dataset LabeledDataset(dataset) labeled_dataloader create_production_dataloader( labeled_dataset, batch_size16, num_workers4, ) safe_train_loop( labeled_dataloader, model, optimizer, criterion, device, epochs3, ) logger.info(训练完成) # 清理 import shutil shutil.rmtree(/tmp/demo_data, ignore_errorsTrue)常见陷阱与注意事项1.persistent_workers与 epoch 长度persistent_workersTrue会让 worker 在 epoch 之间保持存活避免重复初始化的开销。但要注意如果 Dataset 在训练过程中发生变化如在线数据增强的参数更新worker 不会感知到这些变化。2.pin_memory的适用条件pin_memoryTrue可以加速 CPU 到 GPU 的数据传输但只在目标设备是 CUDA 时有效。如果使用 CPU 训练设置pin_memoryTrue反而会增加内存开销。3.prefetch_factor的调优prefetch_factor控制每个 worker 预取的批次数量。默认值为 2。增大可以提高数据吞吐量但会增加内存使用。对于大数据集或高分辨率图像建议适当增大对于小数据集保持默认即可。4. Windows 下的if __name__ __main__要求在 Windows 上多进程使用 spawn 模式所有创建 DataLoader 和启动训练的代码必须放在if __name__ __main__:块中否则会导致无限递归导入。5. CUDA 张量不能在 Dataset 中返回不要在__getitem__中将数据移到 GPU。DataLoader 的 worker 运行在 CPU 进程中返回的数据通过共享内存传输到主进程。CUDA 张量不能跨进程传递。应在训练循环中将 batch 移到 GPU# 错误在 __getitem__ 中使用 CUDA # def __getitem__(self, idx): # return self.data[idx].cuda() # 正确在训练循环中移到 GPU for batch in dataloader: batch batch.to(device)6.worker_init_fn中的随机种子如果不设置worker_init_fn所有 worker 可能使用相同的随机种子导致数据增强的随机性失效。务必在worker_init_fn中为每个 worker 设置不同的随机种子。总结DataLoader worker exited unexpectedly是一个多因素导致的错误需要系统性地排查。核心解决策略如下首先用num_workers0排查单进程模式下错误信息更清晰能暴露 Dataset 中的实际问题。确保 Dataset 可序列化不在__init__中存储文件句柄、数据库连接等不可 pickle 的对象。添加完善的错误处理在__getitem__中使用 try-except避免单个样本异常导致 worker 崩溃。调整系统资源限制增大共享内存、使用file_system共享策略、限制内存使用。适配操作系统macOS/Windows 使用 spawn 模式代码需放在if __name__ __main__下。使用生产级配置persistent_workers、prefetch_factor、worker_init_fn等参数的合理配置。通过以上方法可以系统性地解决 DataLoader worker 崩溃问题确保训练过程的稳定性和可靠性。