如何在 for 循環中使用 Python 中的多處理來生成嵌套字典？

Question

我目前正在生成一個嵌套字典，它通過使用嵌套的 for 循環來保存一些 arrays。 不幸的是，這需要相當長的時間。 我意識到我正在使用的服務器有幾個可用的內核，所以我想知道 Python 的多處理庫是否有助於加快字典的創建。

嵌套的 for 循環看起來像這樣（實際的計算量更大、更復雜）：

import numpy as np

data_dict = {}
for s in range(1,5):
    data_dict[s] = {}
    for d in range(1,5):
        if s * d > 4:
            data_dict[s][d] = np.zeros((s,d))
        else:
            data_dict[s][d] = np.ones((s,d))

所以這就是我嘗試的：

from multiprocessing import Pool
import numpy as np

data_dict = {}
def process():
    #sci=fits.open('{}.fits'.format(name))
    for s in range(1,5):
        data_dict[s] = {}
        for d in range(1,5):
            if s * d > 4:
                data_dict[s][d] = np.zeros((s,d))
            else:
                data_dict[s][d] = np.ones((s,d))

if __name__ == '__main__':
    pool = Pool()                         # Create a multiprocessing Pool
    pool.map(process)  

但是 pool.map （最后一行）似乎需要一個可迭代的，我不知道在那里插入什么。

Answer 1

在我看來，真正的問題是需要什么樣的處理來計算字典的條目以及有多少條目。

這種處理對於理解multiprocessing處理是否可以顯着加快字典的創建至關重要。 如果您的計算受 I/O 限制，則應使用多線程，而如果受 CPU 限制，則應使用多處理。 你可以在這里找到更多關於這個的信息。

假設每個條目的值可以獨立計算並且這種計算受 CPU 限制，讓我們對單進程和多進程實現（基於multiprocessing庫）之間的差異進行基准測試。

以下代碼用於在某些場景中測試這兩種方法，改變每個條目所需的計算復雜度和條目數（對於多進程實現，使用了 7 個進程）。

import timeit
import numpy as np

def some_fun(s, d, n=1):
    """A function with an adaptable complexity"""
    a = s * np.ones(np.random.randint(1, 10, (2,))) / (d + 1)
    for _ in range(n):
        a += np.random.random(a.shape)
    return a

# Code to create dictionary with only one process
setup_simple = "from __main__ import some_fun, n_first_level, n_second_level, complexity"

code_simple = """
data_dict = {}
for s in range(n_first_level):
    data_dict[s] = {}
    for d in range(n_second_level):
        data_dict[s][d] = some_fun(s, d, n=complexity)
"""

# Code to create a dictionary with multiprocessing: we are going to use all the available cores except 1
setup_mp = """import numpy as np
import multiprocessing as mp
import itertools
from functools import partial
from __main__ import some_fun, n_first_level, n_second_level, complexity

n_processes = mp.cpu_count() - 1
# Uncomment if you want to know how many concurrent processes are you going to use
# print(f'{n_processes} concurrent processes')
"""

code_mp = """
with mp.Pool(processes=n_processes) as pool:
    dict_values = pool.starmap(partial(some_fun, n=complexity), itertools.product(range(n_first_level), range(n_second_level)))
data_dict = {
    k: dict(zip(range(n_second_level), dict_values[k * n_second_level: (k + 1) * n_second_level]))
    for k in range(n_first_level)
}
"""

# Time the code with different settings
print('Execution time on 10 repetitions: mean [std]')
for label, complexity, n_first_level, n_second_level in (
    ("TRIVIAL FUNCTION", 0, 10, 10),
    ("TRIVIAL FUNCTION", 0, 500, 500),
    ("SIMPLE FUNCTION", 5, 500, 500),
    ("COMPLEX FUNCTION", 50, 100, 100),
    ("HEAVY FUNCTION", 1000, 10, 10),
):
    print(f'\n{label}, {n_first_level * n_second_level} dictionary entries')
    for l, t in (
        ('Single process', timeit.repeat(stmt=code_simple, setup=setup_simple, number=1, repeat=10)),
        ('Multiprocess', timeit.repeat(stmt=code_mp, setup=setup_mp, number=1, repeat=10)),
    ):
        print(f'\t{l}: {np.mean(t):.3e} [{np.std(t):.3e}] seconds')

這些是結果：

Execution time on 10 repetitions: mean [std]

TRIVIAL FUNCTION, 100 dictionary entries
    Single process: 7.752e-04 [7.494e-05] seconds
    Multiprocess: 1.163e-01 [2.024e-03] seconds

TRIVIAL FUNCTION, 250000 dictionary entries
    Single process: 7.077e+00 [7.098e-01] seconds
    Multiprocess: 1.383e+00 [7.752e-02] seconds

SIMPLE FUNCTION, 250000 dictionary entries
    Single process: 1.405e+01 [1.422e+00] seconds
    Multiprocess: 2.858e+00 [5.742e-01] seconds

COMPLEX FUNCTION, 10000 dictionary entries
    Single process: 1.557e+00 [4.330e-02] seconds
    Multiprocess: 5.383e-01 [5.330e-02] seconds

HEAVY FUNCTION, 100 dictionary entries
    Single process: 3.181e-01 [5.026e-03] seconds
    Multiprocess: 1.171e-01 [2.494e-03] seconds

如您所見，假設您有一個 CPU 受限計算，多進程方法在大多數情況下都能獲得更好的結果。 僅當您對每個條目的計算量非常輕和/或條目數量非常有限時，才應首選單進程方法。

另一方面， multiprocessing提供的改進是有代價的：例如，如果您對每個條目的計算使用大量的 memory，則可能會導致OutOfMemory錯誤，這意味着您必須改進代碼並使其更多避免它的復雜性，在 memory 占用和執行時間減少之間找到適當的平衡。 如果您環顧四周，有很多問題在詢問如何解決由於multiprocessing的非最佳使用而導致的 memory 問題。 換句話說，這意味着您的代碼將不太容易閱讀和維護。

總而言之，您應該判斷執行時間的改進是否值得，即使有可能。