多处理代码使用 numpy 工作，但使用 pytorch 死锁

Question

在尝试通过 pytorch 使用多处理时，我遇到了似乎是死锁的问题。等效的 numpy 代码按我的预期工作。

我制作了我的代码的简化版本：一个由 4 个工作人员组成的池，执行 1000 次数组范围的广播操作（所以每个工作人员大约 250 个）。有问题的数组是 100,000 x 3，广播操作是用单个 1 x 3 行数组减去所有行。大数组是共享/全局数组，每次迭代的行数组都不一样。

使用 numpy 的代码完全按照预期运行，池化的 worker 显示出比等效 for 循环快 4 倍的速度。

然而，pytorch 中的代码遇到了死锁（我假设）：没有一个 worker 完成一次数组广播操作。

下面的 numpy 代码打印以下内容：

在 my_subtractor 上完成 for 循环：耗时 8.1504 秒。
my_subtractor 上的完成池：耗时 2.2247 秒。

另一方面，pytorch 代码会打印这个然后停止：

在 my_subtractor 上完成 for 循环：耗时 3.1082 秒。
BLA
BLA
BLA
BLA

“BLA”打印语句只是为了表明每个工作人员都陷入了 - 显然 - 一个死锁状态。其中恰好有 4 个：每个工作人员进入 - 并陷入 - 一个迭代。

如果你有足够的野心去重现，请注意它在 Windows 上不起作用，因为它没有被 if __name__ == '__main__': 包裹（我在某处读到你需要这个，因为 Windows 处理启动过程的方式）。您还需要创建一个名为 my_globals.py 的空文件。

这里是 numpy 代码

from time import time
import numpy as np
import my_globals
from multiprocessing import Pool as ThreadPool

# shared memory by virtue of being global
my_globals.minuend = np.random.rand(100000,3)

# array to be iterated over in for loop / pool of workers
subtrahends = np.random.rand(10000,3)

# function called at each iteration (broadcast operation)
def my_subtractor(subtrahend):
    my_globals.minuend - subtrahend
    return 0

# launch for loop
ts = time()
for idx, subtrahend in enumerate(subtrahends):
    my_subtractor(subtrahend)
te = time()
print('Finished for loop over my_subtractor: took %2.4f seconds.' % (te - ts))

# launch equivalent pool of workers
ts = time()
pool = ThreadPool(4)
pool.map(my_subtractor, subtrahends)
pool.close()
pool.join()
te = time()
print('Finished pool over my_subtractor: took %2.4f seconds.' % (te - ts))

这是等效的pytorch代码：

from time import time
import torch
import my_globals

from torch.multiprocessing import Pool as ThreadPool

# necessary on my system because it has low limits for number of file descriptors; not recommended for most systems,
# see: https://pytorch.org/docs/stable/multiprocessing.html#file-descriptor-file-descriptor
torch.multiprocessing.set_sharing_strategy('file_system')

# shared memory by virtue of being global
my_globals.minuend = torch.rand(100000,3)

# array to be iterated over in for loop / pool of workers
subtrahends = torch.rand(10000,3)

# function called at each iteration (broadcast operation)
def my_subtractor(subtrahend, verbose=True):
    if verbose:
        print("BLA") # -- prints for every worker in the pool (so 4 times total)
    my_globals.minuend - subtrahend
    if verbose:
        print("ALB") # -- doesn't print for any worker
    return 0

# launch for loop
ts = time()
for idx, subtrahend in enumerate(subtrahends):
    my_subtractor(subtrahend, verbose=False)
te = time()
print('Finished for loop over my_subtractor: took %2.4f seconds.' % (te - ts))

# launch equivalent pool of workers
ts = time()
pool = ThreadPool(4)
pool.map(my_subtractor, subtrahends)
pool.close()
pool.join()
te = time()
print('Finished pool over my_subtractor: took %2.4f seconds.' % (te - ts))

Answer 1

您可以尝试设置 OMP_NUM_THREADS=1 环境变量来尝试解决此问题。它帮助我解决了 DataLoader+OpenCV 死锁问题。