Python多进程内存管理处理多维数组循环

Question

这里有一个使用Python进行多进程编程的示例程序。我发现每个进程的内存使用量比它应该使用的内存高出大约2到3倍。如果我只计算一个进程，使用的内存大约是应该的1.3倍，而随着进程数量的增加，这个比例还会更高。

举个例子，对于一个大小为1000*1000*1000的浮点数组，它应该使用8GB的内存，但我发现当8个处理器同时运行时，内存使用量竟然达到了25GB！我读到多进程是使用共享内存的，所以我不太确定内存是在哪里泄漏的。以下是代码：

#To use the code, please take care of your RAM.
#If you have higher RAM, kindly try for the bigger arrays to see the difference clearly.

from numpy import *
import multiprocessing as mp

a = arange(0, 2500, 5)
b = arange(0, 2500, 5)
c = arange(0, 2500, 5)  
a0 = 540. #random values
b0 = 26.
c0 = 826.
def rand_function(a, b, c, a0, b0, c0):
    Nloop = 100.
    def loop(Nloop, out):
        res_total = zeros((500, 500, 500), dtype = 'float') 
        n = 1
        while n <= Nloop:
            rad = sqrt((a-a0)**2 + (b-b0)**2 + (c-c0)**2)
            res_total = res_total + rad
            n +=1 
        out.put(res_total)
    out = mp.Queue() 
    jobs = []
    Nprocs = mp.cpu_count()
    print "No. of processors : ", Nprocs
    for i in range(Nprocs):
        p = mp.Process(target = loop, args=(Nloop/Nprocs, out)) 
        jobs.append(p)
        p.start()

    final_result = zeros((500,500,500), dtype = 'float')

    for i in range(Nprocs):
        final_result = final_result + out.get()

    p.join()
test = rand_function(a,b,c,a0, b0, c0)

有没有人能告诉我内存泄漏的原因？以及如何解决这个问题？非常感谢！

Answer 1

顺便提一下，如果你想要同时对多个数组进行求和，使用 multiprocessing.Pool 会更好，这样你就不用自己处理循环的输出值了。而且，你的代码并没有给所有的工作进程分配相同的任务，我不确定这是不是故意的。

import numpy as np
import multiprocessing as mp


def loop(arg):
    max_n, a, b, c, a0, b0, c0 = arg
    res_total = np.zeros(shape, dtype=np.float)
    print 'starting'
    for _ in range(max_n):
        rad = np.sqrt((a - a0) ** 2 + (b - b0) ** 2 + (c - c0) ** 2)
        res_total = res_total + rad
    print 'done'
    return res_total


def rand_function(a, b, c, a0, b0, c0):
    c_cpu = mp.cpu_count()
    n_loop = 10
    print "No. of processors : ", c_cpu
    pool = mp.Pool(c_cpu)
    out = pool.map(loop, [(n_loop / c_cpu, a, b, c, a0, b0, c0) 
                             for _ in range(c_cpu)])

    print 'collating'
    final_result = np.zeros(shape, dtype='float')
    for i in out:
        final_result += i
    print final_result.shape


shape = (50, 50, 50)
rand_function(np.arange(0, 250, 5), np.arange(0, 250, 5), 
                  np.arange(0, 250, 5), 540, 26, 826)

在我的电脑上，每个工作进程大约使用一GB的内存。你原来的代码一开始每个工作进程用了大约1.4GB（然后增长到2GB）。我怀疑这和输出队列的修改有关，这可能会触发操作系统的写时复制功能（不过我不太确定）。

Answer 2

一些会占用（很多）内存的东西：

1. res_total
2. 在 res_total = res_total + rad 这个表达式的右边，会创建一个临时数组，这个数组会和 res_total 同时存在一段时间。使用 += 可以避免这种情况。
3. out.put(res_total) 会把数组进行序列化，使用的内存大致和之前一样。

这应该能解释为什么内存使用量可能会比单独的 1 要高得多。