Windows下的num_workers > 0报错
在Windows下新搭建的PyTorch环境中使用之前的某个课程项目进行测试时，发现设置DataLoader的num_workers大于0会导致模型训练卡死。由于GPU无占用，在JupyterLab日志中找到形如AttributeError: Can't get attribute xxx on <module '__main__' (built-in)>的报错，最后定位到DataLoader上。
注意到不少回答都认为Windows下应该将num_workers设为0，但并未解释原因，实际测试发现训练过程远比进行课程项目时慢，于是进一步搜索希望能够找到更好的解决方法。
根据这个帖子的讨论以及这篇博客介绍，由于Windows多进程机制与Linux存在区别，Windows下multiprocessing不支持使用fork创建进程，而是只能通过spawn创建新的进程；这一点差异就导致导入模块时会使用模块锁，避免多个进程导入同一模块时产生竞争。
解决方案也很简单，将原来项目的内部类(是一个数据集类)放到一个外部文件/模块中，并在原来的代码中导入它。此时multiprocessing便可以正确加载该模块，确保模块只被加载一次，然后在各个进程之间共享。
Windows下的DataLoader性能问题
通过上一节的改动，DataLoader已经可以多线程地加载数据集了。但留意到增加worker数量并不能提升训练速度，CUDA核心占用依然出现间断，且间断时CPU和磁盘使用增加。可以推测认为是DataLoader加载速度跟不上训练速度，GPU经常停下来等待数据加载。
不过这也好理解：Linux下通过fork创建进程的效率显著高于Windows，在这个issue下的讨论也普遍认为Windows的进程机制拖慢了PyTorch的效率。
“那么具体一点，能不能衡量一下Windows拖慢了多少训练速度呢？”由于之前干过一段时间性能测试，这个问题马上冒了出来。好消息是在上面提到的这篇博客中，作者给出了ta测试DataLoader的代码。而我对课程项目的代码做了简单修改，作为附加的测试项目，希望能测试真实训练场景下的性能。
Windows与(Subsystem) Linux下DataLoader性能对比
测试环境
- CPU：6核心12线程
- 内存：32GB
- GPU：显存16GB
- 操作系统：
- Windows 10
- Ubuntu 22.04(以WSL运行)
- 软件：
- Python 3.10
- PyTorch 2.2
- CUDA 11.8
测试项目
- 上述博客中的MNIST数据集仅加载测试
num_workers从0至12以2递增- 仅测试加载数据集的用时
- 课程项目的AlexNet训练-验证5个epoch
- 涉及加载train、validation两个数据集
- 每个epoch在train集上训练，并在validation集上验证
- (项目写的比较挫，其他hyperparameter就不写出来丢人了，主要只是测加载用时的影响)
测试记录
测试项目1
| num_workers | WSL用时(s) | Windows用时(s) |
|---|---|---|
| 0 | 16.22 | 16.68 |
| 2 | 8.38 | 14.49 |
| 4 | 4.67 | 10.71 |
| 6 | 4.00 | 10.56 |
| 8 | 3.67 | 11.36 |
| 10 | 3.47 | 12.22 |
| 12 | 4.06 | 14.44 |
测试项目2
| num_workers | WSL用时(s) | Windows用时(s) |
|---|---|---|
| 0 | 264.60 | 271.56 |
| 2 | 167.93 | 279.79 |
| 4 | 108.87 | 288.02 |
| 6 | 92.92 | 361.40 |
| 8 | 87.95 | 444.24 |
| 10 | 89.07 | 524.86 |
| 12 | 88.19 | 614.13 |
测试结果
- 同等配置下，Linux下(即使是作为WSL运行)的DataLoader性能始终高于Windows
- Linux下当
num_workers增加到CPU核心数时，带来的性能提升逐渐接近边际，继续增大反而有可能导致性能下降 - Windows下增加worker数量可以少量提升DataLoader性能，但对于整个训练过程来说，不断创建进程的开销会显著拖慢训练速度；设置
num_workers=0可能仍然是最优解