使用多处理保存多个matplotlib图形

Question

我有一段代码，可以从多个文件中读取数据，这些文件的名字是 001.txt、002.txt，一直到 411.txt。我想从每个文件中读取数据，然后把它们画成图，最后保存为 001.jpg、002.jpg，一直到 411.jpg。

我可以通过循环这些文件来完成这个任务，但我想用 multiprocess 模块来加快速度。
不过，当我使用下面的代码时，电脑就卡住了——我什么都点不了，鼠标能动，声音也在，但我只能强制关机。

显然，我在使用 multiprocess 模块和 matplotlib 时搞错了。我之前用类似的代码生成数据，并保存到文本文件中，没有遇到问题。那么我到底漏掉了什么呢？

    import multiprocessing

    def do_plot(number):
        fig = figure(number)

        a, b = random.sample(range(1,9999),1000), random.sample(range(1,9999),1000)
        # generate random data
        scatter(a, b)

        savefig("%03d" % (number,) + ".jpg")
        print "Done ", number
        close()



    for i in (0, 1, 2, 3):
        jobs = []
    #    for j in chunk:
        p = multiprocessing.Process(target = do_plot, args = (i,))
        jobs.append(p)
        p.start()
    p.join()

Answer 1

使用 multiprocessing 时，最重要的一点是，模块的主要代码只能在主进程中运行。我们可以通过测试 if __name__ == '__main__' 来实现这一点，下面是示例：

import matplotlib.pyplot as plt
import numpy.random as random
from multiprocessing import Pool


def do_plot(number):
    fig = plt.figure(number)

    a = random.sample(1000)
    b = random.sample(1000)

    # generate random data
    plt.scatter(a, b)

    plt.savefig("%03d.jpg" % (number,))
    plt.close()

    print("Done ", number)


if __name__ == '__main__':
    pool = Pool()
    pool.map(do_plot, range(4))

另外，我还把创建单独进程的方式换成了进程池，这样在处理很多图片时会更有效，因为它只会使用与你的电脑核心数量相同的进程。