原帖使用 Windows 10、PyTorch 1.4.0;num_workers=0 后又出现约 9 GB CPU 分配失败或 GPU OOM:
RuntimeError: DataLoader worker (pid(s) 15332) exited unexpectedly
单进程是定位入口,内存需求才是下一步
原题已经提供 num_workers=0 后申请 9663676416 字节失败的证据,约为 9 GiB 单次分配;另一路还在 4 GiB GPU 上 OOM。因此不能停在常见的“设 workers=0”。先判断读进来的是否是预期训练裁块,再看 batch 和中间特征规模。
worker exited 是子进程退出的外层提示。原帖后续错误已经表明,还需要处理实际内存需求,单把 worker 数设成 0 并不是完整修复。
在创建训练 loader 后,先用原数据集构建一个 batch_size=1、单进程的诊断 loader,打印真正送入网络的尺寸。training_data_loader 是原训练 loader:
from torch.utils.data import DataLoader
probe = DataLoader(training_data_loader.dataset, batch_size=1, num_workers=0)
inputs, targets = next(iter(probe))
print("josh050608")
print("input:", inputs.shape, inputs.dtype)
print("target:", targets.shape, targets.dtype)
print("input bytes:", inputs.numel() * inputs.element_size())
输入字节数不包含中间特征和梯度,不能拿它当训练峰值。若图片远大于预期,检查超分辨率数据集是否需要随机裁块、是否误把整图作为训练样本,以及 upscale_factor 与裁块尺寸的关系。对超分辨率,输入与高清目标必须成对裁切,不能各自随机裁剪。
先降低 batch 和有效训练裁块尺寸,保持 CPU 加载数据,在主进程把一个 batch 移到 GPU。Windows 脚本还应把训练入口放进 if __name__ == '__main__':,数据集类保持在模块顶层。
先验证 dataset[0]、再验证 next(iter(probe))、最后做一个训练 step,分别定位读取、组批还是前向耗尽内存。单进程跑通后再逐步增加 workers;若只在多进程失败,继续检查可序列化对象和 worker 内存复制。
超分辨率必须成对裁块
若高清目标是输入的 s 倍,输入裁块 [h,w] 应对应目标 [s*h,s*w]。两个独立随机裁剪会破坏像素对应关系,虽然可能省内存,却会使监督目标错误。下面用合成图演示坐标对应,不读取数据、不申请大内存:
import torch
scale = 2
low = torch.arange(3 * 16 * 20, dtype=torch.float32).reshape(3, 16, 20)
high = low.repeat_interleave(scale, dim=1).repeat_interleave(scale, dim=2)
top, left, crop_h, crop_w = 3, 4, 8, 8
low_patch = low[:, top:top+crop_h, left:left+crop_w]
high_patch = high[:, top*scale:(top+crop_h)*scale,
left*scale:(left+crop_w)*scale]
assert low_patch.shape == (3, 8, 8)
assert high_patch.shape == (3, 16, 16)
assert torch.equal(high_patch[:, ::scale, ::scale], low_patch)
print("paired patches:", low_patch.shape, high_patch.shape)
原工程可能使用 Y 通道或先降采样得到低清输入,应按其真实数据生成顺序选裁块位置,不能直接假设输入是三通道。把 dataset[0] 的结果与数据集内部 crop 参数对照;若原本想训练小块,实际却返回整张高分辨率图,应在数据预处理处改,而不是事后改变张量 view。
| 哪一步失败 | 优先查看 |
|---|---|
dataset[0] 就内存不足 |
单图解码、预处理、中间数组是否过大 |
| 单样本正常,组 batch 失败 | batch_size、collate 与 CPU 内存 |
| batch 正常,前向或反向 OOM | 裁块、激活、模型规模及 GPU 峰值 |
| 单进程正常,多 worker 失败 | Windows 入口保护、可序列化定义与进程内存 |
Windows 的 if __name__ == '__main__': 应包住创建 loader 和启动训练的入口,类定义放模块顶层。先以小裁块和 batch=1 完成一次损失、反向与更新,再逐级扩大;每次只改一个资源参数,记录哪一步达到内存上限。CPU 张量输入大小只是下界,无法代表模型训练峰值。
参考:原始求助帖;技术依据:PyTorch 数据加载(2.5)、Python multiprocessing、PyTorch CUDA 语义(2.5)。