如何使用多线程来加快嵌套的循环计算?

时间:2018-10-21 19:26:48

标签: python multithreading optimization multiprocessing numba

我正在尝试对大型数组进行数值积分,并且计算需要很长时间。我试图通过使用numba和jit装饰器来加速代码,但是不支持numpy.trapz。

我的新想法是创建n个线程以并行运行计算,但是我想知道如何做到这一点,或者它是否可行?

引用以下代码

我可以让sz [2]许多线程在调用ZO_SteadState来计算值的同时运行吗?

    for i in range(sz[1]):
        phii = phi[i]
        for j in range(sz[2]):
            s = tau[0, i, j, :].reshape(1, n4)
            [R3, PHI3, S3] = meshgrid(rprime, phiprime, s)
            BCoeff = Bessel0(bm * R3)

            SS[0, i, j] = ZO_SteadyState(alpha, b,bm,BCoeff,Bessel_Denom, k2,maxt,phii, PHI2, PHI3, phiprime,R3,rprime,s,S3, T,v)

相关计算。

@jit()
def ZO_SteadyState(alpha, b,bm,BCoeff,Bessel_Denom, k2,maxt,phii, PHI2, PHI3, phiprime,R3,rprime,s,S3, T,v):
    g = 1000000 * exp(-(10 ** 5) * (R3 - (b / maxt) * S3) ** 2) * (
            exp(-(10 ** 5) * (PHI3 - 0) ** 2) + exp(-(10 ** 5) * (PHI3 - 2 * np.pi) ** 2) + exp(
        -(10 ** 5) * (PHI3 - 2 * np.pi / 3) ** 2) + exp(
        -(10 ** 5) * (PHI3 - 4 * np.pi / 3) ** 2))  # stationary point heat source.

    y = R3 * ((np.sqrt(2) / b) * (1 / (np.sqrt((H2 ** 2 / bm ** 2) + (1 - (v ** 2 / (bm ** 2 * b ** 2))))))
              * (BCoeff / Bessel_Denom)) * np.cos(v * (phii - PHI3)) * g

    x = (np.trapz(y, phiprime, axis=0)).reshape(1, 31, 300)

    # integral transform of heat source. integral over y-axis
    gbarbar = np.trapz(x, rprime, axis=1)

    PHI2 = np.meshgrid(phiprime, s)[0]

    sz2 = PHI2.shape
    f = h2 * 37 * Array_Ones((sz2[0], sz[1]))  # boundary condition.

    fbar = np.trapz(np.cos(v * (phii - PHI2)) * f, phiprime, 1).reshape(1, n4)  # integrate over y

    A = (alpha / k) * gbarbar + ((np.sqrt(2) * alpha) / k2) * (
                1 / (np.sqrt((H2 ** 2 / bm ** 2) + (1 - (v ** 2 / (bm ** 2 * b ** 2)))))) * fbar

    return np.trapz(exp(-alpha * bm ** 2 * (T[0, i, j] - s)) * A, s)

2 个答案:

答案 0 :(得分:1)

另一个概念实现,其中包含进程产生的进程(编辑:已测试jit):

import numpy as np

# better pickling
import pathos 
from contextlib import closing


from numba import jit

#https://stackoverflow.com/questions/47574860/python-pathos-process-pool-non-daemonic
import multiprocess.context as context
class NoDaemonProcess(context.Process):
    def _get_daemon(self):
        return False
    def _set_daemon(self, value):
        pass
    daemon = property(_get_daemon, _set_daemon)

class NoDaemonPool(pathos.multiprocessing.Pool):
    def Process(self, *args, **kwds):
        return NoDaemonProcess(*args, **kwds)




# matrix dimensions
x = 100 # i
y = 500 # j

NUM_PROCESSES = 10 # total NUM_PROCESSES*NUM_PROCESSES will be spawned

SS = np.zeros([x, y], dtype=float)

@jit
def foo(i):
    return (i*i + 1)
@jit
def bar(phii, j):
    return phii*(j+1)

# The code which is implemented down here:
'''
for i in range(x):
    phii = foo(i)
    for j in range(y):
        SS[i, j] = bar(phii, j)
'''

# Threaded version:
# queue is in global scope


def outer_loop(i):

    phii = foo(i)

    # i is in process scope
    def inner_loop(j):
        result = bar(phii,j)
        # the data is coordinates and result
        return (i, j, result)


    with closing(NoDaemonPool(processes=NUM_PROCESSES)) as pool:
        res = list(pool.imap(inner_loop, range(y)))
    return res

with closing(NoDaemonPool(processes=NUM_PROCESSES)) as pool:    
    results = list(pool.imap(outer_loop, range(x)))

result_list = []
for r in results:
    result_list += r


# read results from queue
for res in result_list:
    if res:
        i, j, val = res
        SS[i,j] = val


# check that all cells filled
print(np.count_nonzero(SS)) # 100*500

编辑:说明。

此代码中所有复杂的原因是,我想比OP要求做更多的并行化。如果只有内部循环并行化,那么外部循环将保留,因此对于外部循环的每次迭代,都会创建新的处理池并执行内部循环的计算。只要对我来说,该公式不依赖于外循环的其他迭代,我决定对所有内容进行并行化:现在,将外循环的计算分配给池中的进程,之后每个“外循环”进程创建自己的新池,并生成其他进程以执行内部循环的计算。

我可能是错的,并且外循环一定不能并行化;在这种情况下,您只能保留内部进程池。

使用进程池可能不是最佳解决方案,尽管在创建和销毁池时会浪费时间。更加有效(但需要模式手工)的解决方案将是一劳永逸地实例化N个流程,然后使用多重处理Queue()将数据馈入其中并接收结果。因此,您应该首先测试该多处理解决方案是否能为您带来足够的加速(如果运行池的构造和破坏时间比Z0_SteadyState短,就会发生这种情况。)

下一个难题是人造的无守护进程池。守护进程用于正常停止应用程序:当主程序退出时,守护进程将以静默方式终止。但是,守护进程无法生成子进程。在您的示例中,您需要等待每个进程结束以检索数据,因此我将它们设置为非守护进程,以允许生成子进程来计算内部循环。

数据交换:我认为与实际计算相比,填充矩阵所需的数据量和执行该数据所需的时间较少。因此,我使用了池和pool.imap函数(比.map()快一点。您也可以尝试.imap_unordered(),但是在您的情况下应该没有太大的区别)。因此,内部池将一直等到所有结果都被计算出来并作为列表返回。因此,外部池返回必须连接的列表列表。然后在单快速循环中根据这些结果重建矩阵。

注意with closing():在此语句下的内容完成后,它将自动关闭池,避免僵尸进程占用内存。

此外,您可能会注意到,我在另一个函数中怪异地定义了一个函数,并且在进程内部,我可以访问一些尚未传递到那里的变量:iphii。发生这种情况的原因是,进程可以访问使用copy-on-change策略(默认为fork模式)从其启动的全局范围。这不涉及酸洗,而且速度很快。

最后一条评论是关于使用pathos库而不是标准的multiprocessingconcurrent.futuressubprocess等。原因是pathos具有更好的酸洗性使用的库,因此它可以序列化标准库无法执行的功能(例如,lambda函数)。我不知道您的功能,所以我使用了更强大的工具来避免进一步的问题。

最后一件事:多处理与线程。您可以将pathos的处理池从ThreadPoolExecutor更改为标准concurrent.futures,就像刚开始启动该代码时一样。但是,在执行期间,在我的系统上CPU仅加载100%(即使用了一个内核,看起来好像所有8个内核都以15-20%加载)。我不太了解线程和进程之间的区别,但是对我来说,进程似乎可以利用所有内核(每个内核100%,总计800%)。

答案 1 :(得分:0)

这是我可能会做的总体想法。没有足够的上下文来提供更可靠的示例。您必须将所有变量设置到该类中。

import multiprocessing

pool = multiprocessing.Pool(processes=12)
runner = mp_Z0(variable=variable, variable2=variable2)

for i, j, v in pool.imap(runner.run, range(sz[1]):
    SS[0, i, j] = v


class mp_Z0:

    def __init__(self, **kwargs):
        for k, v in kwargs:
            setattr(self, k, v)


    def run(self, i):
        phii = self.phi[i]
        for j in range(self.sz[2]):
            s = self.tau[0, i, j, :].reshape(1, self.n4)
            [R3, PHI3, S3] = meshgrid(self.rprime, self.phiprime, s)
            BCoeff = Bessel0(self.bm * R3)

            return (i, j, ZO_SteadyState(self.alpha, self.b, self.bm, BCoeff, Bessel_Denom, self.k2, self.maxt, phii, self.PHI2, PHI3, self.phiprime, R3, self.rprime, self.s, S3, self.T, self.v))

这是一个不使用类的示例(假设所有内容都在本地名称空间中):

import multiprocessing

pool = multiprocessing.Pool(processes=12)    
def runner_function(i):
        phii = phi[i]
        for j in range(sz[2]):
            s = tau[0, i, j, :].reshape(1, n4)
            [R3, PHI3, S3] = meshgrid(rprime, phiprime, s)
            BCoeff = Bessel0(bm * R3)

            return (i, j, ZO_SteadyState(alpha, b, bm, BCoeff, Bessel_Denom, k2, maxt, phii, PHI2, PHI3,
                                         phiprime, R3, rprime, s, S3, T, v))

for i, j, v in pool.imap(runner_function, range(sz[1]):
    SS[0, i, j] = v