Python 及 PyTorch 中 Dataloder 引出的多线程问题

2023年7月11日

摘要

在 Mac 及 Windows 上使用 PyTorch 时，使用 Dataloader 经常会出现一种多线程相关的报错，报错内容也不像常规内容一样具有较强的指向性，且此时代码内容也并没有什么逻辑性错误，较难发现问题根源。本文提供了一种解决方案。

报错内容

在正常的使用 Dataloader 时，例如：

test_loader = DataLoader(dataset=test_set, num_workers=4, batch_size=1)

却出现了以下报错信息：

DataLoader worker (pid 31335) exited unexpectedly with exit code 1. Details are lost due to multiprocessing. Rerunning with num_workers=0 may give better error trace.
RuntimeError: DataLoader worker (pid(s) 32685) exited unexpectedly

仅根据报错信息分析的话，问题貌似出在了 num_workers=4 上，虽然之前的用法也完全看不出有什么问题，但还是根据提示改为了 num_workers=0 试试看。然而，问题依旧，且最不解的是，Linux 系统下同样的代码却完全能够正常运行。

问题分析

原谅作者没有什么文化，自己并没有能力分析出这种问题背后的原因，但根据多方资料汇总，最终还是能够提供一些有用的信息。

在 macOS 和 Windows 的 PyTorch 的 Dataloader 中，一旦我们开启了多线程加载数据集，都将调用 spawn 命令来生成子线程，而在 Linux 中，使用的是 fork 命令。这两个命令的不同造成了 Linux 下能运行的代码在另外两个平台上无法运行。二者的区别如下：

• fork - 除了必要的启动资源外，其他变量，包，数据等都继承自父进程，也就是共享了父进程的一些内存页，因此启动较快，但是由于大部分都用的父进程数据，所以是不安全的进程。
• spawn - 从头构建一个子进程，父进程的数据等拷贝到子进程空间内，拥有自己的 Python 解释器，所以需要重新加载一遍父进程的包，因此启动较慢，由于数据都是自己的，安全性较高。

根据以上介绍，如果经验丰富，大致就能猜到问题的原因了。

由于出问题的代码是一个普通的测试脚本，并未做各种软件工程相关的规范约束，因此，在 macOS 及 Windows 中，由于 spawn 命令每次被 Dataloader 或是其他启动多线程的命令调用时，都会将原本的父进程程序作为一个包导入一次，而导入过程中会将脚本命令依次解析，如果没有特殊的注释，又回碰到同样的多线程启动命令，每个子线程都会再次创建新的子线程，以此迭代，最终炸掉线程池。而 Linux 中的 fork 由于是继承父线程的上下文，则不存在此问题。

解决方案

只需要规范编写代码，假如原本的包含多线程的代码块为 block()，则可改为：

if __name__ == '__main__':
    block()

即可解决。

由于仅有顶层线程的名称属性是 __main__，子线程的名称属性都不是，所以 spawn 创建子线程时导入的代码都不会执行此部分，线程只会创建一次，问题解决。

总结

在今后的开发中，应该无论如何也将该脚本独立执行才需要执行的部份放入 if __name__ == '__main__': 中，不仅可以规避多线程相关的问题，也可避免潜在的今后被其他程序导入时出现的冲突。