如何在这个 for 循环中实现多线程？

Question

考虑一下这段代码sn-p

from tqdm import trange


def main_game(depth1, depth2):
    # some operator with complexity O(20^max(depth1,depth2))
    return depth1+depth2


DEPTH_MAX = 5
total = 0
for depth1 in range(1, DEPTH_MAX + 1):
    for depth2 in range(1, DEPTH_MAX + 1):
        for i in trange(100):
            total += main_game(depth1, depth2)

print(total)

我在 main_game() 中使用 minimax 算法，分支因子 = 10

现在，由于第三个 for 循环有一个耗时的函数（时间复杂度高达 100*O(20^5)），有什么办法可以让它运行得更快吗？我正在考虑并行化（例如多线程）。有什么建议吗？

Answer 1

使用multiprocessing，然后使用Pool().starmap()。 starmap() 以并行方式为您的函数提供准备好的参数元组。并同步收集结果。 如果结果的顺序无关紧要，您可以使用异步版本.starmap_async().get()。

还有Pool().apply() 和Pool.map() 及其_async() 版本，但您实际上只需要学习Pool().starmap()。这只是一些语法差异。

import multiprocessing as mp
n_cpu = mp.cpu_count()

# let's say your function is a diadic function (takes two arguments)
def main_game(depth1, depth2):
    return depth1 + depth2

DEPTH_MAX = 5
depths = list(range(1, DEPTH_MAX + 1))

# let's pre-prepare the arguments - because that goes fast!
depth1_depth2_pairs = [(d1, d2) for d1 in depths for d2 in depths]

# 1: Init multiprocessing.Pool()
pool = mp.Pool(n_cpu)
# 2: pool.starmap()
results = pool.starmap(main_game, depth_1_depth_2_pairs)
# 3: pool.close()
pool.close()

total = sum(results) # this does your `total +=`

## in this case, you could even use
results = pool.starmap_async(main_game, depth_1_depth_2_pairs).get()
## because the order doesn't matter, if you sum them all up
## which is commutative.

这一切你可以使用with 构造写得更好一些（即使发生错误，它也会自动关闭，因此它不仅可以节省你的打字时间，而且更安全。

import multiprocessing as mp

n_cpu = mp.cpu_count()

def main_game(depth1, depth2):
    return depth1 + depth2

DEPTH_MAX = 5
depths = range(1, DEPTH_MAX + 1)
depth1_depth2_pairs = [(d1, d2) for d1 in depths for d2 in depths]

with mp.Pool(n_cpu) as pool:
    results = pool.starmap_async(main_game, depth_1_depth_2_pairs).get()

total = sum(results)